重构:OneForAll入口逻辑更加清晰

oneforall/common/database.py

@@ -6,6 +6,7 @@ SQLite数据库初始化和操作
 """
 
 import records
+
 import config
 from records import Connection
 from config import logger
@@ -28,7 +29,7 @@ class Database(object):
             return db_path
         protocol = 'sqlite:///'
         if not db_path:  # 数据库路径为空连接默认数据库
-            db_path = f'{protocol}{config.result_save_path}/result.sqlite3'
+            db_path = f'{protocol}{config.result_save_dir}/result.sqlite3'
         else:
             db_path = protocol + db_path
         db = records.Database(db_path)  # 不存在数据库时会新建一个数据库
@@ -185,6 +186,17 @@ class Database(object):
         self.query(f'delete from "{table_name}" where '
                    f'subdomain is null or valid == 0')
 
+    def deal_table(self, deal_table_name, backup_table_name):
+        """
+        收集任务完成时对表进行处理
+
+        :param str deal_table_name: 待处理的表名
+        :param str backup_table_name: 备份的表名
+        """
+        self.copy_table(deal_table_name, backup_table_name)
+        self.remove_invalid(deal_table_name)
+        self.deduplicate_subdomain(deal_table_name)
+
     def get_data(self, table_name):
         """
         获取表中的所有数据
@@ -205,11 +217,14 @@ class Database(object):
         table_name = table_name.replace('.', '_')
         query = f'select id, url, subdomain, port, ips, status, reason,' \
                 f'valid, new, title, banner from "{table_name}"'
-        if valid == 0 or valid == 1:
-            where = f' where valid = {valid}'
+        if valid:
+            where = f' where valid = 1'
             query += where
         logger.log('TRACE', f'获取{table_name}表中的所有数据')
         return self.query(query)
 
     def close(self):
+        """
+        关闭数据库连接
+        """
         self.conn.close()

oneforall/common/domain.py

@@ -38,7 +38,7 @@ class Domain(object):
 
         :return: 导出结果
         """
-        extract_cache_file = config.data_storage_path.joinpath('public_suffix_list.dat')
+        extract_cache_file = config.data_storage_dir.joinpath('public_suffix_list.dat')
         tldext = tldextract.TLDExtract(extract_cache_file)
         result = self.match()
         if result:

oneforall/common/module.py

@@ -226,7 +226,7 @@ class Module(object):
         if not config.save_module_result:
             return False
         logger.log('TRACE', f'将{self.source}模块发现的子域结果保存为json文件')
-        path = config.result_save_path.joinpath(self.domain, self.module)
+        path = config.result_save_dir.joinpath(self.domain, self.module)
         path.mkdir(parents=True, exist_ok=True)
         name = self.source + '.json'
         path = path.joinpath(name)

oneforall/common/request.py

@@ -5,9 +5,11 @@ import aiohttp
 import tqdm
 from aiohttp import ClientSession
 from bs4 import BeautifulSoup
+
 import config
 from common import utils
 from config import logger
+from common.database import Database
 
 
 def get_limit_conn():
@@ -220,6 +222,36 @@ async def bulk_request(datas, port):
     return new_datas
 
 
-def run_bulk_query(datas, port):
-    new_datas = asyncio.run(bulk_request(datas, port))
-    return new_datas
+def run_request(domain, data, port):
+    """
+    调用子域请求入口函数
+
+    :param str domain: 待请求的主域
+    :param list data: 待请求的子域数据
+    :param str port: 待请求的端口范围
+    :return: 请求后得到的结果列表
+    :rtype: list
+    """
+    loop = asyncio.get_event_loop()
+    asyncio.set_event_loop(loop)
+    request_coroutine = bulk_request(data, port)
+    data = loop.run_until_complete(request_coroutine)
+    # 在关闭事件循环前加入一小段延迟让底层连接得到关闭的缓冲时间
+    loop.run_until_complete(asyncio.sleep(0.25))
+    count = utils.count_valid(data)
+    logger.log('INFOR', f'经验证{domain}有效子域{count}个')
+    return data
+
+
+def save_data(name, data):
+    """
+    保存请求结果到数据库
+
+    :param str name: 保存表名
+    :param list data: 待保存的数据
+    """
+    db = Database()
+    db.drop_table(name)
+    db.create_table(name)
+    db.save_db(name, data, 'request')
+    db.close()

oneforall/common/resolve.py

@@ -10,6 +10,7 @@ from dns.resolver import Resolver
 
 import config
 from config import logger
+from common.database import Database
 
 
 def dns_resolver():
@@ -99,11 +100,27 @@ def update_data(data_list, results_dict):
     return data_list
 
 
-def init_worker():
-    signal.signal(signal.SIGINT, signal.SIG_IGN)
+def save_data(name, data):
+    """
+    保存解析结果到数据库
+
+    :param str name: 保存表名
+    :param list data: 待保存的数据
+    """
+    db = Database()
+    db.drop_table(name)
+    db.create_table(name)
+    db.save_db(name, data, 'resolve')
+    db.close()
 
 
-def query_progress(pr_queue, total):
+def resolve_progress(pr_queue, total):
+    """
+    解析进度
+
+    :param pr_queue: 进度队列
+    :param int total: 待解析的子域个数
+    """
     bar = tqdm.tqdm()
     bar.total = total
     bar.desc = 'Resolve Progress'
@@ -143,10 +160,9 @@ async def aio_resolve(subdomain_list, process_num, coroutine_num):
     m = Manager()
     pr_queue = m.Queue()
     loop = asyncio.get_event_loop()
-    loop.run_in_executor(None, query_progress, pr_queue, len(subdomain_list))
+    loop.run_in_executor(None, resolve_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
@@ -156,8 +172,8 @@ async def bulk_resolve(data_list):
     """
     批量解析A记录并返回解析结果
 
-    :param list data_list: 待查的数据列表
-    :return: 查询过得到的结果列表
+    :param list data_list: 待解析的数据列表
+    :return: 解析得到的结果列表
     """
     logger.log('INFOR', '正在异步查询子域的A记录')
     # semaphore = asyncio.Semaphore(config.limit_resolve_conn)
@@ -169,3 +185,19 @@ async def bulk_resolve(data_list):
     data_list = update_data(data_list, results_dict)
     logger.log('INFOR', '完成异步查询子域的A记录')
     return data_list
+
+
+def run_resolve(data):
+    """
+    调用子域解析入口函数
+
+    :param list data: 待解析的子域数据列表
+    :return: 解析得到的结果列表
+    :rtype: list
+    """
+    loop = asyncio.get_event_loop()
+    asyncio.set_event_loop(loop)
+    resolve_coroutine = bulk_resolve(data)
+    # 在关闭事件循环前加入一小段延迟让底层连接得到关闭的缓冲时间
+    loop.run_until_complete(asyncio.sleep(0.25))
+    return loop.run_until_complete(resolve_coroutine)

oneforall/common/utils.py

@@ -5,6 +5,7 @@ import time
 import random
 import ipaddress
 import platform
+
 import config
 from pathlib import Path
 from records import Record, RecordCollection
@@ -157,7 +158,7 @@ def check_path(path, name, format):
     :return: 保存路径
     """
     filename = f'{name}.{format}'
-    default_path = config.result_save_path.joinpath(filename)
+    default_path = config.result_save_dir.joinpath(filename)
     if path is None:
         path = default_path
     try:
@@ -250,30 +251,32 @@ def check_response(method, resp):
     return False
 
 
-def mark_subdomain(old_data, new_data):
+def mark_subdomain(old_data, now_data):
     """
     标记新增子域并返回新的数据集
 
-    :param old_data: 之前数据集
-    :param new_data: 现在数据集
-    :return: 已标记的新的数据集
+    :param list old_data: 之前子域数据
+    :param list now_data: 现在子域数据
+    :return: 标记后的的子域数据
+    :rtype: list
     """
     # 第一次收集子域的情况
+    mark_data = now_data.copy()
     if not old_data:
-        for index, item in enumerate(new_data):
+        for index, item in enumerate(mark_data):
             item['new'] = 1
-            new_data[index] = item
-        return new_data
+            mark_data[index] = item
+        return mark_data
     # 非第一次收集子域的情况
     old_subdomains = {item.get('subdomain') for item in old_data}
-    for index, item in enumerate(new_data):
+    for index, item in enumerate(mark_data):
         subdomain = item.get('subdomain')
         if subdomain in old_subdomains:
             item['new'] = 0
         else:
             item['new'] = 1
-        new_data[index] = item
-    return new_data
+        mark_data[index] = item
+    return mark_data
 
 
 def remove_string(string):
@@ -335,3 +338,7 @@ def python_version():
 
 def count_valid(data):
     return len(list(filter(lambda item: item.get('valid') == 1, data)))
+
+
+def get_subdomains(data):
+    return set(map(lambda item: item.get('subdomain'), data))