update

README.md

@@ -19,12 +19,14 @@ actor or a non-existent CVE — it **refuses to fabricate** (evidence-gated hone
 ## How it works (RAG)
 1. Your message is routed to an entity: actor / CVE / technique / malware / report.
 2. The matching record is retrieved from the bundled corpus (`corpus.json`, built from
-   real MITRE ATT&CK + CISA KEV/NVD data) and formatted into the exact template the model
-   was trained on.
-3. TextScout analyzes that record. Empty lookup → honest refusal.
+   real MITRE ATT&CK + CISA KEV/NVD data). **Any CVE not in the snapshot is fetched LIVE
+   from NVD + CISA KEV** — so fresh CVEs that were never in train/valid/test work too.
+3. The record is formatted into the exact template the model was trained on; TextScout
+   analyzes it. Empty lookup (unknown actor, non-existent CVE) → honest refusal.
 
 ## Files
-`Dockerfile` · `app.py` · `requirements.txt` · `corpus.json` (retrieval data).
+`Dockerfile` · `app.py` · `requirements.txt` · `corpus.json` (retrieval data) · `robot.png`.
+Optional Space secret `NVD_API_KEY` raises the NVD rate limit (works without it).
 
 ## Push updates
 With the `hf` CLI authenticated, from the `osint-project` root:

app.py

