fix: fix some error
This commit is contained in:
@@ -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))
|
||||||
|
|
||||||
|
Reference in New Issue
Block a user