detection/pickle_detection.py

@@ -71,10 +71,6 @@ dangerous_names = [
 
 class pickleScanner:
 
-    ReduceCount = 0
-    maliciousModule = []
-    dispatch = {}
-
     def __init__(
         self, file, *, fix_imports=True, encoding="ASCII", errors="strict", buffers=None
     ):
@@ -87,6 +83,8 @@ class pickleScanner:
         self.proto = 0
         self.fix_imports = fix_imports
         self.file = file
+        self.ReduceCount = 0
+        self.maliciousModule = []
 
     def find_class(self, module, name):
         if module.decode() in dangerous_modules or name.decode() in dangerous_names:
@@ -103,8 +101,6 @@ class pickleScanner:
         self.stack = []
         self.append = self.stack.append
         self.proto = 0
-        read = self.read
-        dispatch = self.dispatch
         # 扫描所有的opcodes
         opcode = self.read(1)
         while opcode:
@@ -131,24 +127,24 @@ class pickleScanner:
         b"R",  # REDUCE - same as 'r', but for args tuple
     }
 
-    def output(self):
+    def output(self) -> dict:
         return {
             "ReduceCount": self.ReduceCount,
             "maliciousModule": self.maliciousModule,
         }
 
 
-def pickleDataDetection(file, output_file=None):
+def pickleDataDetection(filename: str, output_file=None):
     """
     :param file: pickle file path
     """
-    with open(file, "rb") as file:
+    with open(filename, "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))
+            json.dump(res, file, indent=4)
     else:
         print(json.dumps(res))