重构解析模块

oneforall/aiobrute.py

@@ -28,10 +28,6 @@ from common.database import Database
 from config import logger
 
 
-def init_worker():
-    signal.signal(signal.SIGINT, signal.SIG_IGN)
-
-
 def detect_wildcard(domain):
     """
     探测域名是否使用泛解析
@@ -131,25 +127,6 @@ def gen_brute_domains(domain, path):
     return domains
 
 
-def progress(pr_queue, total):
-    bar = tqdm.tqdm()
-    bar.total = total
-    bar.desc = 'Progress'
-    bar.ncols = 60
-    while True:
-        done = pr_queue.qsize()
-        bar.n = done
-        bar.update()
-        if done == total:
-            return
-
-
-async def aiodns_query_a(pr_queue, hostname):
-    results = await resolve.aiodns_query_a(hostname)
-    pr_queue.put(1)
-    return results
-
-
 class AIOBrute(Module):
     """
     OneForAll多进程多协程异步子域爆破模块
@@ -250,7 +227,7 @@ class AIOBrute(Module):
             if self.enable_wildcard and self.wildcard_deal:
                 # 通过对比查询的子域和响应的子域来判断真实子域
                 # 去掉解析到CDN的情况
-                if 'cdn' or 'waf' in name:
+                if 'cdn' in name or 'waf' in name:
                     continue
                 if not name.endswith(self.domain):
                     continue
@@ -273,27 +250,11 @@ class AIOBrute(Module):
         logger.log('INFOR', f'正在爆破{domain}的域名')
         # for task in tqdm.tqdm(tasks, total=len(tasks),
         #                       desc='Progress'):
-        m = Manager()
+        results = await resolve.aio_resolve(tasks, self.process, self.coroutine)
-        pr_queue = m.Queue()
+        self.deal_results(results)
-        loop = asyncio.get_event_loop()
+        self.save_json()
-        loop.run_in_executor(None, progress, pr_queue, len(tasks))
+        self.gen_result()
-        wrapped_query = functools.partial(aiodns_query_a, pr_queue)
+        rx_queue.put(self.results)
-        async with aiomp.Pool(processes=self.process,
-                              initializer=init_worker,
-                              childconcurrency=self.coroutine) as pool:
-            try:
-                results = await pool.map(wrapped_query, tasks)
-            except KeyboardInterrupt:
-                logger.log('ALERT', '爆破终止正在退出')
-                pool.terminate()  # 关闭pool，结束工作进程，不在处理未完成的任务。
-                self.save_json()
-                self.gen_result()
-                rx_queue.put(self.results)
-                return
-            self.deal_results(results)
-            self.save_json()
-            self.gen_result()
-            rx_queue.put(self.results)
 
     def run(self, rx_queue=None):
         self.domains = utils.get_domains(self.target)

oneforall/common/resolve.py

@@ -1,12 +1,14 @@
+import signal
+import socket
 import asyncio
 import functools
-import socket
+from multiprocessing import Manager
 
 import tqdm
+import aiomultiprocess as aiomp
 from dns.resolver import Resolver
 
 import config
-from common import utils
 from config import logger
 
 
@@ -21,22 +23,6 @@ def dns_resolver():
     return resolver
 
 
-async def dns_query_a(hostname):
-    """
-    查询A记录
-
-    :param str hostname: 主机名
-    :return: 查询结果
-    """
-    resolver = dns_resolver()
-    try:
-        answer = resolver.query(hostname, 'A')
-    except Exception as e:
-        logger.log('TRACE', e.args)
-        answer = e
-    return hostname, answer
-
-
 async def aiodns_query_a(hostname):
     """
     异步查询A记录
@@ -57,56 +43,129 @@ async def aiodns_query_a(hostname):
     return hostname, answer
 
 
-def resolve_callback(future, index, datas):
+def convert_results(result_list):
     """
-    解析结果回调处理
+    将结果列表类型转换为结果字典类型
 
-    :param future: future对象
+    :param result_list: 待转换的结果列表
-    :param index: 下标
+    :return: 转换后的结果字典
-    :param datas: 结果集
     """
-    hostname, answer = future.result()
+    result_dict = {}
-    if isinstance(answer, BaseException):
+    for result in result_list:
-        logger.log('TRACE', answer.args)
+        hostname, answer = result
-        name = utils.get_classname(answer)
+        value_dict = {'ips': None, 'reason': None, 'valid': None}
-        datas[index]['reason'] = name + ' ' + str(answer)
+        if isinstance(answer, tuple):
-        datas[index]['valid'] = 0
+            value_dict['ips'] = str(answer[2])[1:-1]
-    elif isinstance(answer, tuple):
+            result_dict[hostname] = value_dict
-        ips = answer[2]
+        elif isinstance(answer, Exception):
-        datas[index]['ips'] = str(ips)[1:-1]
+            value_dict['reason'] = str(answer.args)
+            value_dict['valid'] = 0
+            result_dict[hostname] = value_dict
+        else:
+            value_dict['valid'] = 0
+            result_dict[hostname] = value_dict
+    return result_dict
 
 
-async def bulk_query_a(datas):
+def filter_subdomain(data_list):
     """
-    批量查询A记录
+    过滤出无IPS值的子域到新的子域列表
 
