diff --git a/Pipfile b/Pipfile index 32802b7..7d16102 100644 --- a/Pipfile +++ b/Pipfile @@ -23,6 +23,7 @@ cloudscraper = "*" js2py = "*" tablib = "*" brotli = "*" +psutil = "*" [requires] python_version = "3.8" diff --git a/Pipfile.lock b/Pipfile.lock index 65e318a..044dba7 100644 --- a/Pipfile.lock +++ b/Pipfile.lock @@ -1,7 +1,7 @@ { "_meta": { "hash": { - "sha256": "3a98b87337512897f4730ad9c7f97f01c75aebe2a098b4e4b04e9e6bd641b088" + "sha256": "a2f11736ad9a177433ceec37cc0461993824b6d175f5a11f4193a2a05ee40ce1" }, "pipfile-spec": 6, "requires": { @@ -125,11 +125,11 @@ }, "cloudscraper": { "hashes": [ - "sha256:12eebdfab1e76ea82b0ac2bcbe9a53f2d5c7041bb7109d97465f720cecab8f33", - "sha256:309e237370e5e37d9354874311295bcaf640424771f9f21e21cba91d4276365a" + "sha256:7a843769aaed50905afeff82c3ebdd1a3a5980bfa5436719f05f0db4af127790", + "sha256:ebeb5f9adcb52ed6b763688b05c34264942116b01f8f0b7a5b7de7c1d505cd67" ], "index": "pypi", - "version": "==1.2.26" + "version": "==1.2.28" }, "colorama": { "hashes": [ @@ -189,11 +189,11 @@ }, "js2py": { "hashes": [ - "sha256:168952e3827198f68eacbf84a7b23c80a55794415291f821ec6c96f3fd9c3253", - "sha256:6c45cce8eb66b16f8dc692c6e6a2d509be9544b63cfb386eadf7309a931feb62" + "sha256:306ba83c8ce16b330e532bfe2870763d0a4bbf4d6aab07799ea3c8c1446e1275", + "sha256:80f333d8c654396370974956b37679c1d45cfd4efdcf9731122e1fc809f7a67e" ], "index": "pypi", - "version": "==0.67" + "version": "==0.68" }, "loguru": { "hashes": [ @@ -264,6 +264,23 @@ ], "version": "==2.4.11" }, + "psutil": { + "hashes": [ + "sha256:1413f4158eb50e110777c4f15d7c759521703bd6beb58926f1d562da40180058", + "sha256:298af2f14b635c3c7118fd9183843f4e73e681bb6f01e12284d4d70d48a60953", + "sha256:60b86f327c198561f101a92be1995f9ae0399736b6eced8f24af41ec64fb88d4", + "sha256:685ec16ca14d079455892f25bd124df26ff9137664af445563c1bd36629b5e0e", + "sha256:73f35ab66c6c7a9ce82ba44b1e9b1050be2a80cd4dcc3352cc108656b115c74f", + "sha256:75e22717d4dbc7ca529ec5063000b2b294fc9a367f9c9ede1f65846c7955fd38", + "sha256:a02f4ac50d4a23253b68233b07e7cdb567bd025b982d5cf0ee78296990c22d9e", + "sha256:d008ddc00c6906ec80040d26dc2d3e3962109e40ad07fd8a12d0284ce5e0e4f8", + "sha256:d84029b190c8a66a946e28b4d3934d2ca1528ec94764b180f7d6ea57b0e75e26", + "sha256:e2d0c5b07c6fe5a87fa27b7855017edb0d52ee73b71e6ee368fae268605cc3f5", + "sha256:f344ca230dd8e8d5eee16827596f1c22ec0876127c28e800d7ae20ed44c4b310" + ], + "index": "pypi", + "version": "==5.7.0" + }, "pyjsparser": { "hashes": [ "sha256:2b12842df98d83f65934e0772fa4a5d8b123b3bc79f1af1789172ac70265dd21", @@ -333,10 +350,10 @@ }, "sqlalchemy": { "hashes": [ - "sha256:64a7b71846db6423807e96820993fa12a03b89127d278290ca25c0b11ed7b4fb" + "sha256:c4cca4aed606297afbe90d4306b49ad3a4cd36feb3f87e4bfd655c57fd9ef445" ], "markers": "python_version >= '3.0'", - "version": "==1.3.13" + "version": "==1.3.15" }, "tablib": { "hashes": [ diff --git a/oneforall/aiobrute.py b/oneforall/aiobrute.py index 83c6a71..ddd8d5b 100644 --- a/oneforall/aiobrute.py +++ b/oneforall/aiobrute.py @@ -184,8 +184,8 @@ class AIOBrute(Module): self.module = 'Brute' self.source = 'AIOBrute' self.target = target - self.process = process or config.brute_process_num - self.coroutine = coroutine or config.brute_coroutine_num + self.process = process or utils.get_process_num() + self.coroutine = coroutine or utils.get_coroutine_num() self.wordlist = wordlist or config.brute_wordlist_path self.recursive_brute = recursive or config.enable_recursive_brute self.recursive_depth = depth or config.brute_recursive_depth diff --git a/oneforall/common/resolve.py b/oneforall/common/resolve.py index 7c12861..fc54825 100644 --- a/oneforall/common/resolve.py +++ b/oneforall/common/resolve.py @@ -195,8 +195,8 @@ async def run_aio_resolve(subdomain_list): :param list subdomain_list: 待解析的子域列表 :return: 解析得到的结果列表 """ - process_num = config.brute_process_num - coroutine_num = config.brute_coroutine_num + process_num = utils.get_process_num() + coroutine_num = utils.get_coroutine_num() logger.log('INFOR', '正在异步查询子域的A记录') result_list = await aio_resolve(subdomain_list, process_num, coroutine_num) logger.log('INFOR', '完成异步查询子域的A记录') diff --git a/oneforall/common/utils.py b/oneforall/common/utils.py index d9fc7da..b1784f6 100644 --- a/oneforall/common/utils.py +++ b/oneforall/common/utils.py @@ -1,3 +1,4 @@ +import os import re import sys import time @@ -5,6 +6,8 @@ import random import platform from ipaddress import IPv4Address, ip_address +import psutil + import config from pathlib import Path from records import Record, RecordCollection @@ -399,3 +402,35 @@ def check_ip_public(ip_list): if not ip.is_global: return 0 return 1 + + +def get_process_num(): + process_num = config.brute_process_num + if isinstance(process_num, int): + return max(1, process_num) + else: + return os.cpu_count() + + +def get_coroutine_num(): + coroutine_num = config.brute_coroutine_num + if isinstance(coroutine_num, int): + return max(64, coroutine_num) + elif coroutine_num is None: + mem = psutil.virtual_memory() + total_mem = mem.total + g_size = 1024 * 1024 * 1024 + if total_mem <= 1 * g_size: + return 64 + elif total_mem <= 2 * g_size: + return 128 + elif total_mem <= 4 * g_size: + return 256 + elif total_mem <= 8 * g_size: + return 512 + elif total_mem <= 16 * g_size: + return 1024 + else: + return 2048 + else: + return 64 diff --git a/oneforall/config.py b/oneforall/config.py index 53a8e6f..9529292 100644 --- a/oneforall/config.py +++ b/oneforall/config.py @@ -44,9 +44,10 @@ module_thread_timeout = 360.0 # 每个收集模块线程超时时间(默认6分 enable_brute_module = False # 使用爆破模块(默认False) enable_wildcard_check = True # 开启泛解析检测(默认True) enable_wildcard_deal = True # 开启泛解析处理(默认True) -# 爆破时使用的进程数(根据系统中CPU数量情况设置 不宜大于CPU数量 默认为系统中的CPU数量) -brute_process_num = os.cpu_count() -brute_coroutine_num = 1024 # 爆破时每个进程下的协程数 +# 爆破时使用的进程数(根据计算机中CPU数量情况设置 不宜大于CPU数量) +brute_process_num = None # 默认None为系统中的CPU数量 +# 爆破时每个进程下的协程数(根据计算机中内存大小情况设置 默认为系统中的CPU数量) +brute_coroutine_num = None # 默认None根据内存大小设置 # 爆破所使用的字典路径 默认data/subdomains.txt brute_wordlist_path = data_storage_dir.joinpath('subnames.txt') enable_recursive_brute = False # 是否使用递归爆破(默认禁用) diff --git a/requirements.txt b/requirements.txt index ee07dab..900d0d8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,7 +7,7 @@ Brotli==1.0.7 bs4==0.0.1 certifi==2019.11.28 chardet==3.0.4 -cloudscraper==1.2.26 +cloudscraper==1.2.28 colorama==0.4.3 dnspython==1.16.0 docopt==0.6.2 @@ -16,11 +16,12 @@ exrex==0.10.5 fire==0.2.1 idna==2.9 jdcal==1.4.1 -Js2Py==0.67 +Js2Py==0.68 loguru==0.4.1 lxml==4.5.0 multidict==4.7.5 openpyxl==2.4.11 +psutil==5.7.0 pyjsparser==2.7.1 PySocks==1.7.1 pytz==2019.3 @@ -30,7 +31,7 @@ requests-file==1.4.3 requests-toolbelt==0.9.1 six==1.14.0 soupsieve==2.0 -SQLAlchemy==1.3.13 +SQLAlchemy==1.3.15 tablib==1.1.0 termcolor==1.1.0 tldextract==2.2.2