Local LLM Series, Part 3: Keeping Reference Data Current

"""Refresh the local reference library from public bulk data.

Downloads these public data sets and writes short, one-record-per-topic text
files that ask_grounded.py and ask_grounded_web.py can use as references:

    attack     MITRE ATT&CK Enterprise (techniques, detection guidance)
    kev        CISA Known Exploited Vulnerabilities
    sigma      SigmaHQ detection rules (release bundle)
    atomic     Atomic Red Team tests, as listed in its index
    lolbas     LOLBAS: Windows binaries and scripts attackers abuse
    loldrivers LOLDrivers: vulnerable and malicious Windows drivers (with hashes)

Output, in refs\\:

    attack_techniques.txt  attack_detections.txt  cisa_kev.txt  sigma_rules.txt
    atomic_tests.txt  lolbas.txt  loldrivers.txt  refresh_info.json  _downloads\\

Usage (PowerShell, .venv active):

    python -m pip install pyyaml              # once; Sigma, Atomic and LOLBAS need it
    python refresh_refs.py                    # download everything, rebuild refs\\
    python refresh_refs.py --only kev sigma   # refresh just those sources
    python refresh_refs.py --offline          # rebuild from the files already in refs\\_downloads
    python refresh_refs.py --file kev=C:\\transfer\\kev.json --file attack=C:\\transfer\\attack.json

Only public, fixed files are downloaded. No question, case detail or IOC is ever
sent anywhere. If a download fails or looks truncated, that source's existing
reference file is left untouched and the others still refresh.

It also creates a runbooks\\ folder (once) for your own procedures and case notes.
"""
import argparse
import hashlib
import io
import json
import os
import re
import sys
import zipfile
from datetime import datetime, timezone
from pathlib import Path
from urllib.request import Request, urlopen

ATTACK_URLS = [
    "https://raw.githubusercontent.com/mitre-attack/attack-stix-data/master/enterprise-attack/enterprise-attack.json",
]
KEV_URLS = [
    "https://www.cisa.gov/sites/default/files/feeds/known_exploited_vulnerabilities.json",
    "https://raw.githubusercontent.com/cisagov/kev-data/develop/known_exploited_vulnerabilities.json",
]
SIGMA_URLS = [
    "https://github.com/SigmaHQ/sigma/releases/latest/download/sigma_all_rules.zip",
]
ATOMIC_URLS = [
    "https://raw.githubusercontent.com/redcanaryco/atomic-red-team/master/atomics/Indexes/index.yaml",
]
LOLBAS_URLS = [  # the project's JSON API first, then the GitHub bundle of its YAML files
    "https://lolbas-project.github.io/api/lolbas.json",
    "https://github.com/LOLBAS-Project/LOLBAS/archive/refs/heads/master.zip",
]
LOLDRIVERS_URLS = [
    "https://raw.githubusercontent.com/magicsword-io/LOLDrivers/main/loldrivers.io/content/api/drivers.json",
    "https://www.loldrivers.io/api/drivers.json",
]

SEP = "====="            # record separator that ask_grounded.py understands
MAX_RECORD = 1400        # stays under ask_grounded.py's 1,500-character chunk size
MIN_TECHNIQUES = 400     # sanity floors: a smaller result means a bad download
MIN_KEV = 1000
MIN_SIGMA = 1000
MIN_ATOMIC = 500
MIN_LOLBAS = 100
MIN_LOLDRIVERS = 200
USER_AGENT = "refresh_refs/2.0 (local SOC reference refresh)"


# ---------- helpers ----------

def clean_text(text):
    """Strip citation markers, markdown links and HTML tags; collapse whitespace."""
    text = re.sub(r"\(Citation:[^)]*\)", "", str(text or ""))
    text = re.sub(r"\[([^\]]+)\]\([^)]*\)", r"\1", text)
    text = re.sub(r"</?[a-zA-Z][^>]*>", "", text)
    text = re.sub(r"[\x00-\x08\x0b-\x1f\x7f]", " ", text)
    return re.sub(r"\s+", " ", text).strip()


def trim(text, limit):
    """Cut text to at most `limit` characters, preferring a sentence boundary."""
    if len(text) <= limit:
        return text
    cut = text[:limit]
    end = max(cut.rfind(". "), cut.rfind("; "))
    if end > limit * 0.5:
        return cut[: end + 1]
    return cut.rsplit(" ", 1)[0] + " ..."


