Merge branch 'tests/final-tests' of https://git.mamahaha.work/sangge/BackDoorBuster into tests/final-tests

detection/__main__.py

@@ -3,6 +3,8 @@ 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
@@ -373,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
             )
@@ -387,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:
@@ -425,9 +440,18 @@ 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,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__":