From 95feda67d9c620a0cd44d29072e0c3af8562f0f6 Mon Sep 17 00:00:00 2001
From: dqy <1016751306@qq.com>
Date: Thu, 30 May 2024 22:15:49 +0800
Subject: [PATCH 1/2] =?UTF-8?q?feat:=20=E7=BE=8E=E5=8C=96=E8=BE=93?=
 =?UTF-8?q?=E5=87=BA?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 detection/__main__.py | 311 ++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 296 insertions(+), 15 deletions(-)

diff --git a/detection/__main__.py b/detection/__main__.py
index 5f41266..0faabd6 100644
--- a/detection/__main__.py
+++ b/detection/__main__.py
@@ -1,5 +1,5 @@
 import os
-from typing import Dict, List, Tuple
+from typing import Dict, List, Tuple, Optional
 from reportlab.lib.pagesizes import letter
 from reportlab.lib.styles import getSampleStyleSheet
 from reportlab.platypus import Paragraph, Spacer, SimpleDocTemplate
@@ -7,29 +7,156 @@ from .Regexdetection import find_dangerous_functions
 from .GPTdetection import detectGPT
 from .utils import *
 import sys
+from colorama import init, Fore, Style
 
 SUPPORTED_EXTENSIONS = {".py", ".js", ".cpp"}
 OUTPUT_FORMATS = ["html", "md", "txt", "pdf"]
+ORDERS = [
+    "__import__",
+    "system",
+    "exec",
+    "popen",
+    "eval",
+    "subprocess",
+    "__getattribute__",
+    "getattr",
+    "child_process",
+]
+
+# Initialize colorama
+init(autoreset=True)
+
+ORANGE = "\033[38;5;214m"
+CYAN = Fore.CYAN
 
 
-def generate_text_content(results):
+def supports_color() -> bool:
+    """
+    Checks if the running terminal supports color output.
+
+    Returns:
+        bool: True if the terminal supports color, False otherwise.
+    """
+    # Windows support
+    if sys.platform == "win32":
+        return True
+    # Check if output is a TTY (terminal)
+    if hasattr(sys.stdout, "isatty") and sys.stdout.isatty():
+        return True
+    return False
+
+
+def supports_emoji() -> bool:
+    """
+    Checks if the running terminal supports emoji output.
+
+    Returns:
+        bool: True if the terminal supports emoji, False otherwise.
+    """
+    # This is a simple check. Modern terminals typically support emoji.
+    return sys.platform != "win32" or os.getenv("WT_SESSION") is not None
+
+
+def highlight_orders(line: str, risk_level: str, use_color: bool) -> str:
+    """
+    Highlights specific orders in the line based on risk level.
+
+    Args:
+        line (str): The line to highlight.
+        risk_level (str): The risk level of the line ("high", "medium", "low").
+        use_color (bool): Whether to use color for highlighting.
+
+    Returns:
+        str: The highlighted line.
+    """
+    risk_colors = {
+        "high": Fore.RED,
+        "medium": Fore.YELLOW,
+        "low": CYAN,
+    }
+    color = risk_colors.get(risk_level, Fore.WHITE) if use_color else ""
+    reset = Style.RESET_ALL if use_color else ""
+
+    for order in ORDERS:
+        line = line.replace(order, f"{color}{order}{reset}")
+    return line
+
+
+def generate_text_content(results: Dict[str, List[Tuple[int, str]]]) -> str:
+    """
+    Generates a formatted text report for security analysis results.
+
+    Args:
+        results (Dict[str, List[Tuple[int, str]]]): The security analysis results categorized by risk levels.
+
+    Returns:
+        str: The formatted text report as a string.
+    """
+    use_color = supports_color()
+    use_emoji = supports_emoji()
+
     text_output = "Security Analysis Report\n"
+    text_output += "=" * 30 + "\n\n"
+
     for risk_level, entries in results.items():
         if entries and risk_level != "none":
-            text_output += f"{risk_level.capitalize()} Risk:\n"
+            risk_color = (
+                {
+                    "high": Fore.RED,
+                    "medium": Fore.YELLOW,
+                    "low": Fore.GREEN,
+                }.get(risk_level, Fore.WHITE)
+                if use_color
+                else ""
+            )
+
+            risk_title = (
+                {
+                    "High": "👹",
+                    "Medium": "👾",
+                    "Low": "👻",
+                }
+                if use_emoji
+                else {
+                    "High": "",
+                    "Medium": "",
+                    "Low": "",
+                }
+            )
+
+            text_output += f"{risk_color}{risk_level.capitalize()} Risk{risk_title[risk_level.capitalize()]}:{Style.RESET_ALL if use_color else ''}\n"
+            text_output += "-" * (len(risk_level) + 6) + "\n"
             for line_num, line in entries:
