#!/usr/bin/env python3
"""Regenerate the ATT&CK revoked-technique remap block embedded in WELA.ps1.

Sigma/hayabusa rule tags keep referring to technique IDs that MITRE has since
revoked (e.g. T1562.001 -> T1685 in ATT&CK v19).  ATT&CK Navigator silently
drops layer entries whose techniqueID is revoked, so WELA rewrites those IDs
before emitting the heatmap.  This script keeps that lookup table in sync with
the latest Enterprise ATT&CK release.

Usage:
    python3 tools/update_attack_remap.py WELA.ps1          # fetch latest, rewrite in place
    python3 tools/update_attack_remap.py                   # print the block only
    python3 tools/update_attack_remap.py --bundle x.json --attack-version 19.2 WELA.ps1
"""
import argparse
import json
import re
import sys
import urllib.request

INDEX_URL = "https://raw.githubusercontent.com/mitre-attack/attack-stix-data/master/index.json"
BEGIN = "# BEGIN ATTACK-REMAP (auto-generated by tools/update_attack_remap.py - do not edit by hand)"
END = "# END ATTACK-REMAP"


def fetch_json(url):
    with urllib.request.urlopen(url) as response:
        return json.load(response)


def latest_enterprise():
    index = fetch_json(INDEX_URL)
    collection = next(c for c in index["collections"] if c["name"] == "Enterprise ATT&CK")
    newest = max(collection["versions"], key=lambda v: [int(x) for x in v["version"].split(".")])
    return newest["version"], fetch_json(newest["url"])


def build_remap(bundle):
    attack_ids = {}
    revoked_by = {}
    for obj in bundle["objects"]:
        if obj.get("type") == "attack-pattern":
            for ref in obj.get("external_references", []):
                if ref.get("source_name") == "mitre-attack":
                    attack_ids[obj["id"]] = ref["external_id"]
        elif obj.get("type") == "relationship" and obj.get("relationship_type") == "revoked-by":
            revoked_by[obj["source_ref"]] = obj["target_ref"]

    direct = {
        attack_ids[src]: attack_ids[dst]
        for src, dst in revoked_by.items()
        if src in attack_ids and dst in attack_ids
    }

    # Collapse revocation chains (e.g. T1150 -> T1547.011) so WELA needs a single lookup.
    def resolve(technique_id):
        seen = set()
        while technique_id in direct and technique_id not in seen:
            seen.add(technique_id)
            technique_id = direct[technique_id]
        return technique_id

    return {k: resolve(k) for k in direct if resolve(k) != k}


def render(attack_version, remap):
    width = max(len(k) for k in remap) + 2
    lines = [
        BEGIN,
        "# Source: MITRE ATT&CK Enterprise v%s (revoked-by relationships, chains collapsed)" % attack_version,
        '$script:AttackVersion = "%s"' % attack_version.split(".")[0],
        "$script:AttackTechniqueRemap = @{",
    ]
    for old in sorted(remap):
        lines.append('    %s = "%s"' % (('"%s"' % old).ljust(width), remap[old]))
    lines += ["}", END]
    return "\n".join(lines)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("target", nargs="?", help="path to WELA.ps1 (omit to print the block)")
    parser.add_argument("--bundle", help="local enterprise-attack STIX bundle instead of downloading")
    parser.add_argument("--attack-version", help="version label to use with --bundle")
    args = parser.parse_args()

    if args.bundle:
        if not args.attack_version:
            parser.error("--attack-version is required with --bundle")
        version = args.attack_version
        with open(args.bundle, encoding="utf-8") as handle:
            bundle = json.load(handle)
    else:
        version, bundle = latest_enterprise()

    remap = build_remap(bundle)
    block = render(version, remap)

    if not args.target:
        print(block)
        return

    source = open(args.target, encoding="utf-8-sig").read()
    pattern = re.compile(re.escape(BEGIN) + ".*?" + re.escape(END), re.S)
    if not pattern.search(source):
        sys.exit("marker block not found in %s" % args.target)
    open(args.target, "w", encoding="utf-8-sig").write(pattern.sub(lambda _: block, source))
    print("updated %s: ATT&CK v%s, %d mappings" % (args.target, version, len(remap)))


if __name__ == "__main__":
    main()