import spaces # MUST come before torch / any CUDA-touching import import torch import gradio as gr from transformers import AutoModelForSequenceClassification, AutoTokenizer MODEL_ID = "CIRCL/vulnerability-attack-technique-classification-roberta-base" # MITRE ATT&CK Enterprise technique names for the 53 parent techniques # the model was trained on (sub-techniques collapsed to parents). TECHNIQUE_NAMES = { "T1003": "OS Credential Dumping", "T1005": "Data from Local System", "T1021": "Remote Services", "T1036": "Masquerading", "T1040": "Network Sniffing", "T1041": "Exfiltration Over C2 Channel", "T1046": "Network Service Discovery", "T1055": "Process Injection", "T1059": "Command and Scripting Interpreter", "T1068": "Exploitation for Privilege Escalation", "T1070": "Indicator Removal", "T1071": "Application Layer Protocol", "T1078": "Valid Accounts", "T1082": "System Information Discovery", "T1083": "File and Directory Discovery", "T1087": "Account Discovery", "T1091": "Replication Through Removable Media", "T1098": "Account Manipulation", "T1105": "Ingress Tool Transfer", "T1106": "Native API", "T1110": "Brute Force", "T1133": "External Remote Services", "T1136": "Create Account", "T1185": "Browser Session Hijacking", "T1189": "Drive-by Compromise", "T1190": "Exploit Public-Facing Application", "T1202": "Indirect Command Execution", "T1203": "Exploitation for Client Execution", "T1204": "User Execution", "T1210": "Exploitation of Remote Services", "T1211": "Exploitation for Defense Evasion", "T1212": "Exploitation for Credential Access", "T1485": "Data Destruction", "T1486": "Data Encrypted for Impact", "T1496": "Resource Hijacking", "T1497": "Virtualization/Sandbox Evasion", "T1498": "Network Denial of Service", "T1499": "Endpoint Denial of Service", "T1505": "Server Software Component", "T1528": "Steal Application Access Token", "T1542": "Pre-OS Boot", "T1543": "Create or Modify System Process", "T1548": "Abuse Elevation Control Mechanism", "T1550": "Use Alternate Authentication Material", "T1552": "Unsecured Credentials", "T1555": "Credentials from Password Stores", "T1557": "Adversary-in-the-Middle", "T1563": "Remote Service Session Hijacking", "T1565": "Data Manipulation", "T1566": "Phishing", "T1574": "Hijack Execution Flow", "T1608": "Stage Capabilities", "T1685": "Disable or Modify Tools", } tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID).to("cuda") model.eval() id2label = model.config.id2label @spaces.GPU(duration=30) def classify_cve(description: str, top_k: int = 10, threshold: float = 0.5) -> dict: """Classify a CVE description into MITRE ATT&CK techniques. Given a free-text vulnerability description, this function predicts the most relevant MITRE ATT&CK Enterprise techniques using a RoBERTa-base multi-label classifier trained on a curated gold set of 1,207 CVEs. Args: description: A free-text vulnerability (CVE) description. top_k: Maximum number of techniques to return (sorted by probability). threshold: Minimum sigmoid probability to include a technique in the "high-confidence" set. All predicted probabilities are returned regardless; this only affects the summary text. """ if not description or not description.strip(): return {}, "Please enter a CVE description." inputs = tokenizer( description, truncation=True, max_length=512, return_tensors="pt", ).to("cuda") with torch.no_grad(): probs = torch.sigmoid(model(**inputs).logits)[0] # Sort by descending probability sorted_indices = probs.argsort(descending=True) top_k = min(top_k, len(sorted_indices)) results = {} high_conf = [] for idx in sorted_indices[:top_k]: tid = id2label[int(idx)] prob = float(probs[idx]) label_full = f"{tid} — {TECHNIQUE_NAMES.get(tid, 'Unknown')}" results[label_full] = prob if prob >= threshold: high_conf.append(f"{tid} ({prob:.3f})") summary = ( f"High-confidence techniques (≥ {threshold:.1f}): {', '.join(high_conf)}" if high_conf else f"No techniques above threshold {threshold:.1f}. Top: {id2label[int(sorted_indices[0])]} ({float(probs[sorted_indices[0]]):.3f})" ) return results, summary CSS = """ #col-container { max-width: 1100px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ EXAMPLES = [ [ "Zoho ManageEngine ServiceDesk Plus before 11306, ServiceDesk Plus MSP " "before 10530, and SupportCenter Plus before 11014 are vulnerable to " "unauthenticated remote code execution." ], [ "A buffer overflow vulnerability in the Simple Network Management " "Protocol (SNMP) service of Cisco IOS allows remote attackers to " "execute arbitrary code on the affected device via crafted SNMP packets." ], [ "An SQL injection vulnerability in the login page of a web application " "allows unauthenticated remote attackers to bypass authentication and " "access the administrator dashboard by injecting malicious SQL " "queries through the username field." ], [ "A path traversal vulnerability in a file download component allows " "remote attackers to read arbitrary files on the system via specially " "crafted requests containing dot-dot sequences in the filename " "parameter." ], [ "A cross-site scripting (XSS) vulnerability in the comment section " "of a popular forum allows attackers to inject malicious JavaScript " "that executes in the browser of any visitor viewing the affected " "page, enabling session token theft." ], ] with gr.Blocks(title="CVE to MITRE ATT&CK Mapper") as demo: gr.Markdown( "# CVE to MITRE ATT&CK Technique Mapper\n" "Enter a vulnerability (CVE) description and get the most relevant " "MITRE ATT&CK Enterprise techniques, predicted by a RoBERTa-base " "multi-label classifier trained on a curated gold set of 1,207 CVEs " "with expert MITRE CTID labels.\n\n" "📖 [Paper](https://arxiv.org/abs/2607.25572) | " "🤖 [Model](https://huggingface.co/CIRCL/vulnerability-attack-technique-classification-roberta-base) | " "💻 [Code (VulnTrain)](https://github.com/vulnerability-lookup/VulnTrain)" ) with gr.Column(elem_id="col-container"): with gr.Row(): description = gr.Textbox( label="CVE / Vulnerability Description", placeholder="Paste a CVE description here...", lines=6, scale=4, ) run_btn = gr.Button("Classify", variant="primary", scale=1) with gr.Accordion("Advanced settings", open=False): top_k = gr.Slider( minimum=1, maximum=53, value=10, step=1, label="Top-K techniques to show", ) threshold = gr.Slider( minimum=0.0, maximum=1.0, value=0.5, step=0.05, label="Confidence threshold (for high-confidence summary)", ) label_output = gr.Label( label="Predicted MITRE ATT&CK Techniques", num_top_classes=10, ) summary_output = gr.Textbox( label="Summary", interactive=False, lines=2, ) run_btn.click( fn=classify_cve, inputs=[description, top_k, threshold], outputs=[label_output, summary_output], api_name="classify", ) description.submit( fn=classify_cve, inputs=[description, top_k, threshold], outputs=[label_output, summary_output], ) gr.Examples( examples=EXAMPLES, inputs=[description], outputs=[label_output, summary_output], fn=classify_cve, cache_examples=True, cache_mode="lazy", ) demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)