解决批量获取子域问题

oneforall/aiobrute.py

@@ -140,8 +140,8 @@ class AIOBrute(Module):
         python3 aiobrute.py --target m.{fuzz}.a.bz --fuzz True --rule [a-z] run
 
     Note:
-        参数segment的设置受CPU性能，网络带宽，运营商限制等问题影响，默认设置500个子域为任务组，
+        参数segment的设置受CPU性能，网络带宽，运营商限制等限制，默认500个子域为任务组，
-        当你觉得你的环境不受以上因素影响，当前爆破速度较慢，那么强烈建议根据字典大小调整大小：
+        当你的环境不受以上因素影响，当前爆破速度较慢，那么强烈建议根据字典大小调整大小：
         十万字典建议设置为5000，百万字典设置为50000
         参数valid可选值1，0，None，分别表示导出有效，无效，全部子域
         参数format可选格式：'csv', 'tsv', 'json', 'yaml', 'html', 'xls', 'xlsx',
@@ -310,7 +310,7 @@ class AIOBrute(Module):
                     self.path = config.result_save_path.joinpath(name)
                 dbexport.export(self.domain,
                                 valid=self.valid,
-                                path=self.path,
+                                dpath=self.path,
                                 format=self.format,
                                 show=self.show)
 

oneforall/collect.py

@@ -81,7 +81,7 @@ class Collect(object):
             if not self.path:
                 name = f'{self.domain}.{self.format}'
                 self.path = config.result_save_path.joinpath(name)
-            dbexport.export(self.domain, path=self.path, format=self.format)
+            dbexport.export(self.domain, dpath=self.path, format=self.format)
         end = time.time()
         self.elapsed = round(end - start, 1)
 

oneforall/common/database.py

@@ -13,10 +13,10 @@ from config import logger
 
 class Database(object):
     def __init__(self, db_path=None):
-        self.conn = self.get_connection(db_path)
+        self.conn = self.get_conn(db_path)
 
     @staticmethod
-    def get_connection(db_path):
+    def get_conn(db_path):
         """
         获取数据库对象
 
@@ -205,3 +205,6 @@ class Database(object):
             logger.log('ERROR', e)
         else:
             return rows
+
+    def close(self):
+        self.conn.close()

oneforall/common/module.py

@@ -274,4 +274,5 @@ class Module(object):
         source, results = self.results
         # 将结果存入数据库中
         db.save_db(self.domain, results, source)
+        db.close()
         lock.release()

oneforall/common/utils.py

@@ -109,9 +109,11 @@ def get_domains(target):
     :param set or str target:
     :return: 域名集合
     """
-    domains = set()
+    domains = list()
     logger.log('INFOR', f'正在获取域名')
-    if isinstance(target, set):
+    if isinstance(target, (set, tuple)):
+        domains = list(target)
+    elif isinstance(target, list):
         domains = target
     elif isinstance(target, str):
         path = pathlib.Path(target)
@@ -120,9 +122,9 @@ def get_domains(target):
                 for line in file:
                     domain = Domain(line.strip()).match()
                     if domain:
-                        domains.add(domain)
+                        domains.append(domain)
-        if Domain(target).match():
+        elif Domain(target).match():
-            domains = {target}
+            domains = [target]
     logger.log('INFOR', f'获取到{len(domains)}个域名')
     return domains
 

oneforall/config.py

@@ -33,7 +33,7 @@ brute_coroutine_num = 128  # 爆破时每个进程下的协程数(不宜大于10
 # 爆破所使用的字典路径 默认data/subdomains.txt
 brute_wordlist_path = data_storage_path.joinpath('subdomains.txt')
 brute_task_segment = 500
-# 参数segment的设置受CPU性能，网络带宽，运营商限制等问题影响，默认设置500个子域为一任务组，
+# 参数segment的设置受CPU性能，网络带宽，运营商限制等限制，默认500个子域为一任务组，
 # 当你觉得你的环境不受以上因素影响，当前爆破速度较慢，那么强烈建议根据字典大小调整大小：
 # 十万字典建议设置为5000，百万字典设置为50000
 enable_recursive_brute = False  # 是否使用递归爆破(默认禁用)