def write_atomic(path, text):
    tmp = path.with_suffix(path.suffix + ".tmp")
    tmp.write_text(text, encoding="utf-8")
    os.replace(tmp, path)


def v_json(data):
    json.loads(data)


def v_zip(data):
    if not zipfile.is_zipfile(io.BytesIO(data)):
        raise ValueError("not a zip file")


def v_yaml(data):
    head = data[:200].lstrip().lower()
    if head.startswith(b"<!doctype") or head.startswith(b"<html"):
        raise ValueError("got a web page instead of the data file")
    if len(data) < 10000:
        raise ValueError("file is too small to be the full index")


def v_json_or_zip(data):
    try:
        v_zip(data)
    except ValueError:
        v_json(data)


def download(urls, dest, validate):
    """Try each URL in order; save the first that downloads and validates. Returns (url, sha256)."""
    errors = []
    for url in urls:
        try:
            req = Request(url, headers={"User-Agent": USER_AGENT})
            with urlopen(req, timeout=300) as resp:
                data = resp.read()
            validate(data)  # must look right before it is kept
            dest.write_bytes(data)
            return url, hashlib.sha256(data).hexdigest()
        except Exception as err:  # network, HTTP, or bad content: try the next source
            errors.append(f"  {url}: {err}")
    raise RuntimeError("All download sources failed:\n" + "\n".join(errors))


def load_local(path, dest, validate):
    data = Path(path).read_bytes()
    validate(data)
    dest.write_bytes(data)
    return str(Path(path)), hashlib.sha256(data).hexdigest()


def pack(header, lines, limit=MAX_RECORD):
    """Pack lines into records no longer than `limit`, repeating the header."""
    records, body = [], []
    room = limit - len(header) - 12
    size = 0
    for line in lines:
        line = trim(line, room)
        if body and size + len(line) + 1 > room:
            records.append(body)
            body, size = [], 0
        body.append(line)
        size += len(line) + 1
    if body:
        records.append(body)
    out = []
    for n, part in enumerate(records, 1):
        title = header if len(records) == 1 else header.replace("\n", f" (part {n} of {len(records)})\n", 1)
        out.append(title + "\n" + "\n".join(part))
    return out


def need_yaml():
    try:
        import yaml
    except ImportError:
        raise RuntimeError("PyYAML is not installed. Run:  python -m pip install pyyaml")
    return yaml, getattr(yaml, "CSafeLoader", yaml.SafeLoader)


def short(value, limit=70):
    return trim(clean_text(value), limit)


# ---------- ATT&CK ----------

def attack_external_id(obj):
    for ref in obj.get("external_references", []):
        if ref.get("source_name") == "mitre-attack" and ref.get("external_id"):
            return ref["external_id"], ref.get("url", "")
    return None, ""


