martynattakit/CodeSentinel-CWE_Classification

1"""2pipeline/classifier.py3RoBERTa-based CWE classifier — wraps the fine-tuned model for inference.4Input:  natural language vulnerability description (str)5Output: list of top-k CWE predictions with confidence scores6"""7 8from __future__ import annotations9import json10from pathlib import Path11from typing import Optional12import torch13from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline14 15# ── Constants ────────────────────────────────────────────────────────────────16 17HF_REPO = "martynattakit/vuln-classifier-roberta"18MAX_LENGTH = 25619TOP_K = 320 21# CWEs where the model is known to be unreliable (from eval report)22LOW_CONFIDENCE_CWES = {23    "CWE-77",   # 0 samples in training — never predicts correctly24    "CWE-863",  # F1 0.60 — overlaps with CWE-86225}26 27CWE_DESCRIPTIONS = {28    "CWE-787": "Out-of-bounds Write",29    "CWE-79":  "Cross-site Scripting (XSS)",30    "CWE-89":  "SQL Injection",31    "CWE-416": "Use After Free",32    "CWE-78":  "OS Command Injection",33    "CWE-20":  "Improper Input Validation",34    "CWE-125": "Out-of-bounds Read",35    "CWE-22":  "Path Traversal",36    "CWE-352": "Cross-Site Request Forgery (CSRF)",37    "CWE-434": "Unrestricted File Upload",38    "CWE-862": "Missing Authorization",39    "CWE-476": "NULL Pointer Dereference",40    "CWE-287": "Improper Authentication",41    "CWE-190": "Integer Overflow",42    "CWE-502": "Deserialization of Untrusted Data",43    "CWE-77":  "Command Injection",44    "CWE-119": "Buffer Overflow (Generic)",45    "CWE-798": "Hardcoded Credentials",46    "CWE-918": "Server-Side Request Forgery (SSRF)",47    "CWE-306": "Missing Authentication",48    "CWE-362": "Race Condition",49    "CWE-269": "Improper Privilege Management",50    "CWE-94":  "Code Injection",51    "CWE-863": "Incorrect Authorization",52    "CWE-276": "Incorrect Default Permissions",53}54 55SEVERITY_MAP = {56    "CWE-787": "HIGH",57    "CWE-79":  "MEDIUM",58    "CWE-89":  "HIGH",59    "CWE-416": "HIGH",60    "CWE-78":  "HIGH",61    "CWE-20":  "MEDIUM",62    "CWE-125": "MEDIUM",63    "CWE-22":  "HIGH",64    "CWE-352": "MEDIUM",65    "CWE-434": "HIGH",66    "CWE-862": "HIGH",67    "CWE-476": "MEDIUM",68    "CWE-287": "HIGH",69    "CWE-190": "MEDIUM",70    "CWE-502": "HIGH",71    "CWE-77":  "HIGH",72    "CWE-119": "HIGH",73    "CWE-798": "CRITICAL",74    "CWE-918": "HIGH",75    "CWE-306": "CRITICAL",76    "CWE-362": "MEDIUM",77    "CWE-269": "HIGH",78    "CWE-94":  "HIGH",79    "CWE-863": "HIGH",80    "CWE-276": "MEDIUM",81}82 83# ── Classifier class ─────────────────────────────────────────────────────────84 85class CWEClassifier:86    """87    Wraps the fine-tuned RoBERTa model for CWE classification.88    Lazy-loaded on first call — fast import, slow first inference.89    """90 91    def __init__(self, repo: str = HF_REPO, device: Optional[str] = None):92        self.repo = repo93        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")94        self._pipeline = None95 96    def _load(self):97        """Lazy load the model on first inference call."""98        if self._pipeline is not None:99            return100 101        print(f"[CWEClassifier] Loading model from {self.repo}...")102        self._pipeline = pipeline(103            "text-classification",104            model=self.repo,105            tokenizer=self.repo,106            device=0 if self.device == "cuda" else -1,107            top_k=TOP_K,108            truncation=True,109            max_length=MAX_LENGTH,110        )111        print("[CWEClassifier] Model loaded.")112 113    def classify(self, text: str) -> dict:114        """115        Classify a vulnerability description.116 117        Args:118            text: Natural language vulnerability description.119                  Should follow the structured format:120                  "This function performs X on Y without Z, which may allow..."121 122        Returns:123            {124                "top1":        { "cwe_id", "description", "severity", "confidence" },125                "top3":        [ { "cwe_id", "description", "severity", "confidence" }, ... ],126                "warning":     str | None,   # set if top1 is a known weak class127                "raw_scores":  { cwe_id: score, ... }128            }129        """130        self._load()131 132        if not text or not text.strip():133            raise ValueError("Input text cannot be empty.")134 135        raw = self._pipeline(text[:MAX_LENGTH * 4])  # rough char limit before tokenizer136        predictions = raw[0]  # list of {label, score}137 138        results = []139        for pred in predictions:140            cwe_id = pred["label"]141            confidence = round(pred["score"], 4)142            results.append({143                "cwe_id":      cwe_id,144                "description": CWE_DESCRIPTIONS.get(cwe_id, "Unknown"),145                "severity":    SEVERITY_MAP.get(cwe_id, "UNKNOWN"),146                "confidence":  confidence,147            })148 149        top1 = results[0]150 151        # Warn if top1 is a known unreliable class152        warning = None153        if top1["cwe_id"] in LOW_CONFIDENCE_CWES:154            warning = (155                f"{top1['cwe_id']} has limited training data — "156                f"confidence may be unreliable. Review top-3 predictions."157            )158 159        # Also warn if top1 confidence is low160        if top1["confidence"] < 0.5 and warning is None:161            warning = (162                f"Low confidence ({top1['confidence']:.0%}) — "163                f"input may not match known vulnerability patterns."164            )165 166        return {167            "top1":       top1,168            "top3":       results,169            "warning":    warning,170            "raw_scores": {p["label"]: round(p["score"], 4) for p in predictions},171        }172 173 174# ── Module-level singleton ───────────────────────────────────────────────────175 176_classifier: Optional[CWEClassifier] = None177 178def get_classifier() -> CWEClassifier:179    """Return the module-level singleton classifier."""180    global _classifier181    if _classifier is None:182        _classifier = CWEClassifier()183    return _classifier184 185 186def classify(text: str) -> dict:187    """Convenience function — classify without instantiating manually."""188    return get_classifier().classify(text)189 190 191# ── CLI test ─────────────────────────────────────────────────────────────────192 193if __name__ == "__main__":194    test_cases = [195        "This function constructs a SQL query by concatenating user-controlled input without parameterization, which may allow an attacker to inject arbitrary SQL commands.",196        "This function reflects user-supplied data into the HTTP response without encoding, which may allow an attacker to inject malicious scripts.",197        "This function performs operations on a memory buffer without verifying bounds, which may allow an attacker to read or write out-of-bounds memory.",198    ]199 200    clf = CWEClassifier()201    for text in test_cases:202        result = clf.classify(text)203        print(f"Input: {text[:60]}...")204        print(f"  Top-1: {result['top1']['cwe_id']} ({result['top1']['severity']}) — {result['top1']['confidence']:.1%}")205        if result["warning"]:206            print(f"  ⚠ {result['warning']}")207        print()208