detection/Regexdetection.py

@@ -38,11 +38,12 @@ def find_dangerous_functions(
     }
     risk_patterns = patterns.get(file_extension, {})
     classified_results = {"high": [], "medium": [], "low": [], "none": []}
-    for line_number, line in enumerate(file_content.split("\n"), start=1):
-        clean_line = remove_comments(line, file_extension)
-        if not clean_line:
-            continue
-        for pattern, risk_level in risk_patterns.items():
-            if re.search(pattern, clean_line):
-                classified_results[risk_level].append((line_number, clean_line))
+    if file_content is not None:
+        for line_number, line in enumerate(file_content.split("\n"), start=1):
+            clean_line = remove_comments(line, file_extension)
+            if not clean_line:
+                continue
+            for pattern, risk_level in risk_patterns.items():
+                if re.search(pattern, clean_line):
+                    classified_results[risk_level].append((line_number, clean_line))
     return classified_results

detection/__main__.py

@@ -1,4 +1,5 @@
 import os
+from pickle import FALSE
 from typing import Dict, List, Tuple, Optional
 from reportlab.lib.pagesizes import letter
 from reportlab.lib.styles import getSampleStyleSheet
@@ -10,6 +11,8 @@ from .utils import *
 import sys
 from colorama import init, Fore, Style
 
+PYCDC_FLAG = True
+PYCDC_ADDR_FLAG = True
 SUPPORTED_EXTENSIONS = {".py", ".js", ".cpp", ".pyc"}
 OUTPUT_FORMATS = ["html", "md", "txt", "pdf"]
 ORDERS = [
@@ -331,6 +334,13 @@ def checkModeAndDetect(mode: str, filePath: str, fileExtension: str, pycdc_addr:
     if fileExtension == ".pyc":
         # 反汇编pyc文件
         file_content = disassemble_pyc(filePath, pycdc_addr)
+        if file_content == "none":
+            global PYCDC_FLAG
+            PYCDC_FLAG = False
+            return ""
+        elif file_content == "invalid":
+            global PYCDC_ADDR_FLAG
+            PYCDC_ADDR_FLAG = False
         if mode == "regex":
             return find_dangerous_functions(file_content, fileExtension)
         elif mode == "llm":
@@ -360,26 +370,28 @@ def process_path(
                     file_results = checkModeAndDetect(
                         mode, file_path, file_extension, pycdc_addr
                     )
-                    for key in file_results:
-                        if key != "none":  # Exclude 'none' risk level
-                            results[key].extend(
-                                [
-                                    (f"{file_path}: Line {line_num}", line)
-                                    for line_num, line in file_results[key]
-                                ]
-                            )
+                    if file_results != "":
+                        for key in file_results:
+                            if key != "none":  # Exclude 'none' risk level
+                                results[key].extend(
+                                    [
+                                        (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 SUPPORTED_EXTENSIONS:
             file_results = checkModeAndDetect(mode, path, file_extension, pycdc_addr)
-            for key in file_results:
-                if key != "none":  # Exclude 'none' risk level
-                    results[key].extend(
-                        [
-                            (f"{path}: Line {line_num}", line)
-                            for line_num, line in file_results[key]
-                        ]
-                    )
+            if file_results != "":
+                for key in file_results:
+                    if key != "none":  # Exclude 'none' risk level
+                        results[key].extend(
+                            [
+                                (f"{path}: Line {line_num}", line)
+                                for line_num, line in file_results[key]
+                            ]
+                        )
         else:
             print("Unsupported file type.")
             return
@@ -420,6 +432,14 @@ def main():
             output_file = args.output.rsplit(".", 1)[0] + ".txt"
     # 如果未指定输出文件，则输出到 stdout；否则写入文件
     process_path(args.path, output_format, args.mode, args.pycdc, output_file)
+    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."
+        )
+        print("Repo: https://github.com/zrax/pycdc.git")
+    if PYCDC_ADDR_FLAG == False:
+        print("ERROR: The specified pycdc.exe path is not valid")
+        print("Please check your pycdc path.")
 
 
 if __name__ == "__main__":

detection/pyc_detection.py

@@ -3,6 +3,7 @@ import uncompyle6
 import io
 import os
 import subprocess
+from contextlib import redirect_stdout, redirect_stderr
 
 
 def run_pycdc(exe_path: str, pyc_file: str) -> str:
@@ -17,15 +18,12 @@ def run_pycdc(exe_path: str, pyc_file: str) -> str:
         str: Output from pycdc.exe.
     """
     if not os.path.isfile(exe_path):
-        print(f"ERROR: The specified pycdc.exe path is not valid: {exe_path}")
-        print("Please check your pycdc path.")
-        exit(1)
+        return "invalid"
 
     command = f'"{exe_path}" "{pyc_file}"'
-    result = subprocess.run(command, capture_output=True, text=True, shell=True)
-
-    if result.returncode != 0:
-        raise Exception(f"Error running pycdc.exe: {result.stderr}")
+    result = subprocess.run(
+        command, capture_output=True, text=True, shell=True, encoding="utf-8"
+    )
 
     return result.stdout
 
@@ -46,10 +44,6 @@ def disassemble_pyc(file_path: str, pycdc_addr=None) -> str:
         return output.getvalue()
     except Exception as e:
         if pycdc_addr is None:
-            print(
-                "ERROR: For Python 3.11 and above, you need to install pycdc and compile it yourself to obtain pycdc.exe."
-            )
-            print("repo: https://github.com/zrax/pycdc.git")
-            exit(1)
+            return "none"
         else:
             return run_pycdc(pycdc_addr, file_path)