优化覆盖率

common/utils.py

@@ -153,7 +153,7 @@ def get_domains(target, targets=None):
     targets_domains = get_from_targets(targets)
     domains = list(target_domains.union(targets_domains))
     if targets_domains:
-        domains = sorted(domains, key=targets_domains.index) # 按照targets原本的index排序
+        domains = sorted(domains, key=targets_domains.index)  # 按照targets原本的index排序
     logger.log('INFOR', f'Get {len(domains)} domains')
     if not domains:
         logger.log('ERROR', f'Did not get a valid domain name')
@@ -324,18 +324,6 @@ def remove_invalid_string(string):
     return re.sub(r'[\000-\010]|[\013-\014]|[\016-\037]', r'', string)
 
 
-def check_value(values):
-    if not isinstance(values, dict):
-        return values
-    for key, value in values.items():
-        if value is None:
-            continue
-        if isinstance(value, str) and len(value) > 32767:
-            # Excel文件中单元格值长度不能超过32767
-            values[key] = value[:32767]
-    return values
-
-
 def export_all_results(path, name, format, datas):
     path = check_path(path, name, format)
     logger.log('ALERT', f'The subdomain result for all main domains: {path}')
@@ -347,8 +335,6 @@ def export_all_results(path, name, format, datas):
             row.pop('response')
         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)

dbexport.py

@@ -15,10 +15,6 @@ from common.database import Database
 from config.log import logger
 
 
-def domain_to_table(table):
-    return table.replace('.', '_') + "_now_result"
-
-
 def export(target, type='target', db=None, alive=False, limit=None, path=None, format='csv', show=False):
     """
     OneForAll export from database module
@@ -52,14 +48,7 @@ def export(target, type='target', db=None, alive=False, limit=None, path=None, f
                 rows = database.export_data(table_name, alive, limit)
                 if rows is None:
                     continue
-                format = utils.check_format(format, len(rows))
-                path = utils.check_path(path, target, format)
-                if show:
-                    print(rows.dataset)
-                data = rows.export(format)
-                utils.save_data(path, data)
-                logger.log('ALERT', f'The subdomain result for {table_name}: {path}')
-                data = rows.as_dict()
+                data = export_data(format, path, rows, show, table_name, target)
                 datas.extend(data)
         database.close()
         if len(domains) > 1:
@@ -67,16 +56,25 @@ def export(target, type='target', db=None, alive=False, limit=None, path=None, f
     elif type == 'table':
         database = Database(db)
         rows = database.export_data(target, alive, limit)
-        format = utils.check_format(format, len(rows))
-        path = utils.check_path(path, target, format)
-        if show:
-            print(rows.dataset)
-        data = rows.export(format)
+        data = export_data(format, path, rows, show, target, target)
         database.close()
-        utils.save_data(path, data)
-        logger.log('ALERT', f'The subdomain result for {target}: {path}')
-        data_dict = rows.as_dict()
-        return data_dict
+        return data
+
+
+def export_data(format, path, rows, show, table_name, target):
+    format = utils.check_format(format, len(rows))
+    path = utils.check_path(path, target, format)
+    if show:
+        print(rows.dataset)
+    data = rows.export(format)
+    utils.save_data(path, data)
+    logger.log('ALERT', f'The subdomain result for {table_name}: {path}')
+    data = rows.as_dict()
+    return data, format, path
+
+
+def domain_to_table(table):
+    return table.replace('.', '_') + "_now_result"
 
 
 if __name__ == '__main__':

test.py

@@ -6,11 +6,18 @@ Example
 """
 
 from oneforall import OneForAll
+from dbexport import export
 
-if __name__ == '__main__':
-    test = OneForAll(target='freebuf.com')
+
+def oneforall(target):
+    test = OneForAll(target=target)
     test.brute = True
     test.req = True
     test.takeover = True
     test.run()
-    result = test.datas
+
+
+if __name__ == '__main__':
+    TARGET = 'freebuf.com'
+    oneforall(target=TARGET)
+    export(target=TARGET)