fix: fix some error

This commit is contained in:
2024-05-27 20:35:13 +08:00
parent fab5e680ef
commit accd50e8ce

View File

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