- 新增通过网络批量下载文件

- 新增TXT文件批量进行文件下载或者批量对IPA、APK文件进行处理
- 修改通过网络文件下载逻辑
- 修改文件下载时的文件命名逻辑
- 删除libs/task/download_task.py文件

libs/core/download.py

@@ -2,10 +2,12 @@
 # -*- coding: utf-8 -*-
 # Author: kelvinBen
 # Github: https://github.com/kelvinBen/AppInfoScanner
+from genericpath import exists
 import re
 import os
 import sys
 import time
+import uuid
 import config
 import logging
 import requests
@@ -14,16 +16,73 @@ import libs.core as cores
 from requests.packages import urllib3
 from requests.adapters import HTTPAdapter    
 
-class DownloadThreads(threading.Thread):    
+class DownloadThreads(threading.Thread):   
 
-    def __init__(self,input_path,file_name,cache_path,types):
+    def __init__(self,threadID, threadName, download_file_queue, download_file_list, types):
         threading.Thread.__init__(self) 
-        self.url = input_path 
+        self.threadID = threadID
+        self.threadName = threadName
+        self.download_file_queue = download_file_queue
+        self.download_file_list = download_file_list
         self.types = types
-        self.cache_path = cache_path
-        self.file_name = file_name
+        self.cache_path = None
+    
+    def __start__(self):
+        # 从队列中取数据，直到队列数据不为空为止
+        while not self.download_file_queue.empty():
+            file_or_url = self.download_file_queue.get()
+            if not file_or_url:
+                logging.error("[x] Failed to get file!")
+                continue
+            self.__auto_update_type__(file_or_url)
 
-    def __requset__(self):
+    # 自动更新文件类型
+    def __auto_update_type__(self,file_or_url):
+        uuid_name =  str(uuid.uuid1()).replace("-","")
+        # 文件后缀为apk 或者 类型为 Android 则自动修正为Android类型
+        if file_or_url.endswith("apk") or self.types == "Android":
+            types = "Android"
+            file_name = uuid_name + ".apk"
+        # 文件后缀为dex 或者 类型为 Android 则自动修正为Android类型
+        elif file_or_url.endswith("dex") or self.types == "Android":
+            types = "Android"
+            file_name = uuid_name + ".dex"
+        # 文件后缀为ipa 或者 类型为 iOS 则自动修正为iOS类型
+        elif file_or_url.endswith("ipa") or self.types == "iOS":
+            types = "iOS"
+            file_name = uuid_name + ".ipa"
+        else:
+            # 路径以http://开头或者以https://开头 且 文件是不存在的自动修正为web类型
+            if (file_or_url.startswith("http://") or file_or_url.startswith("https://")) and (not os.path.exists(file_or_url)):
+                types = "WEB"
+                file_name = uuid_name + ".html"
+            # 其他情况如：types为WEB 或者目录 或者 单独的二进制文件 等交给后面逻辑处理
+        
+        if file_or_url.startswith("http://") or file_or_url.startswith("https://"):
+            # 进行文件下载
+            self.__file_deduplication__(file_name, uuid_name)
+            if self.cache_path:
+                file_path = self.cache_path
+                self.__download_file__(file_or_url,file_path)
+            #TODO 标记下载过的文件，避免重复下载
+        else:
+            types = self.types
+            file_path = file_or_url
+
+        self.download_file_list.append({"path": file_path, "type": types})
+         
+    # 防止文件名重复导致文件被复写
+    def __file_deduplication__(self,file_name, uuid_name):
+        cache_path = os.path.join(cores.download_dir, file_name)
+        if not os.path.exists(cache_path):
+            self.cache_path = cache_path
+            return
+        new_uuid_name = str(uuid.uuid1()).replace("-","")
+        new_file_name = file_name.replace(uuid_name,new_uuid_name) 
+        self.__file_deduplication__(new_file_name,new_uuid_name)
+
+    # 文件下载    
+    def __download_file__(self, url, file_path):
         try:
             session = requests.Session()
             session.mount('http://', HTTPAdapter(max_retries=3))
@@ -33,16 +92,17 @@ class DownloadThreads(threading.Thread):
             urllib3.disable_warnings()
 
             if config.method.upper() == "POST":
-                resp = session.post(url=self.url,params=config.data ,headers=config.headers,timeout=30)
+                resp = session.post(url=url, params=config.data, headers=config.headers, timeout=30)
             else:
