update:完善调用方式,删除多余代码

This commit is contained in:
2024-05-14 21:31:31 +08:00
parent fa86f12a48
commit 0ae787002c
2 changed files with 28 additions and 27 deletions

View File

@@ -3,6 +3,8 @@ from typing import Dict, List, Tuple
from reportlab.lib.pagesizes import letter
from reportlab.lib.styles import getSampleStyleSheet
from reportlab.platypus import Paragraph, Spacer, SimpleDocTemplate
from detection.pickle_detection import pickleDataDetection
from .Regexdetection import find_dangerous_functions
from .GPTdetection import detectGPT
from .utils import *
@@ -176,6 +178,9 @@ def main():
parser.add_argument(
"-m", "--mode", help="Mode of operation:[regex,llm]", default="regex"
)
parser.add_argument(
"-p","--pickle",help="analyze the pickle file",default=None
)
args = parser.parse_args()
output_format = "txt" # Default output format
output_file = None
@@ -191,6 +196,8 @@ def main():
)
output_file = args.output.rsplit(".", 1)[0] + ".txt"
# 如果未指定输出文件,则输出到 stdout否则写入文件
if args.pickle:
pickleDataDetection(args.pickle, output_file)
process_path(args.path, output_format, args.mode, output_file)