path: root/fw/ta_corpus_memscan.py

#!/usr/bin/env python3
"""Sweep every harvested TA body for the three memory-corruption shapes, boundary-free.

The twelve analysed builds are clean: every host-derived length and index is bounded, often
exactly. But the corpus holds ~110 bodies across 17 component types, and a *different build*
may simply lack a bound that the analysed one has. The HDCP version sweep showed that
build-to-build comparison is where this project gets its strongest results.

The existing detectors need IDA function boundaries, which exist only for the twelve. This
one derives its own scope: for each candidate site, the enclosing "function" is approximated
by walking back to the nearest Thumb prologue (`push {...,lr}`), which is enough for a
backward slice and for asking "is this value compared anywhere before use".

Shapes (same three, same meanings as fw/AUDIT-TA-MEMSAFETY-HUNT.md):
  L  computed length into an ARM-mode copy helper, unbounded in scope
  D  copy destination loaded from memory
  X  `base + K*index` where the index came from a load and is never compared

CONTROL: HDCP's confirmed stack overflow must be found in boundary-free mode too. If the
scope approximation breaks the slice, the control fails and nothing here is trustworthy.

Usage:
  ./ta_corpus_memscan.py --selftest
  ./ta_corpus_memscan.py --all [--component gcp]
"""
from __future__ import annotations

import argparse
import json
import re
from collections import defaultdict
from pathlib import Path

import capstone as cs

HERE = Path(__file__).resolve().parent
REPO = HERE.parent
BODIES = REPO / "psp_out" / "versions_win"
THUMB = cs.Cs(cs.CS_ARCH_ARM, cs.CS_MODE_THUMB)
ARM = cs.Cs(cs.CS_ARCH_ARM, cs.CS_MODE_ARM)
BASE = 0x400100


def linear(b: bytes, base: int = BASE):
    """Resynchronising linear Thumb sweep of the whole body.

    capstone's disasm() STOPS at the first undecodable byte, and these images are full of
    literal pools and ARM-mode regions. A single call therefore decoded only the first few
    hundred bytes, which is why the boundary-free control failed. Skip two bytes and resume.
    """
    ins = []
    off = 0
    n = len(b)
    while off < n - 1:
        got = 0
        for i in THUMB.disasm(b[off:], base + off):
            ins.append(i)
            got += i.size
            if got > 0x4000:      # bound the chunk so a bad region cannot stall us
                break
        off += got if got else 2
    return ins, {i.address: k for k, i in enumerate(ins)}


def scope_start(ins: list, at: int) -> int:
    """Nearest preceding prologue, i.e. an approximate function start."""
    for k in range(at, max(0, at - 400) - 1, -1):
        if ins[k].mnemonic.startswith("push"):
            return k
    return max(0, at - 400)


def copy_helpers(b: bytes, ins: list, base: int) -> set[int]:
    out = set()
    for i in ins:
        if i.mnemonic == "blx" and (i.op_str or "").startswith("#"):
            t = int(i.op_str[1:], 16)
            off = t - base
            if not (0 <= off < len(b) - 16):
                continue
            a = list(ARM.disasm(b[off:off + 0x60], t))
            if any(x.mnemonic.startswith(("cdp", "mcr", "mrc", "stc", "ldc")) for x in a):
                continue
            if any(x.mnemonic.startswith(("ldr", "ldm")) for x in a) and \
               any(x.mnemonic.startswith(("str", "stm")) for x in a):
                out.add(t)
    return out


def bounded(ins: list, lo: int, at: int, reg: str) -> bool:
    chain = {reg}
    for k in range(at - 1, lo - 1, -1):
        i = ins[k]
        ops = i.op_str or ""
        if i.mnemonic.startswith(("cmp", "cmn", "tst", "subs", "cbz", "cbnz")):
            if any(re.search(r"\b%s\b" % re.escape(r), ops) for r in chain):
                return True
        if i.mnemonic.startswith(("and", "ubfx", "uxtb", "uxth", "lsr", "bic")):
            m = re.match(r"(\w+)", ops)
            if m and m.group(1) in chain:
                return True
        m = re.match(r"(\w+),\s*(\w+)$", ops)
        if m and m.group(1) in chain and i.mnemonic.startswith("mov"):
            chain.add(m.group(2))
    return False