-                resp = session.get(url=self.url,data=config.data ,headers=config.headers,timeout=30)
+                resp = session.get(url=url, data=config.data, headers=config.headers, timeout=30)
             
             if resp.status_code == requests.codes.ok:
+                # 下载二进制文件
                 if self.types == "Android" or self.types == "iOS":
                     count = 0
                     progress_tmp = 0
                     length = float(resp.headers['content-length'])
-                    with open(self.cache_path, "wb") as f:
+                    with open(file_path, "wb") as f:
                         for chunk in resp.iter_content(chunk_size = 512):
                             if chunk:
                                 f.write(chunk)
@@ -56,13 +116,17 @@ class DownloadThreads(threading.Thread):
                         f.close()
                 else:
                     html = resp.text
-                    with open(self.cache_path,"w",encoding='utf-8',errors='ignore') as f:
+                    with open(file_path,"w",encoding='utf-8',errors='ignore') as f:
                         f.write(html)
                         f.close()
                 cores.download_flag = True
+            else:
+               logging.error("[x] {} download fails, status code is {} !!!".format(url, str(resp.status_code)))
+
         except Exception as e:
-            raise Exception(e)
+            logging.error("[x] {} download fails, the following exception information:".format(url))
+            logging.exception(e)
 
     def run(self):
         threadLock = threading.Lock()
-        self.__requset__()
+        self.__start__()

libs/task/base_task.py

@@ -14,7 +14,7 @@ from libs.task.web_task import WebTask
 from libs.task.net_task import NetTask
 from libs.core.parses import ParsesThreads
 from libs.task.android_task import AndroidTask
-from libs.task.download_task import DownloadTask
+from libs.core.download import DownloadThreads
 
 class BaseTask(object):
 
@@ -22,7 +22,10 @@ class BaseTask(object):
         if cores.user_add_rules:
             config.filter_strs.append(r'.*'+str(cores.user_add_rules)+'.*')
         self.file_queue = Queue()
-        self.file_path_list = []
+        # 文件下载队列
+        self.download_file_queue = Queue()
+        # 文件下载列表
+        self.download_file_list = []
         self.thread_list = []
         self.app_history_list= []
         self.domain_history_list = []
@@ -30,29 +33,49 @@ class BaseTask(object):
         
     # 统一启动
     def start(self, types="Android", user_input_path="", package=""):
-        # 如果输入路径为目录，则自动检索DEX、IPA、APK等文件
+        # 如果输入路径为目录,且类型非web,则自动检索DEX、IPA、APK等文件
         if not(types == "Web") and os.path.isdir(user_input_path):
-            self.__scanner_specified_file__(self.file_path_list, user_input_path)
-        # 如果输入的路径为txt, 则加载txt中的内容
+            self.__scanner_specified_file__(user_input_path)
+        
+        # 如果输入的路径为txt, 则加载txt中的内容实现批量操作
         elif user_input_path.endswith("txt"):
             with open(user_input_path) as f:
                 lines = f.readlines()
                 for line in lines:
-                    if line.startswith("http://") or line.startswith("https://") or line.endswith("apk") or line.endswith(".dex") or line.endswith("ipa"):
-                        self.file_path_list.append(line)
+                    # http:// 或者 https:// 开头 或者 apk/dex/ipa结尾的文件且文件存在
+                    if (line.startswith("http://") or line.startswith("https://")) or ((line.endswith("apk") or line.endswith(".dex") or line.endswith("ipa")) and os.path.exists(line)):
+                        self.download_file_queue.put(line)
                 f.close()
         else:
-            # 如果是文件则追加到文件列表中
-            self.file_path_list.append(user_input_path)
+            # 如果是文件或者类型为web的目录
+            self.download_file_queue.put(user_input_path)
 
         # 长度小于1需重新选择目录
-        if len(self.file_path_list) < 1:
+        if self.download_file_queue.qsize() < 1:
             raise Exception('[x] The specified DEX, IPA and APK files are not found. Please re-enter the directory to be scanned!') 
+
+        # 统一文件下载中心
+        self.__download_file_center__(types)
         
-        # 遍历目录
-        for file_path in self.file_path_list:
+        for download_file in self.download_file_list:
+            file_path = download_file["path"]
+            types = download_file["type"]
             self.__control_center__(file_path,types)
 
