修复路径问题

oneforall/common/utils.py

@@ -159,26 +159,20 @@ def check_path(path, name, format):
     """
     filename = f'{name}.{format}'
     default_path = config.result_save_dir.joinpath(filename)
-    if path is None:
-        path = default_path
-    try:
-        path = Path(path)
-    except Exception as e:
-        logger.log('ERROR', e.args)
-        path = default_path
+    if isinstance(path, str):
+        path = repr(path).replace('\\', '/')  # 将路径中的反斜杠替换为正斜杠
+        path = path.replace('\'', '')  # 去除多余的转义
     else:
-        if not path.exists():
-            logger.log('ALERT', f'不存在{path}目录将会新建')
-            path.mkdir(parents=True, exist_ok=True)
-            path = path.joinpath(filename)
-        if path.is_dir():
-            path = path.joinpath(filename)
-        if path.exists():
-            logger.log('ALERT', f'存在{path}文件将会覆盖')
-    # 意外情况
-    if not path:
         path = default_path
-    logger.log('DEBUG', f'结果保存路径{path}')
+    path = Path(path)
+    if not path.suffix:  # 输入是目录的情况
+        path = path.joinpath(filename)
+    parent_dir = path.parent
+    if not parent_dir.exists():
+        logger.log('ALERT', f'不存在{parent_dir}目录将会新建')
+        parent_dir.mkdir(parents=True, exist_ok=True)
+    if path.exists():
+        logger.log('ALERT', f'存在{path}文件将会覆盖')
     return path
 
 
@@ -216,12 +210,10 @@ def save_data(path, data):
         with open(path, 'w', encoding="utf-8",
                   errors='ignore', newline='') as file:
             file.write(data)
-            logger.log('ALERT', f'结果输出{path}')
             return True
     except TypeError:
         with open(path, 'wb') as file:
             file.write(data)
-            logger.log('ALERT', f'结果输出{path}')
             return True
     except Exception as e:
         logger.log('ERROR', e.args)
@@ -306,6 +298,7 @@ def export_all(format, path, datas):
     timestamp = get_timestamp()
     name = f'all_subdomain_result_{timestamp}'
     path = check_path(path, name, format)
+    logger.log('INFOR', f'所有主域的子域结果 {path}')
     row_list = list()
     for row in datas:
         row.pop('header')

oneforall/dbexport.py

@@ -12,6 +12,7 @@ import fire
 
 from common import utils
 from common.database import Database
+from config import logger
 
 
 def export(table, db=None, valid=False, path=None, format='csv', show=False):
@@ -47,6 +48,7 @@ def export(table, db=None, valid=False, path=None, format='csv', show=False):
         data = rows.export(format)
     database.close()
     utils.save_data(path, data)
+    logger.log('INFOR', f'{table}主域的子域结果 {path}')
     data_dict = rows.as_dict()
     return data_dict
 

oneforall/takeover.py

@@ -135,7 +135,7 @@ class Takeover(Module):
         self.subdomains = utils.get_domains(self.target)
         self.format = utils.check_format(self.format, len(self.subdomains))
         timestamp = utils.get_timestamp()
-        name = f'all_subdomain_{timestamp}'
+        name = f'takeover_check_result_{timestamp}'
         self.path = utils.check_path(self.path, name, self.format)
         if self.subdomains:
             logger.log('INFOR', f'正在检查子域接管风险')
@@ -160,7 +160,8 @@ class Takeover(Module):
         elapsed = round(end - start, 1)
         logger.log('INFOR', f'{self.source}模块耗时{elapsed}秒'
                             f'发现{len(self.results)}个子域存在接管风险')
-        logger.log('DEBUG', f'结束执行{self.source}模块')
+        logger.log('INFOR', f'子域接管风险检查结果 {self.path}')
+        logger.log('INFOR', f'结束执行{self.source}模块')
 
 
 if __name__ == '__main__':