-    :param datas: 待查的数据集
+    :param list data_list: 待过滤的数据列表
-    :return: 查询过得到的结果集
+    :return: 符合条件的子域列表
     """
-    logger.log('INFOR', '正在异步查询子域的A记录')
+    subdomains = []
-    tasks = []
+    for data in data_list:
-    # semaphore = asyncio.Semaphore(config.limit_resolve_conn)
-    for i, data in enumerate(datas):
         if not data.get('ips'):
             subdomain = data.get('subdomain')
-            task = asyncio.ensure_future(aiodns_query_a(subdomain))
+            subdomains.append(subdomain)
-            wrapped_callback = functools.partial(resolve_callback,
+    return subdomains
-                                                 index=i,
+
-                                                 datas=datas)
+
-            task.add_done_callback(wrapped_callback)  # 回调
+def update_data(data_list, results_dict):
-            tasks.append(task)
+    """
-    if tasks:  # 任务列表里有任务不空时才进行解析
+    更新解析结果
-        futures = asyncio.as_completed(tasks)
+
-        for future in tqdm.tqdm(futures,
+    :param list data_list: 待更新的数据列表
-                                total=len(tasks),
+    :param dict results_dict: 解析结果字典
-                                desc='Progress',
+    :return: 更新后的数据列表
-                                ncols=60):
+    """
-            await future
+    for index, data in enumerate(data_list):
-        # await asyncio.wait(tasks)  # 等待所有task完成
+        if not data.get('ips'):
+            subdomain = data.get('subdomain')
+            value_dict = results_dict.get(subdomain)
+            data.update(value_dict)
+            data_list[index] = data
+    return data_list
+
+
+def init_worker():
+    signal.signal(signal.SIGINT, signal.SIG_IGN)
+
+
+def query_progress(pr_queue, total):
+    bar = tqdm.tqdm()
+    bar.total = total
+    bar.desc = 'Resolve Progress'
+    bar.ncols = 60
+    bar.smoothing = 0
+    while True:
+        done = pr_queue.qsize()
+        bar.n = done
+        bar.update()
+        if done == total:
+            break
+    bar.close()
+
+
+async def aio_query(pr_queue, hostname):
+    """
+    异步查询主机名的A记录
+
+    :param pr_queue: 进度队列
+    :param str hostname: 主机名
+    :return: 查询结果
+    """
+    results = await aiodns_query_a(hostname)
+    pr_queue.put(1)
+    return results
+
+
+async def aio_resolve(subdomain_list, process_num, coroutine_num):
+    """
+    异步解析子域A记录
+
+    :param list subdomain_list: 待解析的子域列表
+    :param int process_num: 解析进程数
+    :param int coroutine_num: 每个解析进程下的协程数
+    :return: 解析结果
+    """
+    m = Manager()
+    pr_queue = m.Queue()
+    loop = asyncio.get_event_loop()
+    loop.run_in_executor(None, query_progress, pr_queue, len(subdomain_list))
+    wrapped_query = functools.partial(aio_query, pr_queue)
+    async with aiomp.Pool(processes=process_num,
+                          initializer=init_worker,
+                          childconcurrency=coroutine_num) as pool:
+        results = await pool.map(wrapped_query, subdomain_list)
+        return results
+
+
+async def bulk_resolve(data_list):
+    """
+    批量解析A记录并返回解析结果
+
+    :param list data_list: 待查的数据列表
+    :return: 查询过得到的结果列表
+    """
+    logger.log('INFOR', '正在异步查询子域的A记录')
+    # semaphore = asyncio.Semaphore(config.limit_resolve_conn)
+    query_subdomains = filter_subdomain(data_list)
+    process_num = config.brute_process_num
+    coroutine_num = config.brute_coroutine_num
+    results = await aio_resolve(query_subdomains, process_num, coroutine_num)
+    results_dict = convert_results(results)
+    data_list = update_data(data_list, results_dict)
     logger.log('INFOR', '完成异步查询子域的A记录')
-    return datas
+    return data_list
-
-
-def run_bulk_query(datas):
-    new_datas = asyncio.run(bulk_query_a(datas))
-    return new_datas

oneforall/common/utils.py

@@ -327,3 +327,7 @@ def get_classname(clsobj):
 
 def python_version():
     return sys.version
+
+
+def count_valid(data):
+    return len(list(filter(lambda item: item.get('valid') == 1, data)))

oneforall/oneforall.py

@@ -141,7 +141,7 @@ class OneForAll(object):
         asyncio.set_event_loop(loop)
 
         # 解析子域
-        task = resolve.bulk_query_a(self.data)
+        task = resolve.bulk_resolve(self.data)
         self.data = loop.run_until_complete(task)
 
         # 保存解析结果
@@ -166,8 +166,8 @@ class OneForAll(object):
         self.datas.extend(self.data)
         # 在关闭事件循环前加入一小段延迟让底层连接得到关闭的缓冲时间
         loop.run_until_complete(asyncio.sleep(0.25))
-        valid_count = len(list(filter(lambda item: item.get('valid') == 1, self.data)))
+        count = utils.count_valid(self.data)
-        logger.log('INFOR', f'经验证{self.domain}有效子域{valid_count}个')
+        logger.log('INFOR', f'经验证{self.domain}有效子域{count}个')
 
         # 保存请求结果
         db.clear_table(self.domain)