diff --git a/detection/__main__.py b/detection/__main__.py index b19b53d..ef59a08 100644 --- a/detection/__main__.py +++ b/detection/__main__.py @@ -375,6 +375,13 @@ def process_path( # 扫描动画 for file_path in tqdm(all_files, desc="Scanning files", unit="file"): file_extension = file_path.suffix + if file_extension in [".pkl",".pickle"]: + res = pickleDataDetection(str(file_path), output_file) + results["pickles"].append({ + "file": str(file_path), + "result": res + }) + continue file_results = checkModeAndDetect( mode, str(file_path), file_extension, pycdc_addr ) @@ -389,7 +396,13 @@ def process_path( ) elif os.path.isfile(path): file_extension = os.path.splitext(path)[1] - if file_extension in SUPPORTED_EXTENSIONS: + if file_extension in [".pkl", ".pickle"]: + res = pickleDataDetection(str(path), output_file) + results["pickles"].append({ + "file": str(path), + "result": res + }) + elif file_extension in SUPPORTED_EXTENSIONS: file_results = checkModeAndDetect(mode, path, file_extension, pycdc_addr) if file_results is not None: for key in file_results: diff --git a/detection/pickle_detection.py b/detection/pickle_detection.py index cfbd258..c848b4c 100644 --- a/detection/pickle_detection.py +++ b/detection/pickle_detection.py @@ -142,11 +142,7 @@ def pickleDataDetection(filename: str, output_file=None): pickscan = pickleScanner(file) pickscan.load() res = pickscan.output() - if output_file: - with open(output_file, "w") as file: - json.dump(res, file, indent=4) - else: - print(json.dumps(res)) + return res if __name__ == "__main__":