重构

oneforall/aiobrute.py

@@ -20,6 +20,7 @@ import fire
 import tqdm
 
 import config
+import dbexport
 from common import resolve, utils
 from common.module import Module
 from common.database import Database
@@ -142,6 +143,10 @@ class AIOBrute(Module):
         参数segment的设置受CPU性能，网络带宽，运营商限制等问题影响，默认设置500个子域为任务组，
         当你觉得你的环境不受以上因素影响，当前爆破速度较慢，那么强烈建议根据字典大小调整大小：
         十万字典建议设置为5000，百万字典设置为50000
+        参数valid可选值1，0，None，分别表示导出有效，无效，全部子域
+        参数format可选格式：'csv', 'tsv', 'json', 'yaml', 'html', 'xls', 'xlsx',
+                         'dbf', 'latex', 'ods'
+        参数path为None会根据format参数和域名名称在项目结果目录生成相应文件
 
     :param str target:       单个域名或者每行一个域名的文件路径
     :param int processes:    爆破的进程数(默认CPU核心数)
@@ -153,11 +158,17 @@ class AIOBrute(Module):
     :param str namelist:     指定递归爆破所使用的字典路径(默认使用config.py配置)
     :param bool fuzz:        是否使用fuzz模式进行爆破(默认False，开启须指定fuzz正则规则)
     :param str rule:         fuzz模式使用的正则规则(默认使用config.py配置)
+    :param bool export:      是否导出爆破结果(默认True)
+    :param int valid:        导出子域的有效性(默认None)
+    :param str format:       导出格式(默认xlsx)
+    :param str path:         导出路径(默认None)
+    :param
     """
 
     def __init__(self, target, processes=None, coroutine=64, wordlist=None,
                  segment=500, recursive=False, depth=2, namelist=None,
-                 fuzz=False, rule=None):
+                 fuzz=False, rule=None, export=True, valid=None, format='xlsx',
+                 path=None):
         Module.__init__(self)
         self.domains = set()
         self.domain = str()
@@ -173,6 +184,10 @@ class AIOBrute(Module):
         self.recursive_namelist = namelist or config.recursive_namelist_path
         self.fuzz = fuzz or config.enable_fuzz
         self.rule = rule or config.fuzz_rule
+        self.export = export
+        self.valid = valid
+        self.format = format
+        self.path = path
         self.nameservers = config.resolver_nameservers
         self.ips_times = dict()  # IP集合出现次数
         self.enable_wildcard = False  # 当前域名是否使用泛解析
@@ -279,9 +294,6 @@ class AIOBrute(Module):
                 source, results = rx_queue.get()
                 # 将结果存入数据库中
                 db.save_db(self.domain, results, source)
-            db.copy_table(self.domain)
-            db.deduplicate_subdomain(self.domain)
-            db.remove_invalid(self.domain)
 
             end = time.time()
             self.elapsed = round(end - start, 1)
@@ -291,6 +303,15 @@ class AIOBrute(Module):
                                 f'发现{self.domain}的域名{length}个')
             logger.log('DEBUG', f'{self.source}模块发现{self.domain}的域名:\n'
                                 f'{self.subdomains}')
+            # 数据库导出
+            if self.export:
+                if not self.path:
+                    name = f'{self.domain}_brute.{self.format}'
+                    self.path = config.result_save_path.joinpath(name)
+                dbexport.export(self.domain,
+                                valid=self.valid,
+                                path=self.path,
+                                format=self.format)
 
 
 def do(domain, result):  # 统一入口名字 方便多线程调用

oneforall/collect.py

@@ -1,7 +1,3 @@
-# coding=utf-8
-"""
-被动收集类
-"""
 import time
 import threading
 import importlib
@@ -27,7 +23,6 @@ class Collect(object):
     def get_mod(self):
         """
         获取要运行的模块