-                text_output += f"  Line {line_num}: {line}\n"
+                line = highlight_orders(line, risk_level, use_color)
+                line_text = f"{Style.RESET_ALL if use_color else ''} {Fore.GREEN if use_color else ''}{line_num}{Style.RESET_ALL if use_color else ''}: {line}{Style.RESET_ALL if use_color else ''}\n"
+                text_output += line_text
+            text_output += "\n"
+
     return text_output
 
 
-def output_results(results, output_format, output_file=None):
+def output_results(
+    results: Dict[str, List[Tuple[int, str]]],
+    output_format: str,
+    output_file: Optional[str] = None,
+) -> None:
+    """
+    Outputs the security analysis results in the specified format.
+
+    Args:
+        results (Dict[str, List[Tuple[int, str]]]): The security analysis results categorized by risk levels.
+        output_format (str): The format to output the results in. Supported formats: "pdf", "html", "md", "txt".
+        output_file (Optional[str]): The name of the file to save the output. If None, prints to the terminal.
+    """
+    OUTPUT_FORMATS = {"pdf", "html", "md", "txt"}
+
     if output_file:
-        file_name = os.path.splitext(output_file)
+        file_name, file_ext = os.path.splitext(output_file)
         if output_format not in OUTPUT_FORMATS:
             output_format = "txt"
             output_file = f"{file_name}.txt"
         results_dir = os.path.dirname(output_file)
-        if not os.path.exists(results_dir):
+        if not os.path.exists(results_dir) and results_dir != "":
             os.makedirs(results_dir)
         if output_format == "pdf":
             output_pdf(results, output_file)
@@ -45,6 +172,79 @@ def output_results(results, output_format, output_file=None):
         print(txt_output)
 
 
+# def generate_text_content(results: Dict[str, List[Tuple[int, str]]]) -> str:
+#     """
+#     Generates a formatted text report for security analysis results.
+
+#     Args:
+#         results (Dict[str, List[Tuple[int, str]]]): The security analysis results categorized by risk levels.
+
+#     Returns:
+#         str: The formatted text report as a string.
+#     """
+#     text_output = (
+#         f"{Fore.CYAN}{Style.BRIGHT}Security Analysis Report{Style.RESET_ALL}\n"
+#     )
+#     text_output += "=" * 30 + "\n\n"
+
+#     for risk_level, entries in results.items():
+#         if entries and risk_level != "none":
+#             risk_color = {
+#                 "high": Fore.RED,
+#                 "medium": Fore.YELLOW,
+#                 "low": Fore.GREEN,
+#             }.get(risk_level, Fore.WHITE)
+#             risk_title = {
+#                 "High": "👹",
+#                 "Medium": "👾",
+#                 "Low": "👻",
+#             }
+#             text_output += f"{risk_color}{risk_level.capitalize()} Risk{risk_title[risk_level.capitalize()]}:{Style.RESET_ALL}\n"
+#             text_output += "-" * (len(risk_level) + 6) + "\n"
+#             for line_num, line in entries:
+#                 text_output += f"  Line {line_num}: {line}\n"
+#             text_output += "\n"
+
+#     return text_output
+
+
+# def output_results(
+#     results: Dict[str, List[Tuple[int, str]]],
+#     output_format: str,
+#     output_file: Optional[str] = None,
+# ) -> None:
+#     """
+#     Outputs the security analysis results in the specified format.
+
+#     Args:
+#         results (Dict[str, List[Tuple[int, str]]]): The security analysis results categorized by risk levels.
+#         output_format (str): The format to output the results in. Supported formats: "pdf", "html", "md", "txt".
+#         output_file (Optional[str]): The name of the file to save the output. If None, prints to the terminal.
+#     """
+#     OUTPUT_FORMATS = {"pdf", "html", "md", "txt"}
+
+#     if output_file:
+#         file_name, file_ext = os.path.splitext(output_file)
+#         if output_format not in OUTPUT_FORMATS:
+#             output_format = "txt"
+#             output_file = f"{file_name}.txt"
+#         results_dir = os.path.dirname(output_file)
+#         if not os.path.exists(results_dir) and results_dir != "":
+#             os.makedirs(results_dir)
+#         if output_format == "pdf":
+#             output_pdf(results, output_file)
+#         elif output_format == "html":
+#             output_html(results, output_file)
+#         elif output_format == "md":
+#             output_markdown(results, output_file)
+#         else:  # Default to txt
+#             output_text(results, output_file)
+#     else:
+#         # If no output file is specified, default to text output to the terminal.
+#         txt_output = generate_text_content(results)
+#         print(txt_output)
+
+
 def output_pdf(results: Dict[str, List[Tuple[int, str]]], file_name):
     doc = SimpleDocTemplate(file_name, pagesize=letter)
     story = []