def build_attack(bundle):
    objs = bundle["objects"]
    by_id = {o["id"]: o for o in objs}
    version = next((o.get("x_mitre_version") for o in objs if o.get("type") == "x-mitre-collection"), "unknown")

    techniques = {}
    for o in objs:
        if o.get("type") != "attack-pattern" or o.get("revoked") or o.get("x_mitre_deprecated"):
            continue
        tid, url = attack_external_id(o)
        if tid:
            techniques[o["id"]] = (tid, url, o)
    if len(techniques) < MIN_TECHNIQUES:
        raise RuntimeError(f"Only {len(techniques)} techniques found; expected at least {MIN_TECHNIQUES}.")
    names = {tid: o["name"] for tid, _, o in techniques.values()}

    tech_records = []
    for tid, url, o in sorted(techniques.values(), key=lambda t: t[0]):
        parent = ""
        if o.get("x_mitre_is_subtechnique") and "." in tid:
            pid = tid.split(".")[0]
            parent = f" Sub-technique of {pid} {names.get(pid, '')}.".rstrip()
        tactics = ", ".join(p["phase_name"].replace("-", " ").title() for p in o.get("kill_chain_phases", []))
        platforms = ", ".join(o.get("x_mitre_platforms", []))
        header = f"{tid} {o['name']}\nMITRE ATT&CK Enterprise v{version} technique ID {tid}.{parent}"
        lines = []
        if tactics:
            lines.append(f"Tactics: {tactics}.")
        if platforms:
            lines.append(f"Platforms: {platforms}.")
        lines.append(f"URL: {url}")
        lines.append("Description: " + clean_text(o.get("description", "")))
        tech_records.extend(pack(header, lines))

    det_records = []
    for rel in objs:
        if rel.get("type") != "relationship" or rel.get("relationship_type") != "detects":
            continue
        strat, tech = by_id.get(rel["source_ref"]), techniques.get(rel["target_ref"])
        if not strat or not tech or strat.get("type") != "x-mitre-detection-strategy":
            continue
        tid, _, t = tech
        did, _ = attack_external_id(strat)
        header = (f"{tid} {t['name']} - detection guidance ({did})\n"
                  f"MITRE ATT&CK Enterprise v{version} detection strategy: {clean_text(strat['name'])}.")
        lines = []
        for aref in strat.get("x_mitre_analytic_refs", []):
            an = by_id.get(aref)
            if not an:
                continue
            logs = "; ".join(
                f"{s.get('name', '')} {s.get('channel', '')}".strip()
                for s in an.get("x_mitre_log_source_references", [])
            )
            plat = ", ".join(an.get("x_mitre_platforms", []))
            lines.append(f"{an['name']} ({plat}). Log sources: {logs}. {clean_text(an.get('description', ''))}")
        if lines:
            det_records.extend(pack(header, lines))

    return version, tech_records, det_records, len(techniques)


# ---------- CISA KEV ----------

def build_kev(catalog):
    vulns = catalog.get("vulnerabilities", [])
    if len(vulns) < MIN_KEV:
        raise RuntimeError(f"Only {len(vulns)} KEV entries found; expected at least {MIN_KEV}.")
    version = catalog.get("catalogVersion", "unknown")
    released = catalog.get("dateReleased", "unknown")
    records = []
    for v in sorted(vulns, key=lambda x: (x.get("dateAdded", ""), x.get("cveID", "")), reverse=True):
        header = f"{v['cveID']} {v.get('vendorProject', '')} {v.get('product', '')}: {clean_text(v.get('vulnerabilityName', ''))}"
        cwes = ", ".join(v.get("cwes", []) or [])
        lines = [
            f"CISA Known Exploited Vulnerabilities catalog (version {version}). "
            f"Date added: {v.get('dateAdded', '')}. Remediation due: {v.get('dueDate', '')}. "
            f"Known ransomware campaign use: {v.get('knownRansomwareCampaignUse', 'Unknown')}.",
            "Description: " + clean_text(v.get("shortDescription", "")),
            "Required action: " + clean_text(v.get("requiredAction", "")),
        ]
        if cwes:
            lines.append(f"CWE: {cwes}.")
        if v.get("notes"):
            lines.append("Notes: " + clean_text(v["notes"]))
        records.extend(pack(header, lines))  # long entries split into numbered parts
    return version, released, records, len(vulns)


# ---------- Sigma ----------

def attack_tags(tags):
    """Split Sigma tags into ATT&CK technique IDs, tactics, CVEs and other tags."""
    techniques, tactics, cves, other = [], [], [], []
    for tag in tags or []:
        tag = str(tag).lower()
        m = re.fullmatch(r"attack\.(t\d{4}(?:\.\d{3})?)", tag)
        if m:
            techniques.append(m.group(1).upper())
        elif tag.startswith("attack."):
            tactics.append(tag[7:].replace("-", " ").replace("_", " ").title())
        elif tag.startswith("cve."):
            cves.append("CVE-" + tag[4:])
        else:
            other.append(tag)
    return techniques, tactics, cves, other


def render_detection(det, limit=650):
    """Condense a Sigma detection block: which fields are matched against which values."""
    parts = []
    for name, body in det.items():
        if name in ("condition", "timeframe"):
            continue
        items = body if isinstance(body, list) else [body]
        pieces = []
        for item in items:
            if isinstance(item, dict):
                for field, val in item.items():
                    vals = val if isinstance(val, list) else [val]
                    shown = " | ".join(short(v, 60) for v in vals[:8])
                    more = f" (+{len(vals) - 8} more)" if len(vals) > 8 else ""
                    pieces.append(f"{field}: {shown}{more}")
            else:
                pieces.append(short(item, 80))
        parts.append(f"{name} = " + "; ".join(pieces[:12]))
    cond = det.get("condition")
    if cond:
        parts.append("condition: " + (cond if isinstance(cond, str) else " / ".join(map(str, cond))))
    return trim(" || ".join(parts), limit)


