1.修改默认情况下获取进程数和协程数方式

2.更新依赖

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

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记录')

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

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  # 是否使用递归爆破(默认禁用)