-        :return: None
         """
         if config.enable_all_module:
             # modules = ['brute', 'certificates', 'crawl',
@@ -76,12 +71,11 @@ class Collect(object):
         for thread in threads:
             thread.join()
 
-        db = Database()
-        db.create_table(self.domain)
-        db.copy_table(self.domain)
-        db.deduplicate_subdomain(self.domain)
-        db.remove_invalid(self.domain)
-        # conn.close()
+        # db = Database()
+        # db.create_table(self.domain)
+        # db.copy_table(self.domain, self.domain+'_collect')
+        # db.remove_invalid(self.domain)
+        # db.deduplicate_subdomain(self.domain)
         # 数据库导出
         if self.export:
             if not self.path:

oneforall/common/database.py

@@ -75,27 +75,30 @@ class Database(object):
         if results:
             try:
                 self.conn.bulk_query(
-                    f'insert into "{table_name}" (id, url, subdomain, port, ips, status,'
-                    f'reason, valid, title, banner, module, source, elapsed, count)'
-                    f'values (:id, :url, :subdomain, :port, :ips, :status, :reason, :valid,'
-                    f':title, :banner, :module, :source, :elapsed, :count)',
+                    f'insert into "{table_name}" ('
+                    f'id, url, subdomain, port, ips, status, reason, valid,'
+                    f'title, banner, module, source, elapsed, count)'
+                    f'values (:id, :url, :subdomain, :port, :ips, :status,'
+                    f':reason, :valid, :title, :banner, :module, :source,'
+                    f':elapsed, :count)',
                     results)
             except Exception as e:
                 logger.log('ERROR', e)
 
-    def copy_table(self, table_name):
+    def copy_table(self, table_name, bak_table_name):
         """
         复制表创建备份
 
         :param str table_name: 表名
+        :param str bak_table_name: 新表名
         """
         table_name = table_name.replace('.', '_')
-        new_table_name = table_name + '_bak'
-        logger.log('DEBUG', f'正在将{table_name}表复制到{new_table_name}新表')
+        bak_table_name = bak_table_name.replace('.', '_')
+        logger.log('DEBUG', f'正在将{table_name}表复制到{bak_table_name}新表')
         try:
-            self.conn.query(f'drop table if exists "{new_table_name}"')
-            self.conn.query(
-                f'create table "{new_table_name}" as select * from "{table_name}"')
+            self.conn.query(f'drop table if exists "{bak_table_name}"')
+            self.conn.query(f'create table "{bak_table_name}" '
+                            f'as select * from "{table_name}"')
         except Exception as e:
             logger.log('ERROR', e)
 
@@ -112,9 +115,38 @@ class Database(object):
         except Exception as e:
             logger.log('ERROR', e)
 
+    def drop_table(self, table_name):
+        """
+        删除表
+
+        :param str table_name: 表名
+        """
+        table_name = table_name.replace('.', '_')
+        logger.log('DEBUG', f'正在删除{table_name}表')
+        try:
+            self.conn.query(f'drop table if exists "{table_name}"')
+        except Exception as e:
+            logger.log('ERROR', e)
+
+    def rename_table(self, table_name, new_table_name):
+        """
+        复制表创建备份
+
+        :param str table_name: 表名
+        :param str new_table_name: 新表名
+        """
+        table_name = table_name.replace('.', '_')
+        new_table_name = new_table_name.replace('.', '_')
+        logger.log('DEBUG', f'正在将{table_name}表重命名为{table_name}表')
+        try:
+            self.conn.query(f'alter table "{table_name}" '
+                            f'rename to "{new_table_name}"')
+        except Exception as e:
+            logger.log('ERROR', e)
+
     def deduplicate_subdomain(self, table_name):
         """
