|
| 1 | +import ast |
| 2 | +import os |
| 3 | + |
| 4 | +from fastapi import HTTPException |
| 5 | +import inspect |
| 6 | +import logging |
| 7 | +from typing import List, Optional, Callable |
| 8 | + |
| 9 | + |
| 10 | +from base_detector_registry import BaseDetectorRegistry |
| 11 | +from detectors.common.scheme import ContentAnalysisResponse |
| 12 | + |
| 13 | +logger = logging.getLogger(__name__) |
| 14 | + |
| 15 | +def custom_func_wrapper(func: Callable, func_name: str, s: str, headers: dict) -> Optional[ContentAnalysisResponse]: |
| 16 | + """Convert a some f(text)->bool into a Detector response""" |
| 17 | + sig = inspect.signature(func) |
| 18 | + try: |
| 19 | + if headers is not None: |
| 20 | + result = func(s, headers) |
| 21 | + else: |
| 22 | + result = func(s) |
| 23 | + |
| 24 | + except Exception as e: |
| 25 | + logging.error(f"Error when computing custom detector function {func_name}: {e}") |
| 26 | + raise e |
| 27 | + if result: |
| 28 | + if isinstance(result, bool): |
| 29 | + return ContentAnalysisResponse( |
| 30 | + start=0, |
| 31 | + end=len(s), |
| 32 | + text=s, |
| 33 | + detection_type=func_name, |
| 34 | + detection=func_name, |
| 35 | + score=1.0) |
| 36 | + elif isinstance(result, dict): |
| 37 | + try: |
| 38 | + return ContentAnalysisResponse(**result) |
| 39 | + except Exception as e: |
| 40 | + logging.error(f"Error when trying to build ContentAnalysisResponse from {func_name} response: {e}") |
| 41 | + raise e |
| 42 | + else: |
| 43 | + msg = f"Unsupported result type for custom detector function {func_name}, must be bool or ContentAnalysisResponse, got: {type(result)}" |
| 44 | + logging.error(msg) |
| 45 | + raise TypeError(msg) |
| 46 | + else: |
| 47 | + return None |
| 48 | + |
| 49 | + |
| 50 | +def static_code_analysis(module_path, forbidden_imports=None, forbidden_calls=None): |
| 51 | + """ |
| 52 | + Perform static code analysis on a Python module to check for forbidden imports and function calls. |
| 53 | + Returns a list of issues found. |
| 54 | + """ |
| 55 | + if forbidden_imports is None: |
| 56 | + forbidden_imports = {"os", "subprocess", "sys", "shutil"} |
| 57 | + if forbidden_calls is None: |
| 58 | + forbidden_calls = {"eval", "exec", "open", "compile", "input"} |
| 59 | + |
| 60 | + issues = [] |
| 61 | + with open(module_path, "r") as f: |
| 62 | + source = f.read() |
| 63 | + try: |
| 64 | + tree = ast.parse(source, filename=module_path) |
| 65 | + except Exception as e: |
| 66 | + issues.append(f"Failed to parse {module_path}: {e}") |
| 67 | + return issues |
| 68 | + |
| 69 | + for node in ast.walk(tree): |
| 70 | + # Check for forbidden imports |
| 71 | + if isinstance(node, ast.Import): |
| 72 | + for alias in node.names: |
| 73 | + if alias.name.split(".")[0] in forbidden_imports: |
| 74 | + issues.append(f"- Forbidden import: {alias.name} (line {node.lineno})") |
| 75 | + if isinstance(node, ast.ImportFrom): |
| 76 | + if node.module and node.module.split(".")[0] in forbidden_imports: |
| 77 | + issues.append(f"- Forbidden import: {node.module} (line {node.lineno})") |
| 78 | + # Check for forbidden function calls |
| 79 | + if isinstance(node, ast.Call): |
| 80 | + func_name = "" |
| 81 | + if isinstance(node.func, ast.Name): |
| 82 | + func_name = node.func.id |
| 83 | + elif isinstance(node.func, ast.Attribute): |
| 84 | + func_name = f"{getattr(node.func.value, 'id', '')}.{node.func.attr}" |
| 85 | + if func_name in forbidden_calls: |
| 86 | + issues.append(f"- Forbidden function call: {func_name} (line {node.lineno})") |
| 87 | + return issues |
| 88 | + |
| 89 | + |
| 90 | +class CustomDetectorRegistry(BaseDetectorRegistry): |
| 91 | + def __init__(self): |
| 92 | + super().__init__() |
| 93 | + |
| 94 | + issues = static_code_analysis(module_path = os.path.join(os.path.dirname(__file__), "custom_detectors", "custom_detectors.py")) |
| 95 | + if issues: |
| 96 | + logging.error(f"Detected {len(issues)} potential security issues inside the custom_detectors file: {issues}") |
| 97 | + raise ImportError(f"Unsafe code detected in custom_detectors:\n" + "\n".join(issues)) |
| 98 | + |
| 99 | + import custom_detectors.custom_detectors as custom_detectors |
| 100 | + |
| 101 | + self.registry = {name: obj for name, obj |
| 102 | + in inspect.getmembers(custom_detectors, inspect.isfunction) |
| 103 | + if not name.startswith("_")} |
| 104 | + self.function_needs_headers = {name: "headers" in inspect.signature(obj).parameters for name, obj in self.registry.items() } |
| 105 | + logger.info(f"Registered the following custom detectors: {self.registry.keys()}") |
| 106 | + |
| 107 | + def handle_request(self, content: str, detector_params: dict, headers: dict) -> List[ContentAnalysisResponse]: |
| 108 | + detections = [] |
| 109 | + if "custom" in detector_params and isinstance(detector_params["custom"], (list, str)): |
| 110 | + custom_functions = detector_params["custom"] |
| 111 | + custom_functions = [custom_functions] if isinstance(custom_functions, str) else custom_functions |
| 112 | + for custom_function in custom_functions: |
| 113 | + if self.registry.get(custom_function): |
| 114 | + try: |
| 115 | + func_headers = headers if self.function_needs_headers.get(custom_function) else None |
| 116 | + result = custom_func_wrapper(self.registry[custom_function], custom_function, content, func_headers) |
| 117 | + if result is not None: |
| 118 | + detections.append(result) |
| 119 | + except Exception as e: |
| 120 | + logger.error(e) |
| 121 | + raise HTTPException(status_code=400, detail="Detection error, check detector logs") |
| 122 | + else: |
| 123 | + raise HTTPException(status_code=400, detail=f"Unrecognized custom function: {custom_function}") |
| 124 | + return detections |
| 125 | + |
0 commit comments