fix:修改pickle扫描方法 与其他统一

This commit is contained in:
tritium0041 2024-06-04 19:50:31 +08:00
parent 0f2fb3c925
commit ec30999d2c
2 changed files with 15 additions and 6 deletions

View File

@ -375,6 +375,13 @@ def process_path(
# 扫描动画
for file_path in tqdm(all_files, desc="Scanning files", unit="file"):
file_extension = file_path.suffix
if file_extension in [".pkl",".pickle"]:
res = pickleDataDetection(str(file_path), output_file)
results["pickles"].append({
"file": str(file_path),
"result": res
})
continue
file_results = checkModeAndDetect(
mode, str(file_path), file_extension, pycdc_addr
)
@ -389,7 +396,13 @@ def process_path(
)
elif os.path.isfile(path):
file_extension = os.path.splitext(path)[1]
if file_extension in SUPPORTED_EXTENSIONS:
if file_extension in [".pkl", ".pickle"]:
res = pickleDataDetection(str(path), output_file)
results["pickles"].append({
"file": str(path),
"result": res
})
elif file_extension in SUPPORTED_EXTENSIONS:
file_results = checkModeAndDetect(mode, path, file_extension, pycdc_addr)
if file_results is not None:
for key in file_results:

View File

@ -142,11 +142,7 @@ def pickleDataDetection(filename: str, output_file=None):
pickscan = pickleScanner(file)
pickscan.load()
res = pickscan.output()
if output_file:
with open(output_file, "w") as file:
json.dump(res, file, indent=4)
else:
print(json.dumps(res))
return res
if __name__ == "__main__":