-        去重表中的子域并删除空值和无效值
+        去重表中的子域
 
         :param str table_name: 表名
         """
@@ -122,7 +154,8 @@ class Database(object):
         logger.log('DEBUG', f'正在去重{table_name}表中的子域')
         try:
             self.conn.query(
-                f'delete from "{table_name}" where id not in (select min(id) from "{table_name}" group by subdomain)')
+                f'delete from "{table_name}" where id not in (select min(id) '
+                f'from "{table_name}" group by subdomain)')
         except Exception as e:
             logger.log('ERROR', e)
 
@@ -136,7 +169,8 @@ class Database(object):
         logger.log('DEBUG', f'正在去除{table_name}表中的无效子域')
         try:
             self.conn.query(
-                f'delete from "{table_name}" where subdomain is null or valid == 0')
+                f'delete from "{table_name}" where '
+                f'subdomain is null or valid == 0')
         except Exception as e:
             logger.log('ERROR', e)
 

oneforall/common/module.py

@@ -227,7 +227,7 @@ class Module(object):
                       'ips': None,
                       'status': None,
                       'reason': None,
-                      'valid': 1,
+                      'valid': None,
                       'title': None,
                       'banner': None,
                       'module': self.module,
@@ -247,7 +247,7 @@ class Module(object):
                           'ips': ips,
                           'status': None,
                           'reason': None,
-                          'valid': 1,
+                          'valid': None,
                           'title': None,
                           'banner': None,
                           'module': self.module,

oneforall/common/request.py

@@ -32,7 +32,8 @@ def gen_new_datas(datas, ports):
     new_datas = []
     protocols = ['http://']
     for data in datas:
-        if data.get('valid'):  # 有效的子域才进行http请求探测
+        valid = data.get('valid')
+        if valid is None:  # 子域有效性未知的才进行http请求探测
             subdomain = data.get('subdomain')
             for port in ports:
                 for protocol in protocols:
@@ -83,6 +84,7 @@ def deal_results(datas, results):
             if resp.status == 400 or resp.status >= 500:
                 datas[index]['valid'] = 0
             else:
+                datas[index]['valid'] = 1
                 headers = resp.headers
                 banner = str({'Server': headers.get('Server'),
                               'Via': headers.get('Via'),

oneforall/common/resolve.py

@@ -78,7 +78,6 @@ def resolve_callback(future, index, datas):
                 datas[index]['ips'] = str(ips)
             else:
                 datas[index]['ips'] = 'No answers'
-                datas[index]['valid'] = 0
 
 
 async def bulk_query_a(datas):

oneforall/config.py

@@ -44,7 +44,7 @@ fuzz_rule = ''  # fuzz域名的正则 示例：[a-z][0-9] 第一位是字母 第
 ips_appear_maximum = 10  # 同一IP集合出现次数超过10认为是泛解析
 
 # 代理设置
-enable_proxy = False  # 是否使用代理 全局开关
+enable_proxy = False  # 是否使用代理(全局开关)
 proxy_all_module = False  # 代理所有模块
 proxy_partial_module = ['GoogleQuery', 'AskSearch', 'DuckDuckGoSearch',
                         'GoogleAPISearch', 'GoogleSearch', 'YahooSearch',

oneforall/dbexport.py

@@ -18,12 +18,13 @@ def export(table, db=None, valid=None, path=None, format='xlsx', output=False):
     OneForAll数据库导出模块
 
     Example:
-        python dbexport.py --db result.db --table name --format csv --output False
-        python dbexport.py --db result.db --table name --format csv --path= ./result.csv
+        python dbexport.py --table name --format csv --path= ./result.csv
+        python dbexport.py --db result.db --table name --output False
 
     Note:
         参数valid可选值1，0，None，分别表示导出有效，无效，全部子域
-        参数format可选格式：'csv', 'tsv', 'json', 'yaml', 'html', 'xls', 'xlsx', 'dbf', 'latex', 'ods'
+        参数format可选格式：'csv', 'tsv', 'json', 'yaml', 'html', 'xls', 'xlsx',
+                         'dbf', 'latex', 'ods'
         参数path为None会根据format参数和域名名称在项目结果目录生成相应文件
 
     :param str table:   要导出的表

oneforall/oneforall.py

@@ -66,12 +66,12 @@ class OneForAll(object):
     :param str path:    导出路径(默认None)
     :param bool output: 是否将导出数据输出到终端(默认False)
     """