def build_sigma(data):
    yaml, loader = need_yaml()
    records, rules, skipped = [], 0, 0
    with zipfile.ZipFile(io.BytesIO(data)) as z:
        for name in sorted(z.namelist()):
            if not name.lower().endswith((".yml", ".yaml")):
                continue
            try:
                docs = list(yaml.load_all(z.read(name).decode("utf-8", errors="replace"), Loader=loader))
            except Exception:
                skipped += 1
                continue
            rule = next((d for d in docs if isinstance(d, dict) and d.get("title") and d.get("detection")), None)
            if not rule or str(rule.get("status", "")).lower() in ("deprecated", "unsupported"):
                skipped += 1
                continue
            rules += 1
            techniques, tactics, cves, other = attack_tags(rule.get("tags"))
            ls = rule.get("logsource") or {}
            logsource = ", ".join(f"{k} {v}" for k, v in ls.items() if v) or "not stated"
            header = (f"{clean_text(rule['title'])} (Sigma rule)\n"
                      f"Sigma detection rule. Status: {rule.get('status', '?')}. Level: {rule.get('level', '?')}. "
                      f"Rule file: {name}.")
            lines = []
            if techniques or tactics:
                lines.append(f"ATT&CK techniques: {', '.join(techniques) or 'none tagged'}. Tactics: {', '.join(tactics) or 'none tagged'}.")
            if cves:
                lines.append("CVEs: " + ", ".join(cves) + ".")
            lines.append(f"Log source: {logsource}.")
            lines.append("Description: " + clean_text(rule.get("description", "")))
            lines.append("Detection: " + render_detection(rule["detection"]))
            fps = rule.get("falsepositives")
            if fps:
                fps = fps if isinstance(fps, list) else [fps]
                lines.append("False positives: " + "; ".join(short(x, 90) for x in fps[:4]) + ".")
            records.extend(pack(header, lines))
    if rules < MIN_SIGMA:
        raise RuntimeError(f"Only {rules} Sigma rules found; expected at least {MIN_SIGMA}.")
    return records, rules, skipped


# ---------- Atomic Red Team ----------

def build_atomic(data):
    yaml, loader = need_yaml()
    index = yaml.load(data.decode("utf-8", errors="replace"), Loader=loader)
    seen, records, tactics_of = set(), [], {}
    for tactic, techs in index.items():
        for tid in techs:
            tactics_of.setdefault(tid, []).append(tactic.replace("-", " ").title())
    for tactic, techs in index.items():
        for tid, body in techs.items():
            tname = clean_text((body.get("technique") or {}).get("name", ""))
            for test in body.get("atomic_tests") or []:
                key = test.get("auto_generated_guid") or (tid, test.get("name"))
                if key in seen:
                    continue
                seen.add(key)
                ex = test.get("executor") or {}
                cmd = clean_text(ex.get("command") or ex.get("steps") or "")
                header = (f"{tid} {tname} - Atomic Red Team test: {clean_text(test.get('name', ''))}\n"
                          f"Atomic Red Team test for ATT&CK {tid} {tname}. Tactics: {', '.join(tactics_of[tid])}.")
                lines = [
                    f"Platforms: {', '.join(test.get('supported_platforms', []))}. Executor: {ex.get('name', '?')}"
                    + (". Needs elevation." if ex.get("elevation_required") else "."),
                    "Description: " + clean_text(test.get("description", "")),
                ]
                if cmd:
                    lines.append("Test command (what this looks like on an endpoint): " + trim(cmd, 450))
                records.extend(pack(header, lines))
    if len(seen) < MIN_ATOMIC:
        raise RuntimeError(f"Only {len(seen)} Atomic tests found; expected at least {MIN_ATOMIC}.")
    return records, len(seen)


# ---------- LOLBAS ----------

