Merge branch 'fix/requirements-detection' into tests/final-tests

detection/__main__.py

@@ -6,6 +6,8 @@ from reportlab.lib.styles import getSampleStyleSheet
 from reportlab.platypus import Paragraph, Spacer, SimpleDocTemplate
 
 from detection.pickle_detection import pickleDataDetection
+
+from .requirements_detection import requirement_detection
 from .Regexdetection import find_dangerous_functions
 from .GPTdetection import detectGPT, GPTdetectFileList
 
@@ -371,7 +373,12 @@ 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
+    path: str,
+    output_format: str,
+    mode: str,
+    pycdc_addr: str,
+    output_file=None,
+    requirement_path=None,
 ):
     results = {"high": [], "medium": [], "low": [], "none": [], "pickles": []}
     if os.path.isdir(path):
@@ -381,6 +388,7 @@ def process_path(
             for file_path in Path(path).rglob("*")
             if file_path.suffix in SUPPORTED_EXTENSIONS
         ]
+<<<<<<< HEAD
         if mode == "llm":
             results = GPTdetectFileList(all_files)
         else:
@@ -425,7 +433,8 @@ def process_path(
     else:
         print("Invalid path.")
         sys.exit(1)
-
+    if requirement_path is not None:
+        requirement_detection(requirement_path, output_file)
     output_results(results, output_format, output_file)
 
 
@@ -446,6 +455,18 @@ def main():
         help="Path to pycdc.exe to decompile",
         default=os.getenv("PATH"),
     )
+    parser.add_argument(
+        "-P",
+        "--Pickle",
+        help="Path to pickle file to analyze",
+        default=None,
+    )
+    parser.add_argument(
+        "-r",
+        "--requirement",
+        help="Path to requirement file to analyze",
+        default=None,
+    )
     args = parser.parse_args()
     output_format = "txt"  # Default output format
     output_file = None
@@ -461,7 +482,9 @@ def main():
             )
             output_file = args.output.rsplit(".", 1)[0] + ".txt"
     # 如果未指定输出文件，则输出到 stdout；否则写入文件