-    def __init__(self, target, brute=False, port='medium', valid=1, path=None,
+    def __init__(self, target, brute=None, port='medium', valid=1, path=None,
                  format='xlsx', output=False):
         self.target = target
         self.port = port
         self.domains = set()
-        self.domain = ''
+        self.domain = str()
         self.datas = list()
         self.brute = brute or config.enable_brute_module
         self.valid = valid
@@ -79,6 +79,53 @@ class OneForAll(object):
         self.format = format
         self.output = output
 
+    def main(self):
+        collect = Collect(self.domain, export=False)
+        collect.run()
+        if self.brute:
+            # 由于爆破会有大量dns解析请求 并发爆破可能会导致其他任务中的网络请求异常
+            brute = AIOBrute(self.domain, export=False)
+            brute.run()
+
+        db = Database()
+        db.copy_table(self.domain, self.domain+'_ori')
+        db.remove_invalid(self.domain)
+        db.deduplicate_subdomain(self.domain)
+        self.datas = db.get_data(self.domain).as_dict()
+        loop = asyncio.get_event_loop()
+        asyncio.set_event_loop(loop)
+
+        # 解析域名地址
+        task = resolve.bulk_query_a(self.datas)
+        self.datas = loop.run_until_complete(task)
+
+        # 保存解析结果
+        resolve_table = self.domain + '_res'
+        db.drop_table(resolve_table)
+        db.create_table(resolve_table)
+        db.save_db(resolve_table, self.datas, 'resolve')
+
+        # 请求域名地址
+        task = request.bulk_get_request(self.datas, self.port)
+        self.datas = loop.run_until_complete(task)
+        # 在关闭事件循环前加入一小段延迟让底层连接得到关闭的缓冲时间
+        loop.run_until_complete(asyncio.sleep(0.25))
+        loop.close()
+
+        db.clear_table(self.domain)
+        db.save_db(self.domain, self.datas)
+
+        # 数据库导出
+        if not self.path:
+            name = f'{self.domain}.{self.format}'
+            self.path = config.result_save_path.joinpath(name)
+        dbexport.export(self.domain, db.conn, self.valid, self.path,
+                        self.format, self.output)
+
+        rename_table = self.domain + '_last'
+        db.drop_table(rename_table)
+        db.rename_table(self.domain, rename_table)
+
     def run(self):
         print(banner)
         dt = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
@@ -87,31 +134,7 @@ class OneForAll(object):
         self.domains = utils.get_domains(self.target)
         if self.domains:
             for self.domain in self.domains:
-                collect = Collect(self.domain, export=False)
-                collect.run()
-                if self.brute:
-                    # 由于爆破会有大量dns解析请求 并发常常会导致其他任务中的网络请求超时
-                    brute = AIOBrute(self.domain)
-                    brute.run()
-                db = Database()
-                self.datas = db.get_data(self.domain).as_dict()
-                loop = asyncio.get_event_loop()
-                asyncio.set_event_loop(loop)
-                task = resolve.bulk_query_a(self.datas)
-                self.datas = loop.run_until_complete(task)
-                task = request.bulk_get_request(self.datas, self.port)
-                self.datas = loop.run_until_complete(task)
-                # 在关闭事件循环前加入一小段延迟让底层连接得到关闭的缓冲时间
-                loop.run_until_complete(asyncio.sleep(0.25))
-                loop.close()
-                db.clear_table(self.domain)
-                db.save_db(self.domain, self.datas)
-                # 数据库导出
-                if not self.path:
-                    name = f'{self.domain}.{self.format}'
-                    self.path = config.result_save_path.joinpath(name)
-                dbexport.export(self.domain, db.conn, self.valid, self.path,
-                                self.format, self.output)
+                self.main()
         else:
             logger.log('FATAL', f'获取域名失败')
         logger.log('INFOR', f'结束运行OneForAll')