修复导出问题

oneforall/common/utils.py

@@ -165,17 +165,17 @@ def check_dpath(dpath=None):
     return dpath
 
 
-def check_format(format, line):
+def check_format(format, count):
     """
     检查导出格式
 
     :param format: 传入的导出格式
-    :param line: 行数
+    :param count: 数量
     :return: 导出格式
     """
     formats = ['txt', 'rst', 'csv', 'tsv', 'json', 'yaml', 'html',
                'jira', 'xls', 'xlsx', 'dbf', 'latex', 'ods']
-    if format == 'xls' and line > 65000:
+    if format == 'xls' and count > 65000:
         logger.log('ALERT', 'xls文件限制为最多65000行')
         logger.log('ALERT', '使用xlsx格式导出')
         return 'xlsx'
@@ -258,8 +258,7 @@ def check_value(values):
 
 
 def export_all(format, datas):
-    line = len(datas)
+    format = check_format(format, len(datas))
-    format = check_format(format, line)
     dpath = check_dpath()
     timestamp = get_timestamp()
     fpath = dpath.joinpath(f'all_subdomain_{timestamp}.{format}')

oneforall/dbexport.py

@@ -38,8 +38,7 @@ def export(table, db=None, valid=None, dpath=None, format='csv', show=False):
     dpath = utils.check_dpath(dpath)
     database = Database(db)
     rows = database.export_data(table, valid)  # 意外情况导出全部子域
-    line = len(rows)
+    format = utils.check_format(format, len(rows))
-    format = utils.check_format(format, line)
     if show:
         print(rows.dataset)
     if format == 'txt':

oneforall/takeover.py

@@ -133,9 +133,9 @@ class Takeover(Module):
     def run(self):
         start = time.time()
         logger.log('INFOR', f'开始执行{self.source}模块')
-        self.format = utils.check_format(self.format)
-        self.dpath = utils.check_dpath(self.dpath)
         self.subdomains = utils.get_domains(self.target)
+        self.format = utils.check_format(self.format, len(self.subdomains))
+        self.dpath = utils.check_dpath(self.dpath)
         if self.subdomains:
             logger.log('INFOR', f'正在检查子域接管风险')
             self.fingerprints = get_fingerprint()