+    # 统一文件下载中心
+    def __download_file_center__(self,types):
+        # 杜绝资源浪费
+        if self.download_file_queue.qsize() < cores.threads_num:
+            threads_num = self.download_file_queue.qsize()
+        else:
+            threads_num = cores.threads_num
+
+        for threadID in range(1, threads_num): 
+            threadName = "Thread - " + str(int(threadID))
+            thread = DownloadThreads(threadID, threadName, self.download_file_queue, self.download_file_list, types)
+            thread.start()
+            thread.join()
+
     # 控制中心
     def __control_center__(self,file_path,types):
         logging.info("[*] Processing {}".format(file_path))
@@ -62,8 +85,7 @@ class BaseTask(object):
         self.__history_handle__()
         logging.info("[*] The filtering rules obtained by AI are as follows: {}".format(set(config.filter_no)))
 
-        # AI 修正扫描类型
-        cache_info = DownloadTask().start(file_path, types)
+  
         cacar_path = cache_info["path"]
         types = cache_info["type"]
 
@@ -187,12 +209,12 @@ class BaseTask(object):
                 f.close()
     
     # 扫描指定后缀文件
-    def __scanner_specified_file__(self, file_list, root_dir, file_suffix=['dex','ipa','apk']):
-        dir_or_files = os.listdir(root_dir)
-        for dir_or_file in dir_or_files:
-            dir_or_file_path = os.path.join(root_dir,dir_or_file) 
+    def __scanner_specified_file__(self, base_dir, file_suffix=['dex','ipa','apk']):
+        files = os.listdir(base_dir)
+        for file in files:
+            dir_or_file_path = os.path.join(base_dir,file) 
             if os.path.isdir(dir_or_file_path):
-                self.__scanner_specified_file__(file_list,dir_or_file_path,file_suffix)
+                self.__scanner_specified_file__(dir_or_file_path,file_suffix)
             else: 
                 if dir_or_file_path.split(".")[-1] in file_suffix:
-                    file_list.append(dir_or_file_path)
+                    self.download_file_queue.put(dir_or_file_path)

libs/task/download_task.py

@@ -1,76 +0,0 @@
-#! /usr/bin/python3
-# -*- coding: utf-8 -*-
-# Author: kelvinBen
-# Github: https://github.com/kelvinBen/AppInfoScanner
-import os
-import re
-import time
-import config
-import hashlib
-import logging
-from queue import Queue
-import libs.core as cores
-from libs.core.download import DownloadThreads
-
-class DownloadTask(object):
-    def __init__(self):
-        self.download_file_queue = Queue()
-        self.thread_list = []
-
-    def start(self, path, types):
-        self.__local_or_remote__(path, types)
-
-        for threadID in range(1, cores.threads_num): 
-            name = "Thread - " + str(int(threadID))
-            thread = DownloadThreads(threadID,name,self.download_file_queue)
-            thread.start()
-            thread.join()     
-
-    # 判断文件是本地加载还是远程加载
-    def __local_or_remote__(self,path,types):
-        # 处理本地文件
-        if not(path.startswith("http://") or path.startswith("https://")):
-            if not os.path.isdir(path): # 不是目录
-                return {"path":path,"type":types}
-            else: # 目录处理
-                return {"path":path,"type":types}
-        else:
-           self.__net_header__(path,types)
-            # self.download_file_queue.put(path)
-
-    # 处理网络请求
-    def __net_header__(self, path, types):
-        create_time = time.strftime("%Y%m%d%H%M%S", time.localtime())
-        if path.endswith("apk") or types == "Android":
-            types = "Android"
-            file_name = create_time+ ".apk"
-        elif path.endswith("ipa") or types == "iOS":
-            types = "iOS"
-            file_name = create_time + ".ipa"
-        else:
-            types = "WEB"
-            file_name = create_time + ".html"
-
-        logging.info("[*] Detected that the task is not local, preparing to download file......")
-        cache_path = os.path.join(cores.download_dir, file_name)
-        self.download_file_queue.put({"path":path, "cache_path":cache_path, "types":types})        
-        # thread = DownloadThreads(path,file_name,cache_path,types)
-        # thread.start()
-        # thread.join()
-        return {"path":cache_path,"type":types}
-    
-    def __update_type__(self, path, types, file_name=None):
-        create_time = time.strftime("%Y%m%d%H%M%S", time.localtime())
-        if path.endswith("apk") or types == "Android":
-            types = "Android"
-            if not file_name:
-                file_name = create_time+ ".apk"
-        elif path.endswith("ipa") or types == "iOS":
-            types = "iOS"
-            if not file_name:
-                file_name = create_time + ".ipa"
-        else:
-            types = "WEB"
-            if not file_name:
-                file_name = create_time + ".html"
-        return types,file_name