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

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:
-            print("The pickle file maybe contains malicious code")
-            print(f"The number of REDUCE opcodes is {self.ReduceCount}")
-            print("The malicious options are: ", self.maliciousModule)
-        else:
-            print("The pickle file is safe")
+        return {
+        "ReduceCount": self.ReduceCount,
+        "maliciousModule": self.maliciousModule
+    }
 
+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))
 
-
-
-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()
+if __name__ == '__main__':
+    pickleDataDetection("test.pkl")