detection/backdoor_detection.py

@@ -1,10 +1,8 @@
 import os
 from typing import Dict, List, Tuple
 from reportlab.lib.pagesizes import letter
-from reportlab.pdfgen import canvas
 from reportlab.lib.styles import getSampleStyleSheet
 from reportlab.platypus import Paragraph, Spacer, SimpleDocTemplate
-from reportlab.lib import colors
 from .Regexdetection import find_dangerous_functions
 from .GPTdetection import detectGPT
 from .utils import *
@@ -25,7 +23,7 @@ def generate_text_content(results):
 
 def output_results(results, output_format, output_file=None):
     if output_file:
-        file_name, file_extension = os.path.splitext(output_file)
+        file_name = os.path.splitext(output_file)
         if output_format not in OUTPUT_FORMATS:
             output_format = "txt"
             output_file = f"{file_name}.txt"
@@ -124,6 +122,8 @@ def checkModeAndDetect(mode: str,filePath: str,fileExtension: str):
         return find_dangerous_functions(read_file_content(filePath), fileExtension)
     elif mode == "llm":
         return detectGPT(read_file_content(filePath))
+    else:
+        return find_dangerous_functions(read_file_content(filePath), fileExtension)
 
 
 def process_path(path: str, output_format: str, mode: str, output_file=None):
@@ -172,7 +172,9 @@ def main():
     parser = argparse.ArgumentParser(description="Backdoor detection tool.")
     parser.add_argument("path", help="Path to the code to analyze")
     parser.add_argument("-o", "--output", help="Output file path", default=None)
-    parser.add_argument("-m", "--mode", help="Mode of operation:[regex,llm]", default="regex")
+    parser.add_argument(
+        "-m", "--mode", help="Mode of operation:[regex,llm]", default="regex"
+    )
     args = parser.parse_args()
     output_format = "txt"  # Default output format
     output_file = None
@@ -187,6 +189,7 @@ def main():
                 "Your input file format was incorrect, the output has been saved as a TXT file."
             )
             output_file = args.output.rsplit(".", 1)[0] + ".txt"
+    # 如果未指定输出文件，则输出到 stdout；否则写入文件
     process_path(args.path, output_format, args.mode, output_file)