martynattakit/CodeSentinel-CWE_Classification

1"""2pipeline/atlas_matcher.py3ATLAS pattern matcher — RAG over hand-crafted MITRE case studies.4NOT a classifier. Returns the closest matching ATLAS technique5with the evidence that led to the match.6 7Input:  raw user input (str)8Output: matched technique dict with cited evidence, or None9"""10 11from __future__ import annotations12import json13import os14from pathlib import Path15from typing import Optional16 17# ── Constants ────────────────────────────────────────────────────────────────18 19# Path to atlas_cases.json relative to project root20ATLAS_CASES_PATH = Path(__file__).parent.parent / "data" / "atlas_cases.json"21 22# Minimum signal matches required before we attempt a match23MIN_SIGNAL_MATCHES = 224 25# Confidence levels based on signal match count26CONFIDENCE_THRESHOLDS = {27    "HIGH":   5,28    "MEDIUM": 3,29    "LOW":    2,30}31 32# ── ATLASMatcher class ───────────────────────────────────────────────────────33 34class ATLASMatcher:35    """36    Lightweight RAG matcher using keyword signals + semantic overlap.37    No vector database needed — corpus is small enough for direct matching.38    """39 40    def __init__(self, cases_path: Path = ATLAS_CASES_PATH):41        self.cases_path = cases_path42        self._cases = None43 44    def _load(self):45        """Load ATLAS cases from JSON file."""46        if self._cases is not None:47            return48 49        if not self.cases_path.exists():50            raise FileNotFoundError(51                f"ATLAS cases not found at {self.cases_path}. "52                "Make sure data/atlas_cases.json exists."53            )54 55        with open(self.cases_path, "r") as f:56            self._cases = json.load(f)57 58        print(f"[ATLASMatcher] Loaded {len(self._cases)} ATLAS techniques.")59 60    def match(self, text: str) -> Optional[dict]:61        """62        Find the best matching ATLAS technique for the given input.63 64        Args:65            text: Raw user input (code or natural language).66 67        Returns:68            Match dict or None if no confident match found:69            {70                "atlas_id":      str,71                "technique":     str,72                "tactic":        str,73                "confidence":    "HIGH" | "MEDIUM" | "LOW",74                "matched_signals": [str, ...],75                "description":   str,76                "mitigations":   [str, ...],77                "real_world":    str | None,78                "reasoning":     str,79            }80        """81        self._load()82 83        text_lower = text.lower()84 85        best_match = None86        best_score = 087 88        for case in self._cases:89            signals = [s.lower() for s in case.get("signals", [])]90            matched = [s for s in signals if s in text_lower]91            score = len(matched)92 93            if score > best_score:94                best_score = score95                best_match = (case, matched)96 97        # Require minimum signal matches98        if best_score < MIN_SIGNAL_MATCHES or best_match is None:99            return None100 101        case, matched_signals = best_match102 103        # Determine confidence level104        confidence = "LOW"105        for level, threshold in CONFIDENCE_THRESHOLDS.items():106            if best_score >= threshold:107                confidence = level108                break109 110        # Build reasoning string111        signals_str = ", ".join(f'"{s}"' for s in matched_signals[:5])112        reasoning = (113            f"Matched {best_score} signal(s) from the input: {signals_str}. "114            f"This pattern is consistent with {case['technique']} "115            f"({case['atlas_id']}) under the {case['tactic']} tactic."116        )117 118        return {119            "atlas_id":        case["atlas_id"],120            "technique":       case["technique"],121            "tactic":          case["tactic"],122            "confidence":      confidence,123            "matched_signals": matched_signals,124            "description":     case["description"],125            "mitigations":     case.get("mitigations", []),126            "real_world":      case.get("real_world"),127            "reasoning":       reasoning,128        }129 130    def match_top_k(self, text: str, k: int = 3) -> list[dict]:131        """132        Return top-k ATLAS matches ranked by signal overlap.133        Useful for debugging or showing multiple possible techniques.134        """135        self._load()136 137        text_lower = text.lower()138        scored = []139 140        for case in self._cases:141            signals = [s.lower() for s in case.get("signals", [])]142            matched = [s for s in signals if s in text_lower]143            score = len(matched)144            if score >= MIN_SIGNAL_MATCHES:145                scored.append((score, case, matched))146 147        scored.sort(key=lambda x: -x[0])148 149        results = []150        for score, case, matched in scored[:k]:151            confidence = "LOW"152            for level, threshold in CONFIDENCE_THRESHOLDS.items():153                if score >= threshold:154                    confidence = level155                    break156 157            results.append({158                "atlas_id":        case["atlas_id"],159                "technique":       case["technique"],160                "tactic":          case["tactic"],161                "confidence":      confidence,162                "signal_count":    score,163                "matched_signals": matched,164            })165 166        return results167 168 169# ── Module-level singleton ───────────────────────────────────────────────────170 171_matcher: Optional[ATLASMatcher] = None172 173def get_matcher() -> ATLASMatcher:174    """Return the module-level singleton matcher."""175    global _matcher176    if _matcher is None:177        _matcher = ATLASMatcher()178    return _matcher179 180 181def match(text: str) -> Optional[dict]:182    """Convenience function — match without instantiating manually."""183    return get_matcher().match(text)184 185 186# ── CLI test ─────────────────────────────────────────────────────────────────187 188if __name__ == "__main__":189    test_inputs = [190        "The LLM API accepts user prompts without sanitization, allowing prompt injection to bypass system instructions and jailbreak the model.",191        "The training data pipeline accepts contributions from external sources without validation, enabling data poisoning attacks.",192        "The model inference API does not rate limit requests, allowing an attacker to extract the model through repeated queries.",193        "SQL injection in the login form allows unauthenticated database access.",  # should return None194    ]195 196    matcher = ATLASMatcher()197    for text in test_inputs:198        result = matcher.match(text)199        print(f"Input: {text[:80]}...")200        if result:201            print(f"  Match:      {result['atlas_id']} — {result['technique']}")202            print(f"  Confidence: {result['confidence']}")203            print(f"  Signals:    {result['matched_signals']}")204            print(f"  Reasoning:  {result['reasoning']}")205        else:206            print("  No ATLAS match (below confidence threshold)")207        print()208