def lolbas_entries(data):
    """Read LOLBAS entries from the JSON API file or from the GitHub zip of YAML files."""
    if zipfile.is_zipfile(io.BytesIO(data)):
        yaml, loader = need_yaml()
        entries = []
        with zipfile.ZipFile(io.BytesIO(data)) as z:
            for name in sorted(z.namelist()):
                parts = name.split("/")
                if name.lower().endswith(".yml") and "yml" in parts[:3] and "Archive-Old-Version" not in parts:
                    doc = yaml.load(z.read(name).decode("utf-8", errors="replace"), Loader=loader)
                    if isinstance(doc, dict) and doc.get("Name") and doc.get("Commands"):
                        entries.append(doc)
        return entries
    return json.loads(data)


def build_lolbas(data):
    entries = [e for e in lolbas_entries(data) if isinstance(e, dict) and e.get("Name")]
    if len(entries) < MIN_LOLBAS:
        raise RuntimeError(f"Only {len(entries)} LOLBAS entries found; expected at least {MIN_LOLBAS}.")
    records = []
    for e in sorted(entries, key=lambda x: str(x["Name"]).lower()):
        header = (f"{clean_text(e['Name'])} (LOLBAS)\n"
                  f"LOLBAS entry: a legitimate Windows binary, script or library that attackers abuse (living off the land).")
        lines = ["Description: " + clean_text(e.get("Description", ""))]
        paths = [clean_text(p.get("Path", "")) for p in e.get("Full_Path") or [] if isinstance(p, dict)]
        if paths:
            lines.append("Typical paths: " + "; ".join(paths[:4]) + ".")
        for c in e.get("Commands") or []:
            lines.append(
                f"Abuse: {clean_text(c.get('Usecase', ''))} [{clean_text(c.get('Category', ''))}, "
                f"ATT&CK {clean_text(c.get('MitreID', ''))}, privileges {clean_text(c.get('Privileges', ''))}]. "
                f"Command: {clean_text(c.get('Command', ''))}. {clean_text(c.get('Description', ''))}")
        for d in e.get("Detection") or []:
            if isinstance(d, dict):
                lines.append("Detection: " + "; ".join(f"{k} {clean_text(v)}" for k, v in d.items() if v))
        records.extend(pack(header, lines))
    return records, len(entries)


# ---------- LOLDrivers ----------

def build_loldrivers(drivers):
    if not isinstance(drivers, list) or len(drivers) < MIN_LOLDRIVERS:
        raise RuntimeError(f"Only {len(drivers) if isinstance(drivers, list) else 0} LOLDrivers entries; expected at least {MIN_LOLDRIVERS}.")
    records = []
    for d in sorted(drivers, key=lambda x: str((x.get("Tags") or [""])[0]).lower()):
        tags = [clean_text(t) for t in d.get("Tags") or []]
        cmds = d.get("Commands") or {}
        if isinstance(cmds, list):
            cmds = cmds[0] if cmds else {}
        header = (f"{tags[0] if tags else d.get('Id', 'unknown')} (LOLDrivers: {d.get('Category', 'driver')})\n"
                  f"LOLDrivers entry {d.get('Id', '')}: {d.get('Category', '')}. Names: {', '.join(tags)}. "
                  f"ATT&CK {clean_text(d.get('MitreID', ''))}. Verified: {d.get('Verified', '?')}.")
        lines = []
        if cmds.get("Description"):
            lines.append("Description: " + clean_text(cmds["Description"]))
        if cmds.get("Usecase"):
            lines.append("Use case: " + clean_text(cmds["Usecase"]))
        if cmds.get("Command"):
            lines.append("Example command: " + clean_text(cmds["Command"]))
        samples = d.get("KnownVulnerableSamples") or []
        for s in samples[:8]:
            hashes = "; ".join(f"{k} {s[k]}" for k in ("SHA256", "SHA1", "MD5") if s.get(k))
            who = clean_text(s.get("Company") or s.get("Publisher") or "")
            lines.append(f"Sample {clean_text(s.get('Filename', ''))} ({who}, {clean_text(s.get('Signature', ''))}): {hashes}.")
        if len(samples) > 8:
            lines.append(f"{len(samples) - 8} more known samples are listed in the LOLDrivers project.")
        records.extend(pack(header, lines))
    return records, len(drivers)


# ---------- runbooks folder ----------