@@ -73,46 +273,127 @@ def output_pdf(results: Dict[str, List[Tuple[int, str]]], file_name):
 
 
 def output_html(results: Dict[str, List[Tuple[int, str]]], file_name=None):
-    html_output = "<html><head><title>Security Analysis Report</title></head><body>"
-    html_output += "<h1>Security Analysis Report</h1>"
+    """
+    Generates an HTML report for security analysis results.
+
+    Args:
+        results (Dict[str, List[Tuple[int, str]]]): The security analysis results categorized by risk levels.
+        file_name (Optional[str]): The name of the file to save the HTML output. If None, returns the HTML string.
+
+    Returns:
+        Optional[str]: The HTML string if file_name is None, otherwise None.
+    """
+    html_output = """
+    <html>
+    <head>
+        <meta charset="UTF-8">
+        <meta name="viewport" content="width=device-width, initial-scale=1.0">
+        <link rel="icon" href="https://s2.loli.net/2024/05/30/WDc6MekjbuCU9Qo.png">
+        <title>Security Analysis Report</title>
+        <style>
+            body {
+                background-image: url('https://s2.loli.net/2024/05/30/85Mv7leB2IRWNp6.jpg');
+                background-size: 100%, auto;
+                background-attachment: fixed;
+                font-family: Arial, sans-serif;
+            }
+            h1, h2 {
+                color: white;
+            }
+            ul {
+                list-style-type: none;
+                padding: 0;
+            }
+            li {
+                background: rgba(255, 255, 255, 0.8);
+                margin: 5px 0;
+                padding: 10px;
+                border-radius: 5px;
+            }
+        </style>
+    </head>
+    <body>
+        <h1>Security Analysis Report</h1>
+    """
+
     for risk_level, entries in results.items():
         if risk_level != "none":
-            html_output += f"<h2>{risk_level.capitalize()} Risk</h2><ul>"
+            risk_title = {
+                "High": f"<h2>{risk_level.capitalize()} Risk👹</h2><ul>",
+                "Medium": f"<h2>{risk_level.capitalize()} Risk👾</h2><ul>",
+                "Low": f"<h2>{risk_level.capitalize()} Risk👻</h2><ul>",
+            }
+            html_output += risk_title[risk_level.capitalize()]
             for line_num, line in entries:
                 html_output += f"<li>{line_num}: {line}</li>"
             html_output += "</ul>"
+
     html_output += "</body></html>"
+
     if file_name:
-        with open(file_name, "w") as file:
+        with open(file_name, "w", encoding="utf-8") as file:
             file.write(html_output)
+        return None
     else:
         return html_output
 
 
 def output_markdown(results: Dict[str, List[Tuple[int, str]]], file_name=None):
-    md_output = "# Security Analysis Report\n"
+    """
+    Generates a Markdown report for security analysis results.
+
+    Args:
+        results (Dict[str, List[Tuple[int, str]]]): The security analysis results categorized by risk levels.
+        file_name (Optional[str]): The name of the file to save the Markdown output. If None, returns the Markdown string.
+
+    Returns:
+        Optional[str]: The Markdown string if file_name is None, otherwise None.
+    """
+    md_output = "# Security Analysis Report\n\n"
+
     for risk_level, entries in results.items():
         if risk_level != "none":
-            md_output += f"## {risk_level.capitalize()} Risk\n"
+            md_output += f"## {risk_level.capitalize()} Risk\n\n"
+            md_output += "| Line Number | Description |\n"
+            md_output += "|-------------|-------------|\n"
             for line_num, line in entries:
-                md_output += f"- {line_num}: {line}\n"
+                md_output += f"| {line_num} | {line} |\n"
+            md_output += "\n"
+
     if file_name:
         with open(file_name, "w") as file:
             file.write(md_output)
+        return None
     else:
         return md_output
 
 
 def output_text(results: Dict[str, List[Tuple[int, str]]], file_name=None):
+    """
+    Generates a plain text report for security analysis results.
+
+    Args:
+        results (Dict[str, List[Tuple[int, str]]]): The security analysis results categorized by risk levels.
+        file_name (Optional[str]): The name of the file to save the text output. If None, returns the text string.
+
+    Returns:
+        Optional[str]: The text string if file_name is None, otherwise None.
+    """
     text_output = "Security Analysis Report\n"