-    process_path(args.path, output_format, args.mode, args.pycdc, output_file)
+    process_path(
+        args.path, output_format, args.mode, args.pycdc, output_file, args.requirement
+    )
     if PYCDC_FLAG == False:
         print(
             "ERROR: Detected Python 3.11 or above .pyc files. You need to install pycdc and compile it yourself to obtain pycdc."

detection/requirements_detection.py

@@ -3,6 +3,16 @@ import requests
 from bs4 import BeautifulSoup
 from packaging.version import Version, InvalidVersion
 import sys
+from reportlab.lib.pagesizes import letter
+from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
+from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
+from colorama import Fore, Style, init
+from tqdm import tqdm
+import html
+import os
+
+
+init(autoreset=True)  # 初始化colorama，并在每次打印后自动重置颜色
 
 
 def fetch_html(url: str) -> str:
@@ -55,7 +65,6 @@ def version_in_range(version, range_str: str) -> bool:
         except InvalidVersion:
             return False
     else:
-        # 如果没有给版本号，默认使用最新版本
         if range_str[-2] == ",":
             return True
 
@@ -77,64 +86,183 @@ def version_in_range(version, range_str: str) -> bool:
     return True
 
 
-def check_vulnerabilities(requirements: list, base_url: str, output_file: str):
+def check_vulnerabilities(requirements: list, base_url: str) -> str:
-    with open(output_file, "w") as out_file:
+    results = []
-        for req in requirements:
+    for req in tqdm(requirements, desc="Checking vulnerabilities", unit="dependency"):
-            version = ""
+        version = ""
-            # 如果有版本
+        if "==" in req:
-            if "==" in req:
+            package_name, version = req.split("==")
-                package_name, version = req.split("==")
+        else:
-            # 没有版本
+            package_name, version = req, None
-            else:
+        url = f"{base_url}{package_name}"
-                package_name, version = req, None
+        # print(f"\nFetching data for {package_name} from {url}")
-            # 拼接URL
+        html_content = fetch_html(url)
-            url = f"{base_url}{package_name}"
+        if html_content:
-            print(f"Fetching data for {package_name} from {url}")
+            extracted_data = parse_html(html_content)
-            html_content = fetch_html(url)
+            if extracted_data:
-            if html_content:
+                relevant_vulns = []
-                # 解析hmtl
+                for vuln in extracted_data:
-                extracted_data = parse_html(html_content)
+                    if version_in_range(version, vuln["chip"]):
-                if extracted_data:
+                        relevant_vulns.append(vuln)
-                    relevant_vulns = []
+                if relevant_vulns:
-                    for vuln in extracted_data:
+                    result = f"Vulnerabilities found for {package_name}:\n"
-                        if version_in_range(version, vuln["chip"]):
+                    for vuln in relevant_vulns:
-                            relevant_vulns.append(vuln)
+                        result += f"  - {vuln['link']}\n"
-                    if relevant_vulns:
+                    results.append(result)
-                        out_file.write(f"Vulnerabilities found for {package_name}:\n")
+    return "\n".join(results)
-                        for vuln in relevant_vulns:
-                            out_file.write(f"  - {vuln['link']}\n")
-                        out_file.write("\n")
-                else:
-                    print(f"No relevant data found for {package_name}.")
-            else:
-                print(f"Failed to fetch data for {package_name}.")
 
 
-def main():
+def save_to_file(output_path: str, data: str):
-    parser = argparse.ArgumentParser(
+    if output_path.endswith(".html"):
-        description="Check project dependencies for vulnerabilities."
+        save_as_html(output_path, data)
-    )
+    elif output_path.endswith(".pdf"):
-    parser.add_argument(
+        save_as_pdf(output_path, data)
-        "-r",
+    elif output_path.endswith(".md"):
-        "--requirement",
+        save_as_markdown(output_path, data)
-        help="Path to the requirements file of the project",
+    else:
-        required=True,
+        save_as_txt(output_path, data)
-    )
-    parser.add_argument(
-        "-o",
-        "--output",
-        help="Output file path with extension, e.g., './output/report.txt'",
-        required=True,
-    )
-    args = parser.parse_args()
 
+
+def save_as_html(output_path: str, data: str):
+    escaped_data = html.escape(data)
+    html_content = f"""
+    <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>Vulnerability Report</title>
+        <style>
+            body {{
+                font-family: Arial, sans-serif;
+                background-image: url('https://s2.loli.net/2024/05/30/85Mv7leB2IRWNp6.jpg');
+                background-size: cover;
+                color: #333;
+                margin: 0;
+                padding: 0;
+                display: flex;
+                justify-content: center;
+                align-items: center;
+                height: 100vh;
+            }}
+            .container {{
+                background: rgba(255, 255, 255, 0.8);
+                border-radius: 10px;
+                padding: 20px;
+                box-shadow: 0 0 10px rgba(0, 0, 0, 0.1);
+                max-width: 800px;
+                width: 100%;
+                margin: 20px;
+                overflow-y: auto;
+                max-height: 90vh;
+            }}
+            .title {{
+                font-size: 24px;
+                font-weight: bold;
+                text-align: center;
+                margin-bottom: 20px;
+            }}
+            pre {{
+                white-space: pre-wrap;
+                word-wrap: break-word;
+                font-size: 14px;
+                line-height: 1.5;
+                color: #333;
+                background: #f4f4f4;
+                padding: 10px;
+                border-radius: 5px;
+                border: 1px solid #ddd;
+                overflow: auto;
+                font-weight: bold;
+            }}
+        </style>
+    </head>
+    <body>
+        <div class="container">
+            <div class="title">Vulnerability Report</div>
+            <pre>{escaped_data}</pre>
+        </div>
+    </body>
+    </html>
+    """
+    with open(output_path, "w", encoding="utf-8") as file:
+        file.write(html_content)
+
+
+def save_as_pdf(output_path: str, data: str):
+    doc = SimpleDocTemplate(output_path, pagesize=letter)
+    story = []
+    styles = getSampleStyleSheet()
+
+    # Add the title centered
+    title_style = ParagraphStyle(
+        "Title",
+        parent=styles["Title"],
+        alignment=1,  # Center alignment
+        fontSize=24,
+        leading=28,
+        spaceAfter=20,
+        fontName="Helvetica-Bold",
+    )
+    title = Paragraph("Vulnerability Report", title_style)
+    story.append(title)
+
+    # Normal body text style
+    normal_style = ParagraphStyle(
+        "BodyText", parent=styles["BodyText"], fontSize=12, leading=15, spaceAfter=12
+    )
+
+    # Add the vulnerability details
+    for line in data.split("\n"):
+        if line.strip():  # Skip empty lines
+            story.append(Paragraph(line, normal_style))
+
+    doc.build(story)
+
+
+def save_as_markdown(output_path: str, data: str):
+    with open(output_path, "w") as file:
+        file.write("## Vulnerability Report: \n\n")
+        file.write(data)
+
+
+def save_as_txt(output_path: str, data: str):
+    with open(output_path, "w") as file:
+        file.write("Vulnerability Report: \n\n")
+        file.write(data)
+
+
+def print_separator(title, char="-", length=50, padding=2):
+    print(f"{title:^{length + 4*padding}}")  # 居中打印标题，两侧各有padding个空格
+    print(char * (length + 2 * padding))  # 打印分割线，两侧各有padding个字符的空格
+
+
+def modify_file_name(file_path: str) -> str:
+    """
+    Modify the file name by adding '-re' before the file extension.
+
+    Args:
+        file_path (str): The original file path.
+
+    Returns:
+        str: The modified file path.
+    """
+    directory, file_name = os.path.split(file_path)
+    name, ext = os.path.splitext(file_name)
+    new_file_name = f"{name}-re{ext}"
+    new_file_path = os.path.join(directory, new_file_name)
+    return new_file_path
+
+
+def requirement_detection(requirement_path, output_path=None):
     base_url = "https://security.snyk.io/package/pip/"
-    # 分析项目依赖，包括名称和版本(如果有的话)
+    requirements = load_requirements(requirement_path)
-    requirements = load_requirements(args.requirement)
+    results = check_vulnerabilities(requirements, base_url)
-    # 传入依赖信息，url前缀，扫描结果输出位置
+    if output_path is not None:
-    check_vulnerabilities(requirements, base_url, args.output)
+        new_path = modify_file_name(output_path)
-    print("Vulnerability scan complete. Results saved to", args.output)
+        save_to_file(new_path, results)
-
+        print(f"Vulnerability scan complete. Results saved to {output_path}")
-
+        print(f"Requirements scan complete. Results saved to {new_path}")
-if __name__ == "__main__":
+    else:
-    main()
+        print_separator("\nVulnerability Report", "=", 40, 5)
+        print(results)