src/cti_attack/attack_meta.py

"""MITRE ATT&CK technique metadata, used for readable labels and the keyword baseline.

Source: https://github.com/mitre-attack/attack-stix-data (ATT&CK Terms of Use —
free to use with attribution).
"""

from __future__ import annotations

import json
import urllib.request

from . import config


def _download() -> None:
    if config.ATTACK_STIX_RAW.exists():
        return
    print(f"downloading ATT&CK STIX bundle (~35 MB) …")
    urllib.request.urlretrieve(config.ATTACK_STIX_URL, config.ATTACK_STIX_RAW)


def _parse_bundle() -> tuple[dict[str, str], dict[str, str]]:
    """Return ``(names, status)`` for every technique ID in the STIX bundle.

    Revoked and deprecated techniques are **kept**, not filtered out. The TRAM
    corpus was annotated against an older ATT&CK release, and two of its labels
    (``T1562.001``, ``T1574.002``) have since been revoked by MITRE. Dropping
    them here would silently blank their names in every report and hide a real
    provenance issue; instead they are named and flagged.
    """
    _download()
    with open(config.ATTACK_STIX_RAW, encoding="utf-8") as fh:
        bundle = json.load(fh)

    names: dict[str, str] = {}
    status: dict[str, str] = {}
    for obj in bundle.get("objects", []):
        if obj.get("type") != "attack-pattern":
            continue
        ext = next(
            (r for r in obj.get("external_references", [])
             if r.get("source_name") == "mitre-attack"),
            None,
        )
        if not (ext and ext.get("external_id")):
            continue
        tid = ext["external_id"]
        if obj.get("revoked"):
            state = "revoked"
        elif obj.get("x_mitre_deprecated"):
            state = "deprecated"
        else:
            state = "current"
        # a current definition always wins over a revoked one sharing the ID
        if tid in status and status[tid] == "current" and state != "current":
            continue
        names[tid] = obj.get("name", "")
        status[tid] = state

    qualified = {}
    for tid, name in names.items():
        if "." in tid:
            parent = names.get(tid.split(".")[0])
            qualified[tid] = f"{parent}: {name}" if parent else name
        else:
            qualified[tid] = name
    return qualified, status


def build_technique_names() -> dict[str, str]:
    """Map technique ID (``T1027``, ``T1059.003``) -> human-readable name.

    Sub-technique names are qualified with their parent, so ``T1059.003``
    becomes ``Command and Scripting Interpreter: Windows Command Shell``.
    """
    if config.ATTACK_NAMES_JSON.exists():
        return json.loads(config.ATTACK_NAMES_JSON.read_text(encoding="utf-8"))

    qualified, status = _parse_bundle()
    config.ATTACK_NAMES_JSON.write_text(
        json.dumps(qualified, indent=2, ensure_ascii=False), encoding="utf-8")
    config.ATTACK_STATUS_JSON.write_text(
        json.dumps(status, indent=2), encoding="utf-8")
    return qualified


def build_technique_status() -> dict[str, str]:
    """Map technique ID -> ``current`` | ``revoked`` | ``deprecated``."""
    if config.ATTACK_STATUS_JSON.exists():
        return json.loads(config.ATTACK_STATUS_JSON.read_text(encoding="utf-8"))
    _, status = _parse_bundle()
    config.ATTACK_STATUS_JSON.write_text(
        json.dumps(status, indent=2), encoding="utf-8")
    return status


def technique_keywords(technique_ids: list[str]) -> dict[str, list[str]]:
    """Surface forms to search for in text, per technique — the keyword baseline.

    Derived from the technique's own name plus its parent's, lowercased. This is
    deliberately naive: it is the zero-training floor a learned model must clear
    to justify existing.
    """
    names = build_technique_names()
    out: dict[str, list[str]] = {}
    for tid in technique_ids:
        forms = set()
        full = names.get(tid, "")
        for part in full.split(":"):
            part = part.strip().lower()
            if len(part) >= 4:
                forms.add(part)
        out[tid] = sorted(forms)
    return out