RUNBOOK_README = """# Runbooks folder

Put your own procedures, escalation rules and past case notes here as .md or .txt files.
The grounded assistant reads this folder together with refs\\ every time it starts
(or when you press "Reload library" on the web page).

Tips for files the assistant can use well:
- One topic per section. Each heading (#, ##, ###) becomes its own searchable record.
- Put the key words in the heading, for example "## Phishing: user clicked a link".
- Write steps as short sentences with names, thresholds and contacts spelled out.
- Do not store passwords, API keys or customer data here. Anything in this folder
  can be sent to the model as reference text.
- Files whose names start with an underscore (like this one) are skipped.
"""

RUNBOOK_TEMPLATE = """# Runbook: <alert or scenario name>

## Scope
When this applies, and what it does not cover.

## Triage
1. First check, and where to look (tool, query, log source).
2. Second check.

## Escalate when
- Condition that moves this to the next tier, and who to contact.

## Containment
- Approved actions and who must approve them.

## Related
ATT&CK technique IDs, Sigma rule names, and past case numbers.
"""


def scaffold_runbooks(folder):
    folder.mkdir(parents=True, exist_ok=True)
    made = False
    for name, text in (("_README.md", RUNBOOK_README), ("_template.md", RUNBOOK_TEMPLATE)):
        path = folder / name
        if not path.exists():
            path.write_text(text, encoding="utf-8")
            made = True
    return made


# ---------- output ----------

def records_text(source_note, records):
    return "\n".join(f"{SEP}\n{r}" for r in [source_note] + records) + "\n"


def job_attack(data, stamp):
    version, tech, det, count = build_attack(json.loads(data))
    files = {
        "attack_techniques.txt": records_text(
            f"SOURCE NOTE: MITRE ATT&CK Enterprise v{version}, retrieved {stamp}. One record per technique.", tech),
        "attack_detections.txt": records_text(
            f"SOURCE NOTE: MITRE ATT&CK Enterprise v{version} detection guidance, retrieved {stamp}.", det),
    }
    info = {"version": version, "techniques": count, "technique_records": len(tech), "detection_records": len(det)}
    return files, info, f"ATT&CK v{version}: {count} techniques, {len(tech)} technique records, {len(det)} detection records."


def job_kev(data, stamp):
    version, released, recs, count = build_kev(json.loads(data))
    files = {"cisa_kev.txt": records_text(
        f"SOURCE NOTE: CISA Known Exploited Vulnerabilities catalog version {version} "
        f"(released {released}), retrieved {stamp}. One record per CVE.", recs)}
    info = {"catalog_version": version, "released": released, "entries": count}
    return files, info, f"CISA KEV {version}: {count} entries (released {released})."


def job_sigma(data, stamp):
    recs, rules, skipped = build_sigma(data)
    files = {"sigma_rules.txt": records_text(
        f"SOURCE NOTE: SigmaHQ rules (release bundle sigma_all_rules.zip), retrieved {stamp}. Rules are licensed "
        f"under the Detection Rule License 1.1. One record per rule; deprecated rules left out.", recs)}
    return files, {"rules": rules, "records": len(recs)}, f"Sigma: {rules} rules, {len(recs)} records ({skipped} files skipped: deprecated or unreadable)."


def job_atomic(data, stamp):
    recs, tests = build_atomic(data)
    files = {"atomic_tests.txt": records_text(
        f"SOURCE NOTE: Atomic Red Team (Red Canary, MIT license) test index, retrieved {stamp}. "
        f"One record per test. These are simulation commands for detection testing.", recs)}
    return files, {"tests": tests, "records": len(recs)}, f"Atomic Red Team: {tests} tests, {len(recs)} records."


def job_lolbas(data, stamp):
    recs, count = build_lolbas(data)
    files = {"lolbas.txt": records_text(
        f"SOURCE NOTE: LOLBAS project (GPL-3.0), retrieved {stamp}. One record per binary, script or library.", recs)}
    return files, {"entries": count, "records": len(recs)}, f"LOLBAS: {count} entries, {len(recs)} records."


