优化导出

brute.py

@@ -18,7 +18,7 @@ import tenacity
 from dns.exception import Timeout
 from dns.resolver import NXDOMAIN, YXDOMAIN, NoAnswer, NoNameservers
 
-import dbexport
+import export
 from common import utils
 from common import similarity
 from config import settings
@@ -517,7 +517,7 @@ def is_valid_subdomain(ip, ttl, times, wc_ips, wc_ttl, cname):
 
 def save_brute_dict(dict_path, dict_set):
     dict_data = '\n'.join(dict_set)
-    if not utils.save_data(dict_path, dict_data):
+    if not utils.save_to_file(dict_path, dict_data):
         logger.log('FATAL', 'Saving dictionary error')
         exit(1)
 
@@ -758,8 +758,7 @@ class Brute(Module):
                 self.path = settings.result_save_dir.joinpath(name)
             # 数据库导出
             if self.export:
-                dbexport.export(self.domain,
+                export.export_data(self.domain,
-                                type='table',
                                    alive=self.alive,
                                    limit='resolve',
                                    path=self.path,

common/resolve.py

@@ -58,13 +58,13 @@ def save_db(name, data):
     :param list data: data to be saved
     """
     logger.log('INFOR', f'Saving resolved results')
-    utils.save_db(name, data, 'resolve')
+    utils.save_to_db(name, data, 'resolve')
 
 
 def save_subdomains(save_path, subdomain_list):
     logger.log('DEBUG', f'Saving resolved subdomain')
     subdomain_data = '\n'.join(subdomain_list)
-    if not utils.save_data(save_path, subdomain_data):
+    if not utils.save_to_file(save_path, subdomain_data):
         logger.log('FATAL', 'Save resolved subdomain error')
         exit(1)
 

common/utils.py

@@ -163,7 +163,6 @@ def get_domains(target, targets=None):
     domains = list(target_domains.union(targets_domains))
     if targets_domains:
         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')
     logger.log('DEBUG', f'The obtained domains \n{domains}')
@@ -225,7 +224,7 @@ def load_json(path):
         return json.load(fp)
 
 
-def save_db(name, data, module):
+def save_to_db(name, data, module):
     """
     Save request results to database
 
@@ -240,7 +239,7 @@ def save_db(name, data, module):
     db.close()
 
 
-def save_data(path, data):
+def save_to_file(path, data):
     """
     保存数据到文件
 
@@ -261,21 +260,6 @@ def save_data(path, data):
         return False
 
 
-def remove_data(path):
-    """
-    删除保存数据的文件
-
-    :param path: 路径
-    :return: 删除成功与否
-    """
-    try:
-        path.unlink()
-    except Exception as e:
-        logger.log('ERROR', e.args)
-        return False
-    return True
-
-
 def check_response(method, resp):
     """
     检查响应 输出非正常响应返回json的信息
@@ -346,7 +330,7 @@ def export_all_results(path, name, fmt, datas):
         row_list.append(Record(keys, values))
     rows = RecordCollection(iter(row_list))
     content = rows.export(fmt)
-    save_data(path, content)
+    save_to_file(path, content)
 
 
 def export_all_subdomains(alive, path, name, datas):
@@ -362,7 +346,7 @@ def export_all_subdomains(alive, path, name, datas):
         else:
             subdomains.add(subdomain)
     data = '\n'.join(subdomains)
-    save_data(path, data)
+    save_to_file(path, data)
 
 
 def export_all(alive, fmt, path, datas):
@@ -796,5 +780,5 @@ def get_ns_path(in_china=None, enable_wildcard=None, ns_ip_list=None):
         return path
     path = settings.authoritative_dns_path
     ns_data = '\n'.join(ns_ip_list)
-    save_data(path, ns_data)
+    save_to_file(path, ns_data)
     return path

dbexport.py

@@ -1,81 +0,0 @@
-#!/usr/bin/python3
-# coding=utf-8
-
-"""
-OneForAll export from database module
-
-:copyright: Copyright (c) 2019, Jing Ling. All rights reserved.
-:license: GNU General Public License v3.0, see LICENSE for more details.
-"""
-
-import fire
-
-from common import utils
-from common.database import Database
-from config.log import logger
-
-
-def export(target, type='target', db=None, alive=False, limit=None, path=None, fmt='csv', show=False):
-    """
-    OneForAll export from database module
-
-    Example:
-        python3 dbexport.py --target name --fmt csv --dir= ./result.csv
-        python3 dbexport.py --db result.db --target name --show False
-        python3 dbexport.py --target table_name --tb True --show False
-
-    Note:
-        --fmt csv/json (result format)
-        --path   Result directory (default directory is ./results)
-
-    :param str  target:  Table to be exported
-    :param str  type:    Type of target
-    :param str  db:      Database path to be exported (default ./results/result.sqlite3)
-    :param bool alive:   Only export the results of alive subdomains (default False)
-    :param str  limit:   Export limit (default None)
-    :param str  fmt:     Result format (default csv)
-    :param str  path:    Result directory (default None)
-    :param bool show:    Displays the exported data in terminal (default False)
-    """
-
-    if type == 'target':
-        database = Database(db)
-        domains = utils.get_domains(target)
-        datas = []
-        if domains:
-            for domain in domains:
-                table_name = domain_to_table(domain)
-                rows = database.export_data(table_name, alive, limit)
-                if rows is None:
-                    continue
-                data = export_data(fmt, path, rows, show, table_name, target)
-                datas.extend(data)
-        database.close()
-        if len(domains) > 1:
-            utils.export_all(alive, fmt, path, datas)
-    elif type == 'table':
-        database = Database(db)
-        rows = database.export_data(target, alive, limit)
-        data, _, _ = export_data(fmt, path, rows, show, target, target)
-        database.close()
-        return data
-
-
-def export_data(fmt, path, rows, show, table_name, target):
-    fmt = utils.check_format(fmt)
-    path = utils.check_path(path, target, fmt)
-    if show:
-        print(rows.dataset)
-    data = rows.export(fmt)
-    utils.save_data(path, data)
-    logger.log('ALERT', f'The subdomain result for {table_name}: {path}')
-    data = rows.as_dict()
-    return data, fmt, path
-
-
-def domain_to_table(table):
-    return table.replace('.', '_') + "_now_result"
-
-
-if __name__ == '__main__':
-    fire.Fire(export)

export.py

@@ -0,0 +1,70 @@
+#!/usr/bin/python3
+# coding=utf-8
+
+"""
+OneForAll export from database module
+
+:copyright: Copyright (c) 2019, Jing Ling. All rights reserved.
+:license: GNU General Public License v3.0, see LICENSE for more details.
+"""
+
+import fire
+
+from common import utils
+from common.database import Database
+from config.log import logger
+
+
+def export_data(target, db=None, alive=False, limit=None, path=None, fmt='csv', show=False):
+    """
+    OneForAll export from database module
+
+    Example:
+        python3 export.py --target name --fmt csv --dir= ./result.csv
+        python3 export.py --target name --tb True --show False
+        python3 export.py --db result.db --target name --show False
+
+    Note:
+        --fmt csv/json (result format)
+        --path   Result directory (default directory is ./results)
+
+    :param str  target:  Table to be exported
+    :param str  db:      Database path to be exported (default ./results/result.sqlite3)
+    :param bool alive:   Only export the results of alive subdomains (default False)
+    :param str  limit:   Export limit (default None)
+    :param str  fmt:     Result format (default csv)
+    :param str  path:    Result directory (default None)
+    :param bool show:    Displays the exported data in terminal (default False)
+    """
+
+    database = Database(db)
+    domains = utils.get_domains(target)
+    datas = list()
+    if domains:
+        for domain in domains:
+            table_name = domain.replace('.', '_')
+            rows = database.export_data(table_name, alive, limit)
+            if rows is None:
+                continue
+            data, _, _ = do_export(fmt, path, rows, show, table_name, target)
+            datas.extend(data)
+    database.close()
+    if len(domains) > 1:
+        utils.export_all(alive, fmt, path, datas)
+    return datas
+
+
+def do_export(fmt, path, rows, show, table_name, target):
+    fmt = utils.check_format(fmt)
+    path = utils.check_path(path, target, fmt)
+    if show:
+        print(rows.dataset)
+    data = rows.export(fmt)
+    utils.save_to_file(path, data)
+    logger.log('ALERT', f'The subdomain result for {table_name}: {path}')
+    data = rows.as_dict()
+    return data, fmt, path
+
+
+if __name__ == '__main__':
+    fire.Fire(export_data)

modules/altdns.py

@@ -196,4 +196,4 @@ class Altdns(Module):
         resolved_data = resolve.run_resolve(self.domain, self.results)
         request.run_request(self.domain, resolved_data, port)
         logger.log('INFOR', f'Saving altdns results')
-        utils.save_db(self.domain, data, 'altdns')
+        utils.save_to_db(self.domain, data, 'altdns')

oneforall.py

@@ -12,7 +12,7 @@ import fire
 from datetime import datetime
 
 
-import dbexport
+import export
 from brute import Brute
 from common import utils, resolve, request
 from modules.collect import Collect
@@ -129,15 +129,14 @@ class OneForAll(object):
             logger.log('FATAL', 'You must provide either target or targets parameter')
             exit(1)
 
-    def export(self, table):
+    def export_data(self):
         """
         Export data from the database
 
-        :param table: table name
+        :return: exported data
-        :return: export data
         :rtype: list
         """
-        return dbexport.export(table, type='table', alive=self.alive, fmt=self.fmt)
+        return export.export_data(self.domain, alive=self.alive, fmt=self.fmt)
 
     def main(self):
         """
@@ -167,7 +166,7 @@ class OneForAll(object):
         utils.deal_data(self.domain)
         # Export results without resolve
         if not self.dns:
-            return self.export(self.domain)
+            return self.export_data()
 
         self.data = utils.get_data(self.domain)
 
@@ -179,7 +178,7 @@ class OneForAll(object):
 
         # Export results without HTTP request
         if not self.req:
-            return self.export(self.domain)
+            return self.export_data()
 
         # HTTP request
         utils.clear_data(self.domain)
@@ -201,7 +200,7 @@ class OneForAll(object):
             enrich.run()
 
         # Export
-        self.datas.extend(self.export(self.domain))
+        self.datas.extend(self.export_data())
 
         # Scan subdomain takeover
         if self.takeover:
@@ -230,13 +229,16 @@ class OneForAll(object):
         self.config_param()
         self.check_param()
         self.domains = utils.get_domains(self.target, self.targets)
-        if self.domains:
+        count = len(self.domains)
+        logger.log('INFOR', f'Got {count} domains')
+        if not count:
+            logger.log('FATAL', 'Failed to obtain domain')
+            exit(1)
         for domain in self.domains:
             self.domain = utils.get_main_domain(domain)
             self.main()
+        if count > 1:
             utils.export_all(self.alive, self.fmt, self.path, self.datas)
-        else:
-            logger.log('FATAL', 'Failed to obtain domain')
         logger.log('INFOR', 'Finished OneForAll')
 
     @staticmethod

takeover.py

@@ -80,7 +80,7 @@ class Takeover(Module):
             data = str(self.results)
         else:
             data = self.results.export(self.fmt)
-        utils.save_data(self.path, data)
+        utils.save_to_file(self.path, data)
 
     def compare(self, subdomain, cname, responses):
         domain_resp = self.get('http://' + subdomain, check=False, ignore=True)

test.py

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