From 4b3db5aff7d5dbde0260e723b9e27d7a7fe9de50 Mon Sep 17 00:00:00 2001 From: kelvinBen Date: Tue, 10 Aug 2021 01:40:16 +0800 Subject: [PATCH] =?UTF-8?q?-=20=E6=96=B0=E5=A2=9E=E9=80=9A=E8=BF=87?= =?UTF-8?q?=E7=BD=91=E7=BB=9C=E6=89=B9=E9=87=8F=E4=B8=8B=E8=BD=BD=E6=96=87?= =?UTF-8?q?=E4=BB=B6=20-=20=E6=96=B0=E5=A2=9ETXT=E6=96=87=E4=BB=B6?= =?UTF-8?q?=E6=89=B9=E9=87=8F=E8=BF=9B=E8=A1=8C=E6=96=87=E4=BB=B6=E4=B8=8B?= =?UTF-8?q?=E8=BD=BD=E6=88=96=E8=80=85=E6=89=B9=E9=87=8F=E5=AF=B9IPA?= =?UTF-8?q?=E3=80=81APK=E6=96=87=E4=BB=B6=E8=BF=9B=E8=A1=8C=E5=A4=84?= =?UTF-8?q?=E7=90=86=20-=20=E4=BF=AE=E6=94=B9=E9=80=9A=E8=BF=87=E7=BD=91?= =?UTF-8?q?=E7=BB=9C=E6=96=87=E4=BB=B6=E4=B8=8B=E8=BD=BD=E9=80=BB=E8=BE=91?= =?UTF-8?q?=20-=20=E4=BF=AE=E6=94=B9=E6=96=87=E4=BB=B6=E4=B8=8B=E8=BD=BD?= =?UTF-8?q?=E6=97=B6=E7=9A=84=E6=96=87=E4=BB=B6=E5=91=BD=E5=90=8D=E9=80=BB?= =?UTF-8?q?=E8=BE=91=20-=20=E5=88=A0=E9=99=A4libs/task/download=5Ftask.py?= =?UTF-8?q?=E6=96=87=E4=BB=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- libs/core/download.py | 88 ++++++++++++++++++++++++++++++++------ libs/task/base_task.py | 62 ++++++++++++++++++--------- libs/task/download_task.py | 76 -------------------------------- 3 files changed, 118 insertions(+), 108 deletions(-) delete mode 100644 libs/task/download_task.py diff --git a/libs/core/download.py b/libs/core/download.py index 7691c63..7d07104 100644 --- a/libs/core/download.py +++ b/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__() \ No newline at end of file + self.__start__() \ No newline at end of file diff --git a/libs/task/base_task.py b/libs/task/base_task.py index e90386c..700edb3 100644 --- a/libs/task/base_task.py +++ b/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) \ No newline at end of file + self.download_file_queue.put(dir_or_file_path) \ No newline at end of file diff --git a/libs/task/download_task.py b/libs/task/download_task.py deleted file mode 100644 index ebba650..0000000 --- a/libs/task/download_task.py +++ /dev/null @@ -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 \ No newline at end of file