@@ -117,7 +117,7 @@ censys_api_secret = ''
 # 免费的API有效期只有1个月，到期之后可以再次生成，每月可以查询250次。
 binaryedge_api = ''
 
-# Binaryedge可以免费注册获取API：http://api.chinaz.com/ApiDetails/Alexa
+# Chinaz可以免费注册获取API：http://api.chinaz.com/ApiDetails/Alexa
 chinaz_api = ''
 
 # Bing可以免费注册获取API：https://azure.microsoft.com/zh-cn/services/
@@ -175,13 +175,16 @@ subdomains_common = {'i', 'w', 'm', 'en', 'us', 'zh', 'w3', 'app', 'bbs',
                      'store', 'mobile'}
 
 # 日志配置
+# 终端日志输出格式
 stdout_fmt = '<cyan>{time:HH:mm:ss,SSS}</cyan> ' \
           '[<level>{level: <5}</level>] ' \
           '<blue>{module}</blue>:<cyan>{line}</cyan> - ' \
           '<level>{message}</level>'
+# 日志文件记录格式
 logfile_fmt = '<light-green>{time:YYYY-MM-DD HH:mm:ss,SSS}</light-green> ' \
           '[<level>{level: <5}</level>] ' \
-          '<cyan>{process.name}</cyan>:<cyan>{thread.name: <10}</cyan> | ' \
+          '<cyan>{process.name}({process.id})</cyan>:' \
+          '<cyan>{thread.name: <10}({thread.id: <5})</cyan> | ' \
           '<blue>{module}</blue>.<blue>{function}</blue>:' \
           '<blue>{line}</blue> - <level>{message}</level>'
 
@@ -199,5 +202,5 @@ if not os.environ.get('PYTHONIOENCODING'):  # 设置编码
     os.environ['PYTHONIOENCODING'] = 'utf-8'
 
 logger.add(sys.stderr, level='INFOR', format=stdout_fmt, enqueue=True)
