update

app.py

@@ -1,20 +1,8 @@
-"""
-TextScout — red-team threat-intel chatbot (Streamlit, for Hugging Face Spaces).
-
-Per turn: user asks in natural language -> RAG retrieves the matching record
-(bundled MITRE ATT&CK / KEV snapshot, or LIVE NVD + CISA KEV for any CVE not in
-the snapshot) -> builds the EXACT training-format prompt -> the fine-tuned model
-(Llama-3.2-3B + LoRA Maximuz23/Text-OSINT) analyzes it. Empty lookup -> the model
-refuses instead of fabricating.
-
-Nothing is hardcoded: the answer is always model.generate(); only the retrieved
-record (real data) and a one-line help message are non-model text. Record
-templates mirror scripts/build_grounded_records.py; SYSTEM_PROMPT mirrors
-ai-test.ipynb c03 — do not drift either.
-"""
 import os
 import re
 import json
+import base64
+import random
 import pathlib
 import urllib.request
 import urllib.parse
@@ -317,7 +305,6 @@ def _vt(ioc, ioc_type):
     elif ioc_type == "file":
         path = "files/" + ioc
     elif ioc_type == "url":
-        import base64
         path = "urls/" + base64.urlsafe_b64encode(ioc.encode()).decode().strip("=")
     else:
         return None
@@ -398,7 +385,6 @@ def recent_threatfox_ioc():
     rows = d.get("data") if d.get("query_status") == "ok" else None
     if not isinstance(rows, list) or not rows:
         return None
-    import random
     pool = [r for r in rows if r.get("ioc_type") == "ip:port"] or rows
     return random.choice(pool).get("ioc", "").split(":")[0]
 
@@ -457,46 +443,56 @@ HELP = (
 )
 
 
+def _hit(prompt, source):
+    """A successful retrieval: the same text is both fed to the model and shown in the RAG expander."""
+    return prompt, prompt, source, None
+
+
 def route(q):
-    """Return (user_prompt, retrieved_record_or_None, source_label, note_or_None)."""
+    """Map a user query to (user_prompt, retrieved_record_or_None, source_label, note_or_None)."""
     actor_idx, cve_idx, tech_idx, soft_idx, actor_re, soft_re = load_corpus()
 
+    # CVE id -> bundled record, else live NVD + CISA KEV
     m = re.search(r"CVE-\d{4}-\d{4,}", q, re.I)
     if m:
         cid = m.group(0).upper()
         if cid in cve_idx:
-            rec = cve_idx[cid]; p = rec.get("_prompt") or t_cve(rec)
-            return p, p, "bundled MITRE/KEV record", None
+            rec = cve_idx[cid]
+            return _hit(rec.get("_prompt") or t_cve(rec), "bundled MITRE/KEV record")
         try:
             rec = live_cve(cid)
         except Exception:
             return None, None, None, "⚠️ Live NVD/CISA KEV lookup failed — please try again in a moment."
         if rec:
-            p = t_cve(rec); return p, p, "live NVD + CISA KEV", None
-        p = t_cve_empty(cid); return p, p, "live NVD + CISA KEV (no record)", None
+            return _hit(t_cve(rec), "live NVD + CISA KEV")
+        return _hit(t_cve_empty(cid), "live NVD + CISA KEV (no record)")
 
+    # ATT&CK technique id
     m = re.search(r"\bT\d{4}(?:\.\d{3})?\b", q)
     if m and m.group(0).upper() in tech_idx:
-        p = t_tech(tech_idx[m.group(0).upper()]); return p, p, "MITRE ATT&CK", None
+        return _hit(t_tech(tech_idx[m.group(0).upper()]), "MITRE ATT&CK")
 
+    # IOC -> live enrichment; if no keys, pass the raw report straight to the model
     if IOC_RE.search(q):
         ioc, ioc_type = detect_ioc(q)
         if ioc and _have_ioc_keys():
-            p = _safe(enrich_ioc, ioc, ioc_type)
-            if p:
-                return p, p, "live ThreatFox + OTX + VirusTotal", None
+            prompt = _safe(enrich_ioc, ioc, ioc_type)
+            if prompt:
+                return _hit(prompt, "live ThreatFox + OTX + VirusTotal")
         return q, None, "the report you provided", None
 
+    # actor or software name — longest match wins (so "Cobalt Strike" beats "Cobalt")
     a = _longest(actor_re, q)
     s = _longest(soft_re, q)
     if a or s:
         if len(s or "") > len(a or ""):
-            p = t_soft(soft_idx[s.lower()]); return p, p, "MITRE ATT&CK Software", None
-        rec = actor_idx[a.lower()]; p = rec.get("_prompt") or t_actor(rec)
-        return p, p, "MITRE ATT&CK", None
+            return _hit(t_soft(soft_idx[s.lower()]), "MITRE ATT&CK Software")
+        rec = actor_idx[a.lower()]
+        return _hit(rec.get("_prompt") or t_actor(rec), "MITRE ATT&CK")
 
+    # actor-intent but no match -> honest empty lookup; else a report keyword; else help
     if ACTOR_KW.search(q):
-        p = t_actor_empty(_guess_name(q)); return p, p, "MITRE ATT&CK (no record)", None
+        return _hit(t_actor_empty(_guess_name(q)), "MITRE ATT&CK (no record)")
     if REPORT_KW.search(q):
         return q, None, "the report you provided", None
     return None, None, None, HELP
@@ -523,7 +519,7 @@ SUGGESTIONS = {
     "🛠️ Who uses Cobalt Strike?": "Which threat actors use the Cobalt Strike tool?",
     "📖 Explain technique T1059.001": "Explain MITRE ATT&CK technique T1059.001",
     "🔓 Assess CVE-2025-0282 (live)": "Assess CVE-2025-0282 for offensive relevance",
-    "🌐 Look up a fresh live IOC": None,   # None -> pull a current abuse.ch IOC at click time
+    "🌐 Look up a fresh live IOC": None,
 }
 
 if "messages" not in st.session_state: