"""ransomware-kill-chain-mapped-to-controls.py — map a ransomware intrusion chain to MITRE ATT&CK
mitigations, offline, and report which stages your current controls leave uncovered.

What it does: reads the MITRE ATT&CK Enterprise STIX 2.1 bundle (published by MITRE in the
mitre-attack/attack-stix-data repository), looks up each technique in a staged chain, follows
"revoked-by" links so an out-of-date technique ID resolves to its current replacement, and writes one
CSV row per (stage, technique, mitigation). Techniques with no ATT&CK mitigation still get a row, with
the mitigation columns empty, because "nothing to configure, only detect" is itself a finding.
--summary prints the mitigations that touch the most stages. --have takes the mitigation IDs you
already have in place and prints, stage by stage, the techniques none of them address.

It reads only ATT&CK's public metadata (technique names, versions, and mitigation relationships).
It makes no network connections and touches no system other than the files you name.

Inputs:  --bundle  path to an Enterprise ATT&CK STIX bundle, e.g. enterprise-attack-19.2.json
         --chain   optional CSV with columns stage_order,stage,technique_id (default: the chain below)
         --out     optional CSV path for the long-form mapping
         --summary print mitigation coverage by stage count
         --have    comma-separated mitigation IDs you already run, e.g. M1032,M1053,M1030

Run:     curl -LO https://raw.githubusercontent.com/mitre-attack/attack-stix-data/master/enterprise-attack/enterprise-attack-19.2.json
         python code/ransomware-kill-chain-mapped-to-controls.py --bundle enterprise-attack-19.2.json \
             --out datasets/ransomware-kill-chain-mapped-to-controls.csv --summary
         python code/ransomware-kill-chain-mapped-to-controls.py --bundle enterprise-attack-19.2.json \
             --have M1032,M1053,M1051
Standard library only (Python 3.9+).
"""
from __future__ import annotations

import argparse
import csv
import datetime
import json
import sys
from collections import defaultdict

# The default chain: the stages the CISA/MS-ISAC/NSA/FBI #StopRansomware Guide (September 2023) organises
# its prevention and threat-hunting advice around, expressed as ATT&CK Enterprise technique IDs.
# T1562.001 is deliberately listed under its pre-v19 ID to show the revoked-by resolution.
DEFAULT_CHAIN = [
    (1, "Initial access", "T1133"),
    (1, "Initial access", "T1190"),
    (1, "Initial access", "T1078"),
    (1, "Initial access", "T1566"),
    (1, "Initial access", "T1199"),
    (2, "Execution and remote control", "T1059.001"),
    (2, "Execution and remote control", "T1219"),
    (3, "Credential access", "T1003.001"),
    (3, "Credential access", "T1110"),
    (4, "Discovery", "T1482"),
    (4, "Discovery", "T1018"),
    (5, "Lateral movement", "T1021.001"),
    (5, "Lateral movement", "T1021.002"),
    (6, "Defense impairment", "T1562.001"),
    (7, "Exfiltration", "T1567.002"),
    (7, "Exfiltration", "T1048"),
    (8, "Impact", "T1490"),
    (8, "Impact", "T1486"),
]

FIELDS = ["stage_order", "stage", "technique_id", "requested_id", "technique_name", "tactics",
          "technique_version", "technique_modified", "mitigation_id", "mitigation_name",
          "attack_release", "generated_at"]


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


def is_live(obj):
    return not obj.get("revoked") and not obj.get("x_mitre_deprecated")


def load(bundle_path):
    with open(bundle_path, encoding="utf-8") as fh:
        objects = json.load(fh)["objects"]
    by_stix = {o["id"]: o for o in objects}
    release = next((o.get("x_mitre_version") for o in objects if o["type"] == "x-mitre-collection"), "unknown")
    techniques = {}          # ATT&CK ID -> object (live or revoked; revoked resolved below)
    for o in objects:
        if o["type"] == "attack-pattern" and not o.get("x_mitre_deprecated"):
            tid = attack_id(o)
            if tid and (tid not in techniques or is_live(o)):
                techniques[tid] = o
    revoked_by, mitigations = {}, defaultdict(list)
    for o in objects:
        if o["type"] != "relationship" or not is_live(o):
            continue
        if o["relationship_type"] == "revoked-by":
            revoked_by[o["source_ref"]] = o["target_ref"]
        elif o["relationship_type"] == "mitigates":
            src = by_stix.get(o["source_ref"])
            if src and src["type"] == "course-of-action" and is_live(src):
                mitigations[o["target_ref"]].append(src)
    return by_stix, techniques, revoked_by, mitigations, release