def job_loldrivers(data, stamp):
    recs, count = build_loldrivers(json.loads(data))
    files = {"loldrivers.txt": records_text(
        f"SOURCE NOTE: LOLDrivers project (Apache-2.0), retrieved {stamp}. One record per driver, with sample hashes.", recs)}
    return files, {"entries": count, "records": len(recs)}, f"LOLDrivers: {count} drivers, {len(recs)} records."


# name -> (cache file, urls, validator, job, label)
SOURCES = {
    "attack": ("enterprise-attack.json", ATTACK_URLS, v_json, job_attack, "ATT&CK"),
    "kev": ("known_exploited_vulnerabilities.json", KEV_URLS, v_json, job_kev, "CISA KEV"),
    "sigma": ("sigma_all_rules.zip", SIGMA_URLS, v_zip, job_sigma, "Sigma"),
    "atomic": ("atomic_index.yaml", ATOMIC_URLS, v_yaml, job_atomic, "Atomic Red Team"),
    "lolbas": ("lolbas_data", LOLBAS_URLS, v_json_or_zip, job_lolbas, "LOLBAS"),
    "loldrivers": ("loldrivers.json", LOLDRIVERS_URLS, v_json, job_loldrivers, "LOLDrivers"),
}


def main():
    parser = argparse.ArgumentParser(description="Refresh the local reference library from public bulk data.")
    parser.add_argument("--out", default="refs", help="reference folder (default: refs)")
    parser.add_argument("--runbooks", default="runbooks", help="folder for your own runbooks (default: runbooks)")
    parser.add_argument("--only", nargs="+", choices=sorted(SOURCES), metavar="SOURCE",
                        help="refresh just these sources: " + ", ".join(sorted(SOURCES)))
    parser.add_argument("--offline", action="store_true", help="rebuild from the files already in <out>\\_downloads")
    parser.add_argument("--file", action="append", default=[], metavar="SOURCE=PATH",
                        help="use a local copy of a source's data file instead of downloading (repeatable)")
    parser.add_argument("--attack-file", help="same as --file attack=PATH")
    parser.add_argument("--kev-file", help="same as --file kev=PATH")
    args = parser.parse_args()

    local = {}
    for item in args.file:
        name, _, path = item.partition("=")
        if name not in SOURCES or not path:
            parser.error(f"--file expects SOURCE=PATH with SOURCE one of {', '.join(sorted(SOURCES))}")
        local[name] = path
    if args.attack_file:
        local["attack"] = args.attack_file
    if args.kev_file:
        local["kev"] = args.kev_file

    out = Path(args.out)
    cache = out / "_downloads"
    cache.mkdir(parents=True, exist_ok=True)
    stamp = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC")
    info_path = out / "refresh_info.json"
    info = json.loads(info_path.read_text(encoding="utf-8")) if info_path.exists() else {}
    failures = []

    for key in (args.only or list(SOURCES)):
        cache_name, urls, validate, job, label = SOURCES[key]
        dest = cache / cache_name
        try:
            if args.offline:
                if not dest.exists():
                    raise RuntimeError(f"--offline needs {dest}, which does not exist.")
                data = dest.read_bytes()
                source = info.get(key, {}).get("source", "cached file")
                digest = hashlib.sha256(data).hexdigest()
            else:
                if key in local:
                    source, digest = load_local(local[key], dest, validate)
                else:
                    print(f"Downloading {label} ...", flush=True)
                    source, digest = download(urls, dest, validate)
                data = dest.read_bytes()
            files, details, message = job(data, stamp)
            for fname, text in files.items():
                write_atomic(out / fname, text)
            info[key] = {"retrieved": stamp, "source": source, "sha256": digest, **details}
            print(message)
        except Exception as err:
            failures.append(f"{label}: {err}")

    if info:
        write_atomic(info_path, json.dumps(info, indent=2) + "\n")
    if scaffold_runbooks(Path(args.runbooks)):
        print(f"\nCreated the {args.runbooks}\\ folder for your own runbooks (see _README.md inside).")
    if failures:
        print("\nNot refreshed (existing files, if any, were left as they were):")
        for f in failures:
            print("  " + f)
        sys.exit(1)
    print(f"\nReference files are in {out}. Example:")
    print('  python ask_grounded.py -r "refs\\*.txt" -r "runbooks\\*.md" -q "What is the MITRE ATT&CK ID for Kerberoasting?"')


if __name__ == "__main__":
    main()