添加当查询多个主域时结果自动合并导出功能

oneforall/common/request.py

@@ -111,36 +111,33 @@ def get_title(markup):
 
     title = soup.title
     if title:
-        return title.text.strip()
+        return title.text
 
     h1 = soup.h1
     if h1:
-        return h1.text.strip()
+        return h1.text
 
     h2 = soup.h2
     if h2:
-        return h2.text.strip()
+        return h2.text
 
     h3 = soup.h3
     if h2:
-        return h3.text.strip()
+        return h3.text
 
     desc = soup.find('meta', attrs={'name': 'description'})
     if desc:
-        return desc['content'].strip()
+        return desc['content']
 
     word = soup.find('meta', attrs={'name': 'keywords'})
     if word:
-        return word['content'].strip()
+        return word['content']
-
-    if len(markup) <= 200:
-        return markup.strip()
 
     text = soup.text
     if len(text) <= 200:
-        return text.strip()
+        return text
 
-    return None
+    return ''
 
 
 def request_callback(future, index, datas):
@@ -164,9 +161,10 @@ def request_callback(future, index, datas):
                               'Via': headers.get('Via'),
                               'X-Powered-By': headers.get('X-Powered-By')})
                 datas[index]['banner'] = banner[1:-1]
-                datas[index]['title'] = get_title(text)
+                title = get_title(text).strip()
+                datas[index]['title'] = utils.remove_string(title)
                 datas[index]['header'] = str(dict(headers))[1:-1]
-                datas[index]['response'] = text
+                datas[index]['response'] = utils.remove_string(text)
         else:
             datas[index]['reason'] = 'Something error'
             datas[index]['valid'] = 0

oneforall/common/utils.py

@@ -1,10 +1,12 @@
 # coding=utf-8
 import re
+import time
 import random
 import ipaddress
 import platform
 import config
 from pathlib import Path
+from records import Record, RecordCollection
 from common.domain import Domain
 from config import logger
 
@@ -144,7 +146,7 @@ def get_semaphore():
         return 800
 
 
-def check_dpath(dpath):
+def check_dpath(dpath=None):
     """
     检查目录路径
 
@@ -163,21 +165,26 @@ def check_dpath(dpath):
     return dpath
 
 
-def check_format(format):
+def check_format(format, line):
     """
     检查导出格式
 
     :param format: 传入的导出格式
+    :param line: 行数
     :return: 导出格式
     """
     formats = ['txt', 'rst', 'csv', 'tsv', 'json', 'yaml', 'html',
                'jira', 'xls', 'xlsx', 'dbf', 'latex', 'ods']
+    if format == 'xls' and line > 65000:
+        logger.log('ALERT', 'xls文件限制为最多65000行')
+        logger.log('ALERT', '使用xlsx格式导出')
+        return 'xlsx'
     if format in formats:
         return format
     else:
         logger.log('ALERT', f'不支持{format}格式导出')
         logger.log('ALERT', '默认使用csv格式导出')
-        return 'xls'
+        return 'csv'
 
 
 def save_data(fpath, data):
@@ -233,3 +240,46 @@ def mark_subdomain(old_data, new_data):
             item['new'] = 1
         new_data[index] = item
     return new_data
+
+
+def remove_string(string):
+    # Excel文件中单元格值不能直接存储以下非法字符
+    return re.sub(r'[\000-\010]|[\013-\014]|[\016-\037]', r'', string)
+
+
+def check_value(values):
+    for i, value in enumerate(values):
+        # Excel文件中单元格值长度不能超过32767
+        if value is None:
+            continue
+        if isinstance(value, str) and len(value) > 32767:
+            values[i] = value[:32767]
+    return values
+
+
+def export_all(format, datas):
+    line = len(datas)
+    format = check_format(format, line)
+    dpath = check_dpath()
+    timestamp = get_timestamp()
+    fpath = dpath.joinpath(f'all_subdomain_{timestamp}.{format}')
+    row_list = list()
+    for row in datas:
+        row.pop('header')
+        row.pop('response')
+        row.pop('module')
+        row.pop('source')
+        row.pop('elapsed')
+        row.pop('count')
+        keys = row.keys()
+        values = row.values()
+        if format in {'xls', 'xlsx'}:
+            values = check_value(values)
+        row_list.append(Record(keys, values))
+    rows = RecordCollection(iter(row_list))
+    content = rows.export(format)
+    save_data(fpath, content)
+
+
+def get_timestamp():
+    return int(time.time())

oneforall/dbexport.py

@@ -34,10 +34,12 @@ def export(table, db=None, valid=None, dpath=None, format='csv', show=False):
     :param str dpath:    导出目录(默认None)
     :param bool show:   终端显示导出数据(默认False)
     """
-    format = utils.check_format(format)
+
     dpath = utils.check_dpath(dpath)
     database = Database(db)
     rows = database.export_data(table, valid)  # 意外情况导出全部子域
+    line = len(rows)
+    format = utils.check_format(format, line)
     if show:
         print(rows.dataset)
     if format == 'txt':

oneforall/oneforall.py

@@ -80,6 +80,7 @@ class OneForAll(object):
         self.domains = set()
         self.domain = str()
         self.data = list()
+        self.datas = list()
         self.brute = brute
         self.dns = dns
         self.req = req
@@ -159,6 +160,7 @@ class OneForAll(object):
         # 请求子域
         task = request.bulk_get_request(self.data, self.port)
         self.data = loop.run_until_complete(task)
+        self.datas.extend(self.data)
         # 在关闭事件循环前加入一小段延迟让底层连接得到关闭的缓冲时间
         loop.run_until_complete(asyncio.sleep(0.25))
 
@@ -190,6 +192,8 @@ class OneForAll(object):
         if self.domains:
             for self.domain in self.domains:
                 self.main()
+            if len(self.domains) >= 2:
+                utils.export_all(self.format, self.datas)
         else:
             logger.log('FATAL', f'获取域名失败')
         logger.log('INFOR', f'结束运行OneForAll')

oneforall/takeover.py

@@ -83,8 +83,8 @@ class Takeover(Module):
             data = str(self.results)
         else:
             data = self.results.export(self.format)
-        ts = int(time.time())
+        timestamp = utils.get_timestamp()
-        fpath = self.dpath.joinpath(f'takeover_{ts}.{self.format}')
+        fpath = self.dpath.joinpath(f'takeover_{timestamp}.{self.format}')
         utils.save_data(fpath, data)
 
     def compare(self, subdomain, cname, responses):