def resolve(tid, by_stix, techniques, revoked_by):
    obj = techniques.get(tid)
    if obj is None:
        return None
    seen = set()
    while obj.get("revoked") and obj["id"] in revoked_by and obj["id"] not in seen:
        seen.add(obj["id"])
        obj = by_stix[revoked_by[obj["id"]]]
    return obj if is_live(obj) else None


def read_chain(path):
    if not path:
        return DEFAULT_CHAIN
    with open(path, newline="", encoding="utf-8") as fh:
        return [(int(r["stage_order"]), r["stage"], r["technique_id"].strip()) for r in csv.DictReader(fh)]


def build_rows(chain, by_stix, techniques, revoked_by, mitigations, release):
    today = datetime.date.today().isoformat()
    rows = []
    for order, stage, requested in chain:
        tech = resolve(requested, by_stix, techniques, revoked_by)
        if tech is None:
            print(f"warning: {requested} not found in this bundle (or deprecated with no replacement)", file=sys.stderr)
            continue
        tid = attack_id(tech)
        if tid != requested:
            print(f"note: {requested} is revoked in ATT&CK {release}; using {tid} {tech['name']}", file=sys.stderr)
        base = {"stage_order": order, "stage": stage, "technique_id": tid, "requested_id": requested,
                "technique_name": tech["name"],
                "tactics": ";".join(p["phase_name"] for p in tech.get("kill_chain_phases", [])),
                "technique_version": tech.get("x_mitre_version", ""),
                "technique_modified": tech.get("modified", "")[:10],
                "attack_release": release, "generated_at": today}
        mits = sorted(mitigations.get(tech["id"], []), key=attack_id)
        if not mits:
            rows.append({**base, "mitigation_id": "", "mitigation_name": ""})
        for m in mits:
            rows.append({**base, "mitigation_id": attack_id(m), "mitigation_name": m["name"]})
    return rows


def summary(rows):
    stages = sorted({(r["stage_order"], r["stage"]) for r in rows})
    cover = defaultdict(set)
    techs = defaultdict(set)
    names = {}
    for r in rows:
        if r["mitigation_id"]:
            cover[r["mitigation_id"]].add(r["stage_order"])
            techs[r["mitigation_id"]].add(r["technique_id"])
            names[r["mitigation_id"]] = r["mitigation_name"]
    all_techs = {r["technique_id"] for r in rows}
    print(f"\n{len(all_techs)} techniques across {len(stages)} stages; {len(cover)} distinct mitigations\n")
    print(f"{'mitigation':<48} {'stages':>6} {'techniques':>10}")
    for mid in sorted(cover, key=lambda m: (-len(cover[m]), -len(techs[m]), m)):
        print(f"{mid + ' ' + names[mid]:<48} {len(cover[mid]):>6} {len(techs[mid]):>10}")
    bare = sorted({r["technique_id"] + " " + r["technique_name"] for r in rows if not r["mitigation_id"]})
    if bare:
        print("\nno ATT&CK mitigation (detection only): " + ", ".join(bare))


def gaps(rows, have):
    have = {h.strip().upper() for h in have.split(",") if h.strip()}
    by_tech = defaultdict(set)
    meta = {}
    for r in rows:
        by_tech[r["technique_id"]].add(r["mitigation_id"])
        meta[r["technique_id"]] = (r["stage_order"], r["stage"], r["technique_name"])
    print(f"\ncontrols in place: {', '.join(sorted(have)) or 'none'}\n")
    for tid in sorted(by_tech, key=lambda t: (meta[t][0], t)):
        order, stage, name = meta[tid]
        hit = sorted(by_tech[tid] & have)
        status = "covered by " + ", ".join(hit) if hit else "GAP"
        print(f"{order}. {stage:<30} {tid:<10} {name:<40} {status}")


