mirror of
https://github.com/Yamato-Security/WELA.git
synced 2026-09-29 02:24:53 +02:00
113 lines
4.2 KiB
Python
113 lines
4.2 KiB
Python
#!/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()
|