update:完善调用方式，删除多余代码

detection/backdoor_detection.py

@@ -3,6 +3,8 @@ from typing import Dict, List, Tuple
 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 .utils import *
@@ -176,6 +178,9 @@ def main():
     parser.add_argument(
         "-m", "--mode", help="Mode of operation:[regex,llm]", default="regex"
     )
+    parser.add_argument(
+        "-p","--pickle",help="analyze the pickle file",default=None
+    )
     args = parser.parse_args()
     output_format = "txt"  # Default output format
     output_file = None
@@ -191,6 +196,8 @@ def main():
             )
             output_file = args.output.rsplit(".", 1)[0] + ".txt"
     # 如果未指定输出文件，则输出到 stdout；否则写入文件
+    if args.pickle:
+        pickleDataDetection(args.pickle, output_file)
     process_path(args.path, output_format, args.mode, output_file)
 
 

detection/pickle_detection.py

@@ -1,4 +1,5 @@
 import io
+import json
 import os
 import pickletools
 import pickle
@@ -87,10 +88,9 @@ class pickleScanner():
         self.file = file
 
     def find_class(self, module, name):
-        print(module, name)
         if module.decode() in dangerous_modules or name.decode() in dangerous_names:
             # self.maliciousCount += 1
-            self.maliciousModule.append((module, name))
+            self.maliciousModule.append((module.decode(), name.decode()))
 
     def load(self):
         self._unframer = _Unframer(self._file_read, self._file_readline)
@@ -128,30 +128,24 @@ class pickleScanner():
 
 
     def output(self):
-        if self.ReduceCount > 0 or len(self.maliciousModule) > 0:
+        return {
-            print("The pickle file maybe contains malicious code")
+        "ReduceCount": self.ReduceCount,
-            print(f"The number of REDUCE opcodes is {self.ReduceCount}")
+        "maliciousModule": self.maliciousModule
-            print("The malicious options are: ", self.maliciousModule)
+    }
-        else:
-            print("The pickle file is safe")
 
+def pickleDataDetection(file,output_file=None):
+    '''
+    :param file: pickle file path
+    '''
+    with open(file, "rb") as file:
+        pickscan = pickleScanner(file)
+        pickscan.load()
+    res = pickscan.output()
+    if output_file:
+        with open(output_file, "w") as file:
+            file.writelines(json.dumps(res))
+    else:
+        print(json.dumps(res))
 
-
+if __name__ == '__main__':
-
+    pickleDataDetection("test.pkl")
-class test:
-    a = 1
-    b = 2
-    def __reduce__(self):
-        return (__import__("os").system,('calc',))
-
-
-data = pickle.dumps(test(),protocol=2)
-print(data)
-print(pickletools.dis(data))
-with open("test.pkl", "wb") as file:
-    file.write(data)
-
-with open("test.pkl", "rb") as file:
-    pickscan = pickleScanner(file)
-    pickscan.load()
-pickscan.output()