diff --git a/detection/__main__.py b/detection/__main__.py index 24f8703..9881eb5 100644 --- a/detection/__main__.py +++ b/detection/__main__.py @@ -10,6 +10,7 @@ from .utils import * import sys from colorama import init, Fore, Style from tqdm import tqdm +from pathlib import Path PYCDC_FLAG = True PYCDC_ADDR_FLAG = True @@ -362,19 +363,18 @@ def process_path( ): results = {"high": [], "medium": [], "low": [], "none": []} if os.path.isdir(path): - # 获取所有文件 - all_files = [] - for root, dirs, files in os.walk(path): - for file in files: - file_extension = os.path.splitext(file)[1] - if file_extension in SUPPORTED_EXTENSIONS: - file_path = os.path.join(root, file) - all_files.append(file_path) + # 使用rglob获取所有文件 + all_files = [ + file_path + for file_path in Path(path).rglob("*") + if file_path.suffix in SUPPORTED_EXTENSIONS + ] + # 扫描动画 for file_path in tqdm(all_files, desc="Scanning files", unit="file"): - file_extension = os.path.splitext(file_path)[1] + file_extension = file_path.suffix file_results = checkModeAndDetect( - mode, file_path, file_extension, pycdc_addr + mode, str(file_path), file_extension, pycdc_addr ) if file_results is not None: for key in file_results: