diff --git a/detection/pickle_detection.py b/detection/pickle_detection.py index 7016641..cfbd258 100644 --- a/detection/pickle_detection.py +++ b/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))