@@ -1,21 +1,23 @@
 """
 TextScout — red-team threat-intel chatbot (Streamlit, for Hugging Face Spaces).
 
-Flow per turn: user asks in natural language -> RAG step retrieves the matching
-real record from the bundled corpus (MITRE ATT&CK groups/techniques/software,
-CISA KEV/NVD CVEs) -> builds the EXACT training-format prompt -> TextScout
-(Llama-3.2-3B + LoRA adapter Maximuz23/Text-OSINT) analyzes it. If retrieval is
-empty (unknown actor / non-existent CVE) the record is an empty lookup and the
-model refuses instead of fabricating — evidence-gated honesty, shown live.
-
-The record templates below mirror scripts/build_grounded_records.py byte-for-byte;
-the SYSTEM_PROMPT mirrors ai-test.ipynb c03. Do not drift either or the model
-goes off-distribution.
+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 pathlib
+import urllib.request
+import urllib.parse
 
 import streamlit as st
 import torch
@@ -28,6 +30,8 @@ USE_GPU = torch.cuda.is_available()
 BASE_MODEL = "unsloth/Llama-3.2-3B-Instruct-bnb-4bit" if USE_GPU else "unsloth/Llama-3.2-3B-Instruct"
 MAX_NEW_TOKENS_DEFAULT = 384
 HF_TOKEN = os.environ.get("HF_TOKEN")
+ROBOT = str(pathlib.Path(__file__).parent / "robot.png")
+CORPUS_PATH = pathlib.Path(__file__).parent / "corpus.json"
 
 SYSTEM_PROMPT = (
     "You are an expert cybersecurity analyst specializing in Text OSINT and threat "
@@ -45,9 +49,7 @@ REFUSAL_HINTS = re.compile(
     re.I,
 )
 
-CORPUS_PATH = pathlib.Path(__file__).parent / "corpus.json"
 
-# --- model -----------------------------------------------------------------------
 @st.cache_resource(show_spinner="Loading TextScout (Llama-3.2-3B + adapter)… first load is slow.")
 def load_model():
     tok = AutoTokenizer.from_pretrained(BASE_MODEL, token=HF_TOKEN)
@@ -58,7 +60,6 @@ def load_model():
     if USE_GPU:
         kwargs.update(device_map={"": 0}, torch_dtype=torch.float16)
     else:
-        # bf16 ~6GB (fits the free 16GB CPU Space) and runs on CPU, unlike fp16.
         kwargs.update(torch_dtype=torch.bfloat16, low_cpu_mem_usage=True)
     base = AutoModelForCausalLM.from_pretrained(BASE_MODEL, **kwargs)
     base.config.use_cache = True
@@ -156,7 +157,65 @@ def t_soft(s):
             f"Documented user groups: {', '.join(g) if g else 'none recorded'}")
 
 
-# --- retrieval (RAG) -------------------------------------------------------------
+# --- live retrieval: NVD + CISA KEV (fresh CVEs not in the bundle) ---------------
+def _get_json(url, headers=None, timeout=15):
+    req = urllib.request.Request(url, headers={"User-Agent": "TextScout-demo", **(headers or {})})
+    with urllib.request.urlopen(req, timeout=timeout) as r:
+        return json.loads(r.read())
+
+
+@st.cache_data(ttl=3600, show_spinner=False)
+def _kev_catalog():
+    try:
+        data = _get_json("https://www.cisa.gov/sites/default/files/feeds/known_exploited_vulnerabilities.json")
+        return {v["cveID"].upper(): v for v in data.get("vulnerabilities", [])}
+    except Exception:
+        return {}                       # KEV is supplementary; degrade gracefully
+
+
+@st.cache_data(ttl=3600, show_spinner=False)
+def _nvd_lookup(cid):
+    # raises on network error (so we can show "try again" rather than a false refusal);
+    # returns None only when NVD answers with zero records.
+    headers = {"apiKey": os.environ["NVD_API_KEY"]} if os.environ.get("NVD_API_KEY") else None
+    url = "https://services.nvd.nist.gov/rest/json/cves/2.0?cveId=" + urllib.parse.quote(cid)
+    vulns = _get_json(url, headers).get("vulnerabilities") or []
+    return vulns[0]["cve"] if vulns else None
+
+
+def live_cve(cid):
+    """Build a CVE record from live NVD + KEV, or None if genuinely not found."""
+    nvd = _nvd_lookup(cid)
+    kev = _kev_catalog().get(cid)
+    if not nvd and not kev:
+        return None
+    desc, cvss, cwes = "", "", []
+    if nvd:
+        desc = next((x["value"] for x in nvd.get("descriptions", []) if x["lang"] == "en"), "")
+        metrics = nvd.get("metrics", {})
+        for k in ("cvssMetricV31", "cvssMetricV30", "cvssMetricV2"):
+            if metrics.get(k):
+                cd = metrics[k][0]["cvssData"]
+                cvss = f"{cd.get('baseScore', '')} {cd.get('baseSeverity', '')} {cd.get('vectorString', '')}".strip()
+                break
+        cwes = list(dict.fromkeys(
+            x["value"] for w in nvd.get("weaknesses", []) for x in w.get("description", [])
+            if x.get("value", "").startswith("CWE")))
+    return {
+        "cve_id": cid,
+        "name": (kev.get("vulnerabilityName") if kev else "") or "",
+        "vendor": kev.get("vendorProject", "") if kev else "",
+        "product": kev.get("product", "") if kev else "",
+        "nvd_description": desc,
+        "kev_description": kev.get("shortDescription", "") if kev else "",
+        "cvss": cvss,
+        "cwes": cwes,
+        "ransomware_use": kev.get("knownRansomwareCampaignUse", "") if kev else "",
+        "required_action": kev.get("requiredAction", "") if kev else "",
+    }
+
+
+# --- bundled retrieval -----------------------------------------------------------
 IOC_RE = re.compile(r"hxxp|https?://|\b\d{1,3}\[?\.\]?\d{1,3}\[?\.\]?\d{1,3}\[?\.\]?\d{1,3}\b|\b[a-f0-9]{32,64}\b", re.I)
 ACTOR_KW = re.compile(r"\b(actor|group|apt|adversary|profile|who is|what did|techniques does|ttps|campaign)\b", re.I)
 REPORT_KW = re.compile(r"\b(extract|report|ioc|indicator|infrastructure|c2|c&c|phish)\b", re.I)
@@ -173,15 +232,14 @@ def load_corpus():
     actor_idx, soft_idx = {}, {}
     for a in c["actors"]:
         for key in [a["name"], *a.get("aliases", [])]:
-            actor_idx.setdefault(key.lower(), a)
+            if key:
+                actor_idx.setdefault(key.lower(), a)
     for s in c["software"]:
-        if len(s["name"]) >= 5:          # drop short/common tool names (net, at, Reg…)
+        if len(s["name"]) >= 5:
             soft_idx.setdefault(s["name"].lower(), s)
     cve_idx = {x["cve_id"].upper(): x for x in c["cves"]}
     tech_idx = {x["attack_id"].upper(): x for x in c["techniques"]}
-    actor_re = _name_regex(list(actor_idx))
-    soft_re = _name_regex(list(soft_idx))
-    return actor_idx, cve_idx, tech_idx, soft_idx, actor_re, soft_re
+    return actor_idx, cve_idx, tech_idx, soft_idx, _name_regex(list(actor_idx)), _name_regex(list(soft_idx))
 
 
 def _longest(rx, q):
@@ -205,14 +263,14 @@ def _guess_name(q):
 
 HELP = (
     "I'm **TextScout** — I answer from retrieved threat-intel records. Try: "
-    "**profile an actor** (APT28), **assess a CVE** (CVE-2021-44228), **explain a technique** "
-    "(T1059.001), **identify malware** (Cobalt Strike), or **extract IOCs from a report**. "
-    "Pick a suggestion to see it."
+    "**profile an actor** (APT28), **assess a CVE** (any id — fresh ones go to live NVD), "
+    "**explain a technique** (T1059.001), **identify malware** (Cobalt Strike), or "
+    "**extract IOCs from a report**. Pick a suggestion to see it."
 )
 
 
 def route(q):
-    """Return (user_prompt, retrieved_record_or_None, source_label, help_or_None)."""
+    """Return (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()
 
     m = re.search(r"CVE-\d{4}-\d{4,}", q, re.I)
@@ -220,14 +278,20 @@ def route(q):
         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, "NVD + CISA KEV", None
-        p = t_cve_empty(cid); return p, p, "NVD + CISA KEV (no record)", None
+            return p, p, "bundled MITRE/KEV record", None
+        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
 
     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
 
-    if IOC_RE.search(q):                       # report extraction — no external lookup
+    if IOC_RE.search(q):
         return q, None, "the report you provided", None
 
     a = _longest(actor_re, q)
@@ -238,7 +302,7 @@ def route(q):
         rec = actor_idx[a.lower()]; p = rec.get("_prompt") or t_actor(rec)
         return p, p, "MITRE ATT&CK", None
 
-    if ACTOR_KW.search(q):                      # actor intent, no corpus match -> honest refusal
+    if ACTOR_KW.search(q):
         p = t_actor_empty(_guess_name(q)); return p, p, "MITRE ATT&CK (no record)", None
     if REPORT_KW.search(q):
         return q, None, "the report you provided", None
@@ -246,16 +310,13 @@ def route(q):
 
 
 # --- UI --------------------------------------------------------------------------
-st.set_page_config(page_title="TextScout", page_icon="🛰️", layout="centered")
-st.title("🛰️ TextScout")
-st.caption(
-    "Red-team threat-intel assistant. Ask in plain language — TextScout retrieves the matching "
-    "MITRE ATT&CK / CVE record and analyzes it, and **refuses to fabricate** when there's no record."
-)
+st.set_page_config(page_title="TextScout", page_icon=ROBOT, layout="centered")
+st.markdown("<h1 style='text-align:center; margin:0.2rem 0 1.4rem'>TextScout</h1>",
+            unsafe_allow_html=True)
 
 with st.sidebar:
-    st.subheader("Model")
-    st.markdown(f"- **Base:** `{BASE_MODEL}`\n- **Adapter:** `{HF_REPO}`\n- **Device:** `{'GPU' if USE_GPU else 'CPU (free)'}`")
+    st.image(ROBOT, width=64)
+    st.markdown("## TextScout")
     st.divider()
     max_tokens = st.slider("Max new tokens", 128, 512, MAX_NEW_TOKENS_DEFAULT, 32,
                            help="Lower for snappier CPU responses.")
@@ -268,7 +329,7 @@ SUGGESTIONS = {
     "🧬 What does Kimsuky do?": "What techniques does the threat actor Kimsuky use?",
     "🛠️ 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-2021-44228": "Assess CVE-2021-44228 for offensive relevance",
+    "🔓 Assess CVE-2025-0282 (live)": "Assess CVE-2025-0282 for offensive relevance",
     "🌐 Extract IOCs from a report": ("Extract the IOCs from this abuse.ch report and assess red "
                                        "team relevance: hxxp://27.204.192.167:41628/i was flagged "
                                        "as a malware_download host."),
@@ -277,7 +338,6 @@ SUGGESTIONS = {
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-# empty-state suggestion chips
 clicked = None
 if not st.session_state.messages:
     st.markdown("##### Try one of these 👇")
@@ -286,29 +346,28 @@ if not st.session_state.messages:
         if cols[i % 2].button(label, use_container_width=True):
             clicked = text
 
-# render history
 for msg in st.session_state.messages:
-    with st.chat_message(msg["role"], avatar="🛰️" if msg["role"] == "assistant" else None):
+    with st.chat_message(msg["role"], avatar=ROBOT if msg["role"] == "assistant" else None):
         if msg.get("retrieved"):
-            with st.expander(f"🔎 RAG — record retrieved from {msg.get('source','source')}"):
+            with st.expander(f"🔎 RAG — record retrieved from {msg.get('source', 'source')}"):
                 st.code(msg["retrieved"])
         if msg.get("badge"):
             st.caption(msg["badge"])
         st.markdown(msg["content"])
 
-user_input = st.chat_input("Ask TextScout… e.g. profile APT28, assess CVE-2021-44228") or clicked
+user_input = st.chat_input("Ask TextScout… e.g. profile APT28, assess CVE-2025-0282") or clicked
 
 if user_input:
     st.session_state.messages.append({"role": "user", "content": user_input})
-    prompt, retrieved, source, help_text = route(user_input)
-    if help_text:
-        answer, badge = help_text, None
-    else:
-        with st.chat_message("assistant", avatar="🛰️"):
-            with st.spinner(f"🔎 Retrieving from {source} → TextScout analyzing…"):
+    with st.chat_message("assistant", avatar=ROBOT):
+        with st.spinner("🔎 Retrieving + analyzing…"):
+            prompt, retrieved, source, note = route(user_input)
+            if note:
+                answer, badge = note, None
+            else:
                 answer = generate(prompt, max_new_tokens=max_tokens)
-        badge = ("🛡️ No record retrieved → refused (honesty guardrail)"
-                 if REFUSAL_HINTS.search(answer) else "✅ Grounded in the retrieved record")
+                badge = ("🛡️ No record retrieved → refused (honesty guardrail)"
+                         if REFUSAL_HINTS.search(answer) else "✅ Grounded in the retrieved record")
     st.session_state.messages.append({"role": "assistant", "content": answer,
                                       "retrieved": retrieved, "source": source, "badge": badge})
     st.rerun()