detection/__main__.py

@@ -3,8 +3,6 @@ from typing import Dict, List, Tuple, Optional
 from reportlab.lib.pagesizes import letter
 from reportlab.lib.styles import getSampleStyleSheet
 from reportlab.platypus import Paragraph, Spacer, SimpleDocTemplate
-
-from detection.pickle_detection import pickleDataDetection
 from .Regexdetection import find_dangerous_functions
 from .GPTdetection import detectGPT
 from .pyc_detection import disassemble_pyc
@@ -375,13 +373,6 @@ 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
             )
@@ -396,13 +387,7 @@ def process_path(
                         )
     elif os.path.isfile(path):
         file_extension = os.path.splitext(path)[1]
-        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:
+        if 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:
@@ -440,18 +425,9 @@ def main():
         help="Path to pycdc.exe to decompile",
         default=os.getenv("pycdc"),
     )
-    parser.add_argument(
-        "-P",
-        "--Pickle",
-        help="Path to pickle file to analyze",
-        default=None,
-    )
     args = parser.parse_args()
     output_format = "txt"  # Default output format
     output_file = None
-    if args.Pickle:
-        pickleDataDetection(args.Pickle, args.output)
-        return
     if args.output:
         _, ext = os.path.splitext(args.output)
         ext = ext.lower()

detection/pickle_detection.py

@@ -142,7 +142,11 @@ def pickleDataDetection(filename: str, output_file=None):
         pickscan = pickleScanner(file)
         pickscan.load()
     res = pickscan.output()
-    return res
+    if output_file:
+        with open(output_file, "w") as file:
+            json.dump(res, file, indent=4)
+    else:
+        print(json.dumps(res))
 
 
 if __name__ == "__main__":