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:

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__":