diff --git a/detection/__main__.py b/detection/__main__.py index ef59a08..9586c4d 100644 --- a/detection/__main__.py +++ b/detection/__main__.py @@ -1,3 +1,4 @@ +import json import os from typing import Dict, List, Tuple, Optional from reportlab.lib.pagesizes import letter @@ -106,7 +107,11 @@ def generate_text_content(results: Dict[str, List[Tuple[int, str]]]) -> str: text_output += "=" * 30 + "\n\n" for risk_level, entries in results.items(): - if entries and risk_level != "none": + if risk_level == "pickles": + text_output += f"Pickles:\n" + for i in entries: + text_output += f" {i['file']}:{json.dumps(i['result'])}\n" + elif entries and risk_level != "none": risk_color = ( { "high": Fore.RED, @@ -139,6 +144,8 @@ def generate_text_content(results: Dict[str, List[Tuple[int, str]]]) -> str: text_output += line_text text_output += "\n" + + return text_output @@ -363,7 +370,7 @@ def checkModeAndDetect(mode: str, filePath: str, fileExtension: str, pycdc_addr: def process_path( path: str, output_format: str, mode: str, pycdc_addr: str, output_file=None ): - results = {"high": [], "medium": [], "low": [], "none": []} + results = {"high": [], "medium": [], "low": [], "none": [],"pickles": []} if os.path.isdir(path): # 使用rglob获取所有文件 all_files = [ @@ -440,18 +447,9 @@ def main(): help="Path to pycdc.exe to decompile", default=os.getenv("pycdc"), ) - parser.add_argument( - "-P", - "--Pickle", - help="Path to pickle file to analyze", - default=None, - ) args = parser.parse_args() output_format = "txt" # Default output format output_file = None - if args.Pickle: - pickleDataDetection(args.Pickle, args.output) - return if args.output: _, ext = os.path.splitext(args.output) ext = ext.lower()