martynattakit/CodeSentinel-CWE_Classification

1"""2pipeline/router.py3The glue layer — routes input through the correct models and returns4a unified output card. This is the only file that knows about all three5pipeline components.6"""7 8from __future__ import annotations9import re10import time11from typing import Optional12 13# ── ATLAS signal keywords ─────────────────────────────────────────────────────14 15ATLAS_SIGNALS = {16    "model", "llm", "language model", "neural network", "neural net",17    "transformer", "embedding", "embeddings", "inference", "training data",18    "fine-tun", "fine_tun", "finetun",19    "prompt injection", "prompt engineer", "jailbreak", "adversarial",20    "data poison", "model inversion", "membership inference", "model extraction",21    "model stealing", "backdoor", "trojan",22    "api key", "openai", "anthropic", "huggingface", "ollama", "vector db",23    "vector database", "chromadb", "rag", "retrieval", "mcp server",24    "model registry", "mlflow", "weights", "checkpoint", "safetensor",25    "machine learning", "deep learning", "gradient", "loss function",26    "classifier", "tokenizer", "attention", "gpt", "bert", "claude", "gemini",27}28 29# ── Code detection heuristics ─────────────────────────────────────────────────30 31CODE_PATTERNS = [32    r"def\s+\w+\s*\(",33    r"function\s+\w+\s*\(",34    r"void\s+\w+\s*\(",35    r"int\s+\w+\s*\(",36    r"public\s+\w+\s+\w+\s*\(",37    r"private\s+\w+\s+\w+\s*\(",38    r"#include\s*[<\"]",39    r"import\s+[\w.]+",40    r"require\s*\(",41    r"SELECT\s+.+FROM",42    r"<\?php",43    r"func\s+\w+\s*\(",44    r"fn\s+\w+\s*\(",45]46 47_CODE_RE = re.compile("|".join(CODE_PATTERNS), re.IGNORECASE)48 49def _is_code(text: str) -> bool:50    if _CODE_RE.search(text):51        return True52    code_chars = sum(text.count(c) for c in "{}[]();")53    if len(text) > 0 and code_chars / len(text) > 0.05:54        return True55    lines = text.split("\n")56    indented = sum(1 for l in lines if l.startswith("    ") or l.startswith("\t"))57    if len(lines) > 3 and indented / len(lines) > 0.3:58        return True59    return False60 61 62def _has_atlas_signals(text: str) -> bool:63    text_lower = text.lower()64    return any(signal in text_lower for signal in ATLAS_SIGNALS)65 66 67# ── Router ────────────────────────────────────────────────────────────────────68 69class Router:70    def __init__(self):71        self._classifier    = None72        self._code_analyzer = None73        self._atlas_matcher = None74 75    def _get_classifier(self):76        if self._classifier is None:77            from pipeline.classifier import CWEClassifier78            self._classifier = CWEClassifier()79        return self._classifier80 81    def _get_code_analyzer(self):82        if self._code_analyzer is None:83            from pipeline.code_analyzer import CodeAnalyzer84            self._code_analyzer = CodeAnalyzer()85        return self._code_analyzer86 87    def _get_atlas_matcher(self):88        if self._atlas_matcher is None:89            from pipeline.atlas_matcher import ATLASMatcher90            self._atlas_matcher = ATLASMatcher()91        return self._atlas_matcher92 93    def run(self, user_input: str) -> dict:94        if not user_input or not user_input.strip():95            raise ValueError("Input cannot be empty.")96 97        start = time.time()98 99        run_atlas = _has_atlas_signals(user_input)100        is_code   = _is_code(user_input)101 102        description = user_input103        input_type  = "text"104        code_analysis_warning = None105 106        if is_code:107            try:108                description = self._get_code_analyzer().analyze(user_input)109                input_type  = "code"110            except Exception as e:111                error_str = str(e).lower()112                if "cuda" in error_str or "bitsandbytes" in error_str or "gpu" in error_str:113                    description = user_input114                    input_type  = "code"115                    code_analysis_warning = (116                        "GPU not available — code passed directly to classifier. "117                        "For best results, describe the vulnerability in plain English."118                    )119                else:120                    raise121 122        cwe_result = self._get_classifier().classify(description)123 124        atlas_result = None125        if run_atlas:126            try:127                atlas_result = self._get_atlas_matcher().match(user_input)128            except Exception:129                pass130 131        elapsed = round(time.time() - start, 2)132 133        existing_warning = cwe_result.get("warning")134        if code_analysis_warning and existing_warning:135            final_warning = f"{code_analysis_warning} | {existing_warning}"136        elif code_analysis_warning:137            final_warning = code_analysis_warning138        else:139            final_warning = existing_warning140 141        return _build_output_card(142            raw_input=user_input,143            input_type=input_type,144            description=description,145            cwe_result=cwe_result,146            atlas_result=atlas_result,147            elapsed_seconds=elapsed,148            warning_override=final_warning,149        )150 151 152def _build_output_card(153    raw_input: str,154    input_type: str,155    description: str,156    cwe_result: dict,157    atlas_result: Optional[dict],158    elapsed_seconds: float,159    warning_override: Optional[str] = None,160) -> dict:161    top1 = cwe_result["top1"]162 163    return {164        "cwe_id":      top1["cwe_id"],165        "cwe_name":    top1["description"],166        "severity":    top1["severity"],167        "confidence":  top1["confidence"],168        "description": description,169        "alternatives": cwe_result["top3"][1:],170        "atlas_match": atlas_result,171        "input_type":  input_type,172        "warning":     warning_override if warning_override is not None else cwe_result.get("warning"),173        "elapsed_s":   elapsed_seconds,174    }175 176 177# ── Module-level singleton ────────────────────────────────────────────────────178 179_router: Optional[Router] = None180 181def get_router() -> Router:182    global _router183    if _router is None:184        _router = Router()185    return _router186 187 188def route(user_input: str) -> dict:189    return get_router().run(user_input)