feat: 将依赖检测添加到模组

detection/__main__.py

@@ -5,6 +5,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
 from .pyc_detection import disassemble_pyc
@@ -361,7 +363,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": []}
     if os.path.isdir(path):
@@ -375,12 +382,9 @@ def process_path(
         # 扫描动画
         for file_path in tqdm(all_files, desc="Scanning files", unit="file"):
             file_extension = file_path.suffix
-            if file_extension in [".pkl",".pickle"]:
+            if file_extension in [".pkl", ".pickle"]:
                 res = pickleDataDetection(str(file_path), output_file)
-                results["pickles"].append({
-                    "file": str(file_path),
-                    "result": res
-                })
+                results["pickles"].append({"file": str(file_path), "result": res})
                 continue
             file_results = checkModeAndDetect(
                 mode, str(file_path), file_extension, pycdc_addr
@@ -398,10 +402,7 @@ def process_path(
         file_extension = os.path.splitext(path)[1]
         if file_extension in [".pkl", ".pickle"]:
             res = pickleDataDetection(str(path), output_file)
-            results["pickles"].append({
-                "file": str(path),
-                "result": res
-            })
+            results["pickles"].append({"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 +420,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 +448,12 @@ 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
@@ -464,7 +472,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

@@ -9,6 +9,7 @@ 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，并在每次打印后自动重置颜色
@@ -94,7 +95,7 @@ def check_vulnerabilities(requirements: list, base_url: str) -> str:
         else:
             package_name, version = req, None
         url = f"{base_url}{package_name}"
-        # print(f"Fetching data for {package_name} from {url}")
+        # print(f"\nFetching data for {package_name} from {url}")
         html_content = fetch_html(url)
         if html_content:
             extracted_data = parse_html(html_content)
@@ -236,34 +237,32 @@ def print_separator(title, char="-", length=50, padding=2):
     print(char * (length + 2 * padding))  # 打印分割线，两侧各有padding个字符的空格
 
 
-def main():
-    parser = argparse.ArgumentParser(
-        description="Check project dependencies for vulnerabilities."
-    )
-    parser.add_argument(
-        "-r",
-        "--requirement",
-        help="Path to the requirements file of the project",
-        required=True,
-    )
-    parser.add_argument(
-        "-o",
-        "--output",
-        help="Output file path with extension, e.g., './output/report.txt'",
-    )
-    args = parser.parse_args()
+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(args.requirement)
+    requirements = load_requirements(requirement_path)
     results = check_vulnerabilities(requirements, base_url)
-
-    if args.output:
-        save_to_file(args.output, results)
-        print(f"Vulnerability scan complete. Results saved to {args.output}")
+    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("\n\nVulnerability Report", "=", 40, 5)
+        print_separator("\nVulnerability Report", "=", 40, 5)
         print(results)
-
-
-if __name__ == "__main__":
-    main()