diff --git a/detection/Regexdetection.py b/detection/Regexdetection.py index b1b1549..39f6280 100644 --- a/detection/Regexdetection.py +++ b/detection/Regexdetection.py @@ -25,6 +25,16 @@ def find_dangerous_functions( ".cpp": { r"\bsystem\(": "high", }, + ".pyc": { + r"\bexec\b": "high", + r"\beval\b": "high", + r"\bos\.system\b": "high", + r"\bos\.exec\b": "high", + r"\bos\.fork\b": "high", + r"\bos\.kill\b": "high", + r"\bos\.popen\b": "medium", + r"\bos\.spawn\b": "medium", + }, } risk_patterns = patterns.get(file_extension, {}) classified_results = {"high": [], "medium": [], "low": [], "none": []} @@ -36,4 +46,3 @@ def find_dangerous_functions( if re.search(pattern, clean_line): classified_results[risk_level].append((line_number, clean_line)) return classified_results - diff --git a/detection/backdoor_detection.py b/detection/backdoor_detection.py index ef3e32b..0dbbe40 100644 --- a/detection/backdoor_detection.py +++ b/detection/backdoor_detection.py @@ -3,12 +3,13 @@ from typing import Dict, List, Tuple from reportlab.lib.pagesizes import letter from reportlab.lib.styles import getSampleStyleSheet from reportlab.platypus import Paragraph, Spacer, SimpleDocTemplate -from .Regexdetection import find_dangerous_functions -from .GPTdetection import detectGPT -from .utils import * +from detection.Regexdetection import find_dangerous_functions +from detection.GPTdetection import detectGPT +from detection.utils import * +from detection.pyc_detection import disassemble_pyc import sys -SUPPORTED_EXTENSIONS = {".py", ".js", ".cpp"} +SUPPORTED_EXTENSIONS = {".py", ".js", ".cpp", ".pyc"} OUTPUT_FORMATS = ["html", "md", "txt", "pdf"] @@ -119,12 +120,23 @@ def output_text(results: Dict[str, List[Tuple[int, str]]], file_name=None): def checkModeAndDetect(mode: str, filePath: str, fileExtension: str): # TODO:添加更多方式,这里提高代码的复用性和扩展性 - if mode == "regex": - return find_dangerous_functions(read_file_content(filePath), fileExtension) - elif mode == "llm": - return detectGPT(read_file_content(filePath)) + if fileExtension == ".pyc": + # 反汇编pyc文件 + file_content = disassemble_pyc(filePath) + if mode == "regex": + return find_dangerous_functions(file_content, fileExtension) + elif mode == "llm": + return detectGPT(file_content) + else: + return find_dangerous_functions(file_content, fileExtension) else: - return find_dangerous_functions(read_file_content(filePath), fileExtension) + file_content = read_file_content(filePath) + if mode == "regex": + return find_dangerous_functions(file_content, fileExtension) + elif mode == "llm": + return detectGPT(file_content) + else: + return find_dangerous_functions(file_content, fileExtension) def process_path(path: str, output_format: str, mode: str, output_file=None): diff --git a/detection/pyc_detection.py b/detection/pyc_detection.py new file mode 100644 index 0000000..697da18 --- /dev/null +++ b/detection/pyc_detection.py @@ -0,0 +1,22 @@ +from typing import List, Tuple +import uncompyle6 +import io + + +def disassemble_pyc(file_path: str) -> str: + """ + Disassembles a .pyc file using uncompyle6. + + Args: + file_path (str): The path to the .pyc file. + + Returns: + str: The disassembled code as a string. + """ + output = io.StringIO() + try: + uncompyle6.main.decompile_file(file_path, output) + return output.getvalue() + except Exception as e: + print(f"Error occurred while disassembling: {e}") + return ""