def origin(ins: list, lo: int, at: int, reg: str) -> str | None:
    cur, pos = reg, at
    for _ in range(4):
        for k in range(pos - 1, lo - 1, -1):
            i = ins[k]
            m = re.match(r"(%s),\s*(.+)$" % re.escape(cur), i.op_str or "")
            if not m or i.mnemonic.startswith(("cmp", "str", "push", "tst", "b")):
                continue
            rhs = m.group(2)
            if i.mnemonic.startswith("ldr"):
                return "LOAD %s" % rhs
            if re.fullmatch(r"#(0x[0-9a-f]+|\d+)", rhs.strip()):
                return None                      # immediate: fine
            m2 = re.match(r"(r\d+|sb|sl|fp|ip)$", rhs.strip())
            if m2:
                cur, pos = m2.group(1), k
                break
            # Arithmetic: follow the FIRST source register. HDCP's length is
            # `add.w r7, r8, r8, lsl #2` (the times-five), so giving up here loses the
            # very bug this scan is controlled against.
            if i.mnemonic.startswith(("add", "sub", "mul", "lsl", "lsr", "orr", "mla",
                                      "rsb", "asr")):
                m3 = re.match(r"(r\d+|sb|sl|fp|ip)\b", rhs.strip())
                if m3:
                    cur, pos = m3.group(1), k
                    break
            return None
        else:
            return None
    return None


def scan_body(b: bytes, base: int = BASE) -> list[tuple[str, int, str]]:
    ins, _ = linear(b, base)
    helpers = copy_helpers(b, ins, base)
    hits = []
    for n, i in enumerate(ins):
        lo = scope_start(ins, n)
        if i.mnemonic in ("blx", "bl") and (i.op_str or "").startswith("#") \
           and int(i.op_str[1:], 16) in helpers:
            # Do NOT gate on bounded(): r2 is a scratch register, so "compared somewhere
            # in scope" matches almost always and suppressed HDCP's real overflow. Report
            # every memory-derived length and annotate; the value of this sweep is the
            # DIFF between builds of the same component, not a per-site verdict.
            ln = origin(ins, lo, n, "r2")
            if ln:
                note = "" if not bounded(ins, lo, n, "r2") else " (a compare exists in scope)"
                hits.append(("L", i.address, "len <- %s%s" % (ln, note)))
            ds = origin(ins, lo, n, "r0")
            if ds and ds.startswith("LOAD"):
                hits.append(("D", i.address, "dst <- %s" % ds))
        idx = None
        if i.mnemonic == "mla":
            m = re.match(r"(\w+),\s*(\w+),\s*(\w+),\s*(\w+)$", i.op_str or "")
            if m:
                idx = m.group(2)
        elif i.mnemonic.startswith("add"):
            m = re.match(r"(\w+),\s*(\w+),\s*(\w+),\s*lsl\s*#(\d+)$", i.op_str or "")
            if m:
                idx = m.group(3)
        if idx:
            src = origin(ins, lo, n, idx)
            if src and src.startswith("LOAD"):
                note = "" if not bounded(ins, lo, n, idx) else " (a compare exists in scope)"
                hits.append(("X", i.address, "idx %s <- %s%s" % (idx, src, note)))
    return hits


def component_of(rec: dict) -> str:
    d = str(rec.get("description") or "")
    m = re.match(r"AMD[- ]?(?:TEE )?(\w+) (?:Application|Driver)", d)
    return m.group(1).lower() if m else "?"


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--selftest", action="store_true")
    ap.add_argument("--all", action="store_true")
    ap.add_argument("--component")
    args = ap.parse_args()

    if args.selftest:
        b = (REPO / "psp_out" / "hdcp_code.bin").read_bytes()
        hits = scan_body(b)
        want = [h for h in hits if h[0] == "L" and h[1] == 0x402CDE]
        print("SELFTEST (boundary-free) on HDCP's confirmed overflow @0x402cde:")
        print("   %s" % (("FOUND: %s" % (want[0],)) if want
                         else "MISSED (L-hits: %s)" % [hex(h[1]) for h in hits if h[0] == "L"][:6]))
        print("\nverdict: %s" % ("TRUSTWORTHY" if want else "NOT trustworthy"))
        return 0 if want else 1

    # map body files -> component via the harvest inventory
    comp = {}
    for line in (HERE / "windows_inventory.jsonl").read_text().splitlines():
        try:
            r = json.loads(line)
        except Exception:  # noqa: BLE001
            continue
        h = (r.get("sha256") or "")[:8]
        if h:
            comp[h] = (component_of(r), r.get("ta_version") or r.get("version"))

    groups = defaultdict(list)
    for p in sorted(BODIES.glob("ta_*.bin")):
        m = re.match(r"ta_(.+)_([0-9a-f]{8})\.bin$", p.name)
        if not m:
            continue
        c, v = comp.get(m.group(2), ("?", m.group(1)))
        if args.component and c != args.component:
            continue
        groups[c].append((m.group(1), p))

    print("%-10s %-16s %s" % ("component", "build", "shapes found (L=length D=dest X=index)"))
    for c in sorted(groups):
        for ver, p in sorted(groups[c]):
            hits = scan_body(p.read_bytes())
            kinds = defaultdict(int)
            for k, _, _ in hits:
                kinds[k] += 1
            tag = " ".join("%s=%d" % (k, kinds[k]) for k in sorted(kinds)) or "clean"
            print("%-10s %-16s %s" % (c, ver, tag))
            for k, ea, what in hits:
                if k in ("L", "X"):
                    print("      %s @%#x  %s" % (k, ea, what[:70]))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())