def chart(rows, path, min_stages=3):
    """Horizontal bar chart of mitigations by stages touched. Needs matplotlib (tested with 3.10.9)."""
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    n_stages = len({r["stage_order"] for r in rows})
    n_techs = len({r["technique_id"] for r in rows})
    release = rows[0]["attack_release"] if rows else "?"
    cover, techs, names = defaultdict(set), defaultdict(set), {}
    for r in rows:
        if r["mitigation_id"]:
            cover[r["mitigation_id"]].add(r["stage_order"])
            techs[r["mitigation_id"]].add(r["technique_id"])
            names[r["mitigation_id"]] = r["mitigation_name"]
    picked = [m for m in cover if len(cover[m]) >= min_stages]
    picked.sort(key=lambda m: (-len(cover[m]), -len(techs[m]), m))
    labels = [f"{m} {names[m]}" for m in picked]
    values = [len(cover[m]) for m in picked]

    surface, ink, ink2, muted, grid, axis, bar = "#fcfcfb", "#0b0b0b", "#52514e", "#898781", "#e1e0d9", "#c3c2b7", "#2a78d6"
    fig, ax = plt.subplots(figsize=(9, 0.42 * len(picked) + 1.6), dpi=110)
    fig.patch.set_facecolor(surface)
    ax.set_facecolor(surface)
    ax.barh(range(len(picked)), values, height=0.55, color=bar, zorder=3)
    ax.set_yticks(range(len(picked)))
    ax.set_yticklabels(labels, color=ink, fontsize=9.5)
    ax.invert_yaxis()
    ax.set_xlim(0, n_stages + 1.9)
    ax.set_xticks(range(0, n_stages + 1))
    ax.tick_params(axis="x", colors=muted, labelsize=9, length=0)
    ax.tick_params(axis="y", length=0)
    ax.grid(axis="x", color=grid, linewidth=0.8, zorder=0)
    for side in ("top", "right", "bottom"):
        ax.spines[side].set_visible(False)
    ax.spines["left"].set_color(axis)
    for i, m in enumerate(picked):
        ax.text(values[i] + 0.08, i, f"{values[i]} stages, {len(techs[m])} techniques",
                va="center", ha="left", fontsize=8.5, color=ink2)
    ax.set_xlabel(f"ransomware-chain stages touched (of {n_stages})", color=muted, fontsize=9)
    fig.text(0.01, 0.985, "ATT&CK mitigations that reach the most stages of a ransomware intrusion",
             ha="left", va="top", fontsize=12, color=ink)
    fig.text(0.01, 0.935, f"Enterprise ATT&CK v{release}: {n_techs} techniques in {n_stages} stages; "
             f"mitigations touching {min_stages} or more stages", ha="left", va="top", fontsize=9, color=ink2)
    fig.tight_layout(rect=(0, 0, 1, 0.9))
    fig.savefig(path, facecolor=surface)
    print(f"wrote chart to {path}", file=sys.stderr)


def main():
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("--bundle", required=True)
    ap.add_argument("--chain")
    ap.add_argument("--out")
    ap.add_argument("--summary", action="store_true")
    ap.add_argument("--have")
    ap.add_argument("--chart", help="optional PNG path; requires matplotlib")
    args = ap.parse_args()
    loaded = load(args.bundle)
    rows = build_rows(read_chain(args.chain), *loaded)
    if args.chart:
        chart(rows, args.chart)
    if args.out:
        with open(args.out, "w", newline="", encoding="utf-8") as fh:
            writer = csv.DictWriter(fh, fieldnames=FIELDS)
            writer.writeheader()
            writer.writerows(rows)
        print(f"wrote {len(rows)} rows to {args.out}", file=sys.stderr)
    if args.summary:
        summary(rows)
    if args.have is not None:
        gaps(rows, args.have)


if __name__ == "__main__":
    main()