-logger.add(log_path, level='TRACE', format=logfile_fmt, enqueue=True,
+logger.add(log_path, level='DEBUG', format=logfile_fmt, enqueue=True,
            encoding='utf-8')

oneforall/dbexport.py

@@ -8,30 +8,32 @@ OneForAll数据库导出模块
 :license: GNU General Public License v3.0, see LICENSE for more details.
 """
 
+from pathlib import Path
 import fire
+import config
 from common.database import Database
 from config import logger
 
 
-def export(table, db=None, valid=None, path=None, format='xlsx', show=False):
+def export(table, db=None, valid=None, dpath=None, format='xlsx', show=False):
     """
     OneForAll数据库导出模块
 
     Example:
-        python3 dbexport.py --table name --format csv --path= ./result.csv
+        python3 dbexport.py --table name --format csv --dir= ./result.csv
         python3 dbexport.py --db result.db --table name --show False
 
     Note:
         参数port可选值有'small', 'medium', 'large', 'xlarge'，详见config.py配置
         参数format可选格式有'csv', 'tsv', 'json', 'yaml', 'html', 'xls', 'xlsx',
                          'dbf', 'latex', 'ods'
-        参数path为None会根据format参数和域名名称在项目结果目录生成相应文件
+        参数dir为None默认使用OneForAll结果目录
 
     :param str table:   要导出的表
     :param str db:      要导出的数据库路径(默认为results/result.sqlite3)
     :param int valid:   导出子域的有效性(默认None)
     :param str format:  导出格式(默认xlsx)
-    :param str path:    导出路径(默认None)
+    :param str dpath:    导出目录(默认None)
     :param bool show:   终端显示导出数据(默认False)
     """
     formats = ['csv', 'tsv', 'json', 'yaml', 'html',
@@ -39,6 +41,17 @@ def export(table, db=None, valid=None, path=None, format='xlsx', show=False):
     if format not in formats:
         logger.log('FATAL', f'不支持{format}格式导出')
         return
+    if dpath is None:
+        dpath = config.result_save_path
+    if isinstance(dpath, str):
+        dpath = Path(dpath)
+    if not dpath.is_dir():
+        logger.log('FATAL', f'{dpath}不是目录')
+        return
+    if not dpath.exists():
+        logger.log('ALERT', f'不存在{dpath}将会新建此目录')
+        dpath.mkdir(parents=True, exist_ok=True)
+
     database = Database(db)
     if valid is None:
         rows = database.get_data(table)
@@ -48,20 +61,20 @@ def export(table, db=None, valid=None, path=None, format='xlsx', show=False):
         rows = database.get_data(table)  # 意外情况导出全部子域
     if show:
         print(rows.dataset)
-    if not path:
-        path = 'export.' + format
     logger.log('INFOR', f'正在将数据库中{table}表导出')
     data = rows.export(format)
+    database.close()
+    fpath = dpath.joinpath(f'{table}.{format}')
     try:
-        with open(path, 'w') as file:
+        with open(fpath, 'w') as file:
             file.write(data)
             logger.log('INFOR', '成功完成导出')
-            logger.log('INFOR', path)
+            logger.log('INFOR', fpath)
     except TypeError:
-        with open(path, 'wb') as file:
+        with open(fpath, 'wb') as file:
             file.write(data)
             logger.log('INFOR', '成功完成导出')
-            logger.log('INFOR', path)
+            logger.log('INFOR', fpath)
     except Exception as e:
         logger.log('ERROR', e)
 

oneforall/domains.txt

@@ -1 +1,2 @@
 example.com
+freebuf.com

oneforall/oneforall.py

@@ -60,7 +60,6 @@ class OneForAll(object):
         参数port可选值有'small', 'medium', 'large', 'xlarge'，详见config.py配置
         参数format可选格式有'csv', 'tsv', 'json', 'yaml', 'html', 'xls', 'xlsx',
                           'dbf', 'latex', 'ods'
-        参数path为None会根据format参数和域名名称在项目结果目录生成相应文件
 
     :param str target:  单个域名或者每行一个域名的文件路径(必需参数)
     :param bool brute:  使用爆破模块(默认False)
@@ -68,11 +67,10 @@ class OneForAll(object):
     :param str port:    请求验证的端口范围(默认medium)
     :param int valid:   导出子域的有效性(默认1)
     :param str format:  导出格式(默认xlsx)
-    :param str path:    导出路径(默认None)
     :param bool show:   终端显示导出数据(默认False)
     """
     def __init__(self, target, brute=None, verify=None, port='medium', valid=1,
-                 path=None, format='xlsx', show=False):
+                 format='xlsx', show=False):
         self.target = target
         self.port = port
         self.domains = set()
@@ -81,7 +79,6 @@ class OneForAll(object):
         self.brute = brute
         self.verify = verify
         self.valid = valid
-        self.path = path
         self.format = format
         self.show = show
 
@@ -91,9 +88,6 @@ class OneForAll(object):
         if self.verify is None:
             self.verify = config.enable_verify_subdomain
         rename_table = self.domain + '_last'
-        if not self.path:
-            name = f'{self.domain}.{self.format}'
-            self.path = config.result_save_path.joinpath(name)
         collect = Collect(self.domain, export=False)
         collect.run()
         if self.brute:
@@ -109,8 +103,8 @@ class OneForAll(object):
         if not self.verify:
             # 数据库导出
             self.valid = None
-            dbexport.export(self.domain, db.conn, self.valid, self.path,
+            dbexport.export(self.domain, valid=self.valid, format=self.format,
-                            self.format, self.show)
+                            show=self.show)
             db.drop_table(rename_table)
             db.rename_table(self.domain, rename_table)
             return
@@ -134,16 +128,16 @@ class OneForAll(object):
         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)
 
         # 数据库导出
-        dbexport.export(self.domain, db.conn, self.valid, self.path,
+        dbexport.export(self.domain, valid=self.valid, format=self.format,
-                        self.format, self.show)
+                        show=self.show)
         db.drop_table(rename_table)
         db.rename_table(self.domain, rename_table)
+        db.close()
 
     def run(self):
         print(banner)
@@ -162,3 +156,4 @@ class OneForAll(object):
 if __name__ == '__main__':
     fire.Fire(OneForAll)
     # OneForAll('example.com').run()
+    # OneForAll('./domains.txt').run()