+    text_output += "=" * len("Security Analysis Report") + "\n\n"
+
     for risk_level, entries in results.items():
         if risk_level != "none":
             text_output += f"{risk_level.capitalize()} Risk:\n"
+            text_output += "-" * len(f"{risk_level.capitalize()} Risk:") + "\n"
             for line_num, line in entries:
-                text_output += f"  {line_num}: {line}\n"
+                text_output += f"  Line {line_num}: {line}\n"
+            text_output += "\n"
+
     if file_name:
         with open(file_name, "w") as file:
             file.write(text_output)
+        return None
     else:
         return text_output
 
-- 
2.47.2


From f798cf143c254615acaeca9e46f886af20a3911d Mon Sep 17 00:00:00 2001
From: dqy <1016751306@qq.com>
Date: Thu, 30 May 2024 22:27:04 +0800
Subject: [PATCH 2/2] =?UTF-8?q?fix:=20=E5=88=A0=E9=99=A4=E5=8E=9F=E6=9C=89?=
 =?UTF-8?q?=E4=BB=A3=E7=A0=81=E6=B3=A8=E9=87=8A?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 detection/__main__.py | 73 -------------------------------------------
 1 file changed, 73 deletions(-)

diff --git a/detection/__main__.py b/detection/__main__.py
index 0faabd6..644ac18 100644
--- a/detection/__main__.py
+++ b/detection/__main__.py
@@ -172,79 +172,6 @@ def output_results(
         print(txt_output)
 
 
-# def generate_text_content(results: Dict[str, List[Tuple[int, str]]]) -> str:
-#     """
-#     Generates a formatted text report for security analysis results.
-
-#     Args:
-#         results (Dict[str, List[Tuple[int, str]]]): The security analysis results categorized by risk levels.
-
-#     Returns:
-#         str: The formatted text report as a string.
-#     """
-#     text_output = (
-#         f"{Fore.CYAN}{Style.BRIGHT}Security Analysis Report{Style.RESET_ALL}\n"
-#     )
-#     text_output += "=" * 30 + "\n\n"
-
-#     for risk_level, entries in results.items():
-#         if entries and risk_level != "none":
-#             risk_color = {
-#                 "high": Fore.RED,
-#                 "medium": Fore.YELLOW,
-#                 "low": Fore.GREEN,
-#             }.get(risk_level, Fore.WHITE)
-#             risk_title = {
-#                 "High": "👹",
-#                 "Medium": "👾",
-#                 "Low": "👻",
-#             }
-#             text_output += f"{risk_color}{risk_level.capitalize()} Risk{risk_title[risk_level.capitalize()]}:{Style.RESET_ALL}\n"
-#             text_output += "-" * (len(risk_level) + 6) + "\n"
-#             for line_num, line in entries:
-#                 text_output += f"  Line {line_num}: {line}\n"
-#             text_output += "\n"
-
-#     return text_output
-
-
-# def output_results(
-#     results: Dict[str, List[Tuple[int, str]]],
-#     output_format: str,
-#     output_file: Optional[str] = None,
-# ) -> None:
-#     """
-#     Outputs the security analysis results in the specified format.
-
-#     Args:
-#         results (Dict[str, List[Tuple[int, str]]]): The security analysis results categorized by risk levels.
-#         output_format (str): The format to output the results in. Supported formats: "pdf", "html", "md", "txt".
-#         output_file (Optional[str]): The name of the file to save the output. If None, prints to the terminal.
-#     """
-#     OUTPUT_FORMATS = {"pdf", "html", "md", "txt"}
-
-#     if output_file:
-#         file_name, file_ext = os.path.splitext(output_file)
-#         if output_format not in OUTPUT_FORMATS:
-#             output_format = "txt"
-#             output_file = f"{file_name}.txt"
-#         results_dir = os.path.dirname(output_file)
-#         if not os.path.exists(results_dir) and results_dir != "":
-#             os.makedirs(results_dir)
-#         if output_format == "pdf":
-#             output_pdf(results, output_file)
-#         elif output_format == "html":
-#             output_html(results, output_file)
-#         elif output_format == "md":
-#             output_markdown(results, output_file)
-#         else:  # Default to txt
-#             output_text(results, output_file)
-#     else:
-#         # If no output file is specified, default to text output to the terminal.
-#         txt_output = generate_text_content(results)
-#         print(txt_output)
-
-
 def output_pdf(results: Dict[str, List[Tuple[int, str]]], file_name):
     doc = SimpleDocTemplate(file_name, pagesize=letter)
     story = []
-- 
2.47.2