#!/usr/bin/env python3
"""Detect the callee-written in/out length slot pattern, and whether it is re-validated.

This is the one bug class in this hunt with a proven instance, and the class field-flow v2 is
blind to by construction (the length lives in a stack slot written by a callee through a
pointer, not in a struct field written by a `str [base,#imm]`).

The pattern, from `FINDING-XGMI-CALLEE-LENGTH-STACK-COPY.md`:

    mov.w sl, #0x40           ; capacity
    str.w sl, [sp, #8]        ; in/out length slot = 64
    add   r1, sp, #8          ; &slot passed to the callee
    bl    <callee>            ; callee writes the REAL length into the slot
    ldr   r4, [sp, #8]        ; re-load it
    mov   r2, r4              ; use as a copy length ...
    blx   <memcpy>            ; ... with no re-validation   <-- BUG

and the safe version, ASD `sub_D12DEC`:

    ldr   r5, [sp, #8]        ; callee-written length
    ldr   r0, [sp, #0x74]     ; the caller's capacity
    cmp   r5, r0 ; bls ok     ; <-- validated, then copy

So the detector looks for: a stack slot initialised from an immediate, its address taken, a
call, a re-load of the same slot, and that value reaching a copy length -- then asks whether
any `cmp` involving the re-loaded register occurs between the re-load and the copy.

CONTROLS (both required):
  POSITIVE  xgmi `0x402fc4` reported   (no re-validation)
  NEGATIVE  asd  `0xd12e4c` NOT reported (compared against the caller's capacity)

Usage:
  ./ta_slotflow_audit.py --selftest
  ./ta_slotflow_audit.py --module <ta prefix>|asd|sys_drv|hdcp
"""
from __future__ import annotations

import argparse
import json
import re
from pathlib import Path

import capstone as cs

HERE = Path(__file__).resolve().parent
REPO = HERE.parent
THUMB = cs.Cs(cs.CS_ARCH_ARM, cs.CS_MODE_THUMB)
ARM = cs.Cs(cs.CS_ARCH_ARM, cs.CS_MODE_ARM)
CACHE = HERE / "audits" / "fnbounds.json"
MODULES = {"hdcp": ("psp_out/hdcp_code.bin", 0x400100),
           "sys_drv": ("psp_out/sys_drv_code.bin", 0xE00100),
           "asd": ("psp_out/asd_code.bin", 0xD00100)}

SLOT_STORE = re.compile(r"^(\w+),\s*\[sp,\s*#(0x[0-9a-f]+|\d+)\]$")
SLOT_LOAD = re.compile(r"^(\w+),\s*\[sp,\s*#(0x[0-9a-f]+|\d+)\]$")
ADDR_OF = re.compile(r"^(\w+),\s*sp,\s*#(0x[0-9a-f]+|\d+)$")


def load_module(name):
    if name in MODULES:
        bpath, base = MODULES[name]
        b = (REPO / bpath).read_bytes()
        fns = json.loads(CACHE.read_text())[name]
    else:
        base = 0x400100
        meta = json.loads((HERE / "syms" / ("%s.meta.json" % name)).read_text())
        bp = Path(meta["path"])
        b = (bp if bp.is_absolute() else REPO / bp).read_bytes()
        fns = [[f["ea"], f.get("size", 0)]
               for f in json.loads((HERE / "libdumps" / ("%s.json" % name)).read_text())
               if f.get("size", 0) > 0]
    return b, base, fns


def split_merged(b, base, fns):
    """Split IDA-merged extents at interior prologues.

    Both directions matter here: a merged extent can pair a slot init in one function with a
    call in another (false positive), and can let a `cmp` from a neighbouring function mark a
    slot "validated" when its own function never checks it (false NEGATIVE).
    """
    out = []
    for ea, sz in fns:
        o = ea - base
        if o < 0 or sz <= 0 or o + sz > len(b):
            continue
        ins = list(THUMB.disasm(b[o:o + sz], ea))
        cuts = [ea]
        for n, i in enumerate(ins):
            if i.address == ea or not i.mnemonic.startswith("push"):
                continue
            if any(x.mnemonic in ("sub", "subw", "sub.w")
                   and re.match(r"sp,\s*(?:sp,\s*)?#", x.op_str or "")
                   for x in ins[n + 1:n + 4]):
                cuts.append(i.address)
        cuts.append(ea + sz)
        for k in range(len(cuts) - 1):
            out.append([cuts[k], cuts[k + 1] - cuts[k]])
    return out


def copy_helpers(b, base, fns):
    """ARM-mode load+store bodies reached by BLX, decoded inside function extents."""
    tgts = set()
    for ea, size in fns:
        off = ea - base
        if off < 0 or off + size > len(b):
            continue
        for i in THUMB.disasm(b[off:off + size], ea):
            if i.mnemonic == "blx" and (i.op_str or "").startswith("#"):
                tgts.add(int(i.op_str[1:], 16))
    out = set()
    for t in tgts:
        o = t - base
        if not (0 <= o < len(b) - 16):
            continue
        a = list(ARM.disasm(b[o:o + 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 analyse(name, verbose=False):
    b, base, fns = load_module(name)
    helpers = copy_helpers(b, base, fns)
    fns = split_merged(b, base, fns)
    findings = []
    for ea, size in fns:
        off = ea - base
        if off < 0 or size <= 0 or off + size > len(b):
            continue
        ins = list(THUMB.disasm(b[off:off + size], ea))
        # slots initialised from an immediate, and slots whose address is taken
        init, addr_taken, sp_taken = {}, set(), False
        for n, i in enumerate(ins):
            if i.mnemonic.startswith("str"):
                m = SLOT_STORE.match(i.op_str or "")
                if m:
                    src = m.group(1)
                    # XGMI sets the capacity (`mov.w sl,#0x40`) ~12 instructions before it
                    # stores it, so an 8-instruction window loses the positive control.
                    for k in range(n - 1, max(0, n - 40) - 1, -1):
                        mm = re.match(r"(%s),\s*#(0x[0-9a-f]+|\d+)$" % re.escape(src),
                                      ins[k].op_str or "")
                        if mm and ins[k].mnemonic.startswith("mov"):
                            init[int(m.group(2), 0)] = (int(mm.group(2), 0), i.address)
                            break
            if i.mnemonic.startswith("add"):
                m = ADDR_OF.match(i.op_str or "")
                if m:
                    addr_taken.add(int(m.group(2), 0))
            if i.mnemonic.startswith("mov") and re.fullmatch(r"\w+,\s*sp", i.op_str or ""):
                sp_taken = True
        if not init:
            continue
        for n, i in enumerate(ins):
            if not (i.mnemonic in ("blx", "bl") and (i.op_str or "").startswith("#")):
                continue
            if int(i.op_str[1:], 16) not in helpers:
                continue
            # what feeds r2 (the length)?  a re-load of an initialised, address-taken slot
            src_reg, reload_idx, slot = None, None, None
            cur = "r2"
            for k in range(n - 1, max(0, n - 30) - 1, -1):
                m = re.match(r"(%s),\s*(.+)$" % re.escape(cur), ins[k].op_str or "")
                if not m or ins[k].mnemonic.startswith(("cmp", "str", "push", "b")):
                    continue
                if ins[k].mnemonic.startswith("ldr"):
                    mm = SLOT_LOAD.match(ins[k].op_str or "")
                    if mm:
                        slot = int(mm.group(2), 0)
                        src_reg, reload_idx = cur, k
                    break
                m2 = re.match(r"(r\d+|sb|sl|fp|ip)$", m.group(2).strip())
                if m2:
                    cur = m2.group(1)
                    continue
                break
            if slot is None:
                continue
            # The slot must be REACHABLE by a callee: either its own address is taken, or a
            # struct pointer covering it is formed (`add rX, sp, #j` with j <= slot, or
            # `mov rX, sp`). Dropping this along with the immediate-init requirement made v2
            # flag [sp+4] in sub_E08984 -- a constant 48 that no callee ever writes.
            reachable = any(j <= slot for j in addr_taken) or sp_taken
            if not reachable:
                continue
            # v2: the immediate-init requirement is dropped. XGMI initialises its slot from an
            # immediate, but SYS_DRV's equivalent (0xe1ffd2) keeps the out-length in a FIELD of
            # a struct passed by pointer, so it is never immediate-initialised in the caller --
            # and the v1 pattern silently could not match it. What actually matters is: the
            # slot is re-loaded AFTER a call, and the value is used as a length unchecked.
            calls_between = any(ins[k].mnemonic in ("bl", "blx") for k in range(0, reload_idx))
            if not calls_between:
                continue
            cap = init.get(slot, (None, None))[0]
            # is the re-loaded value compared before it is used as the length?
            validated = any(ins[k].mnemonic.startswith(("cmp", "cbz", "cbnz", "subs"))
                            and re.search(r"\b%s\b" % re.escape(src_reg), ins[k].op_str or "")
                            for k in range(reload_idx + 1, n))
            if not validated:
                findings.append((i.address, ea, slot, cap))
    return findings


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--selftest", action="store_true")
    ap.add_argument("--module")
    ap.add_argument("--all", action="store_true")
    a = ap.parse_args()

    if a.selftest:
        xg = analyse("xgmi")
        pos = [f for f in xg if f[0] == 0x402FC4]
        asd = analyse("asd")
        neg = [f for f in asd if f[0] == 0xD12E4C]
        print("SELFTEST (callee-written in/out length slot)")
        print("   POSITIVE xgmi 0x402fc4 reported:      %s" % ("yes" if pos else "NO"))
        if pos:
            print("      slot [sp+%#x], capacity %d" % (pos[0][2], pos[0][3]))
        print("   NEGATIVE asd 0xd12e4c not reported:   %s" % ("yes" if not neg else "NO"))
        ok = bool(pos) and not neg
        print("\nverdict: %s" % ("TRUSTWORTHY" if ok else "NOT trustworthy"))
        return 0 if ok else 1

    mods = (["gcp", "opm", "xgmi", "merit", "fp", "auc", "powerplay", "wspe", "rap", "oc",
             "csu", "tmz", "asd", "sys_drv"] if a.all else [a.module])
    for m in mods:
        try:
            res = analyse(m)
        except Exception as exc:  # noqa: BLE001
            print("%-10s SKIP (%s)" % (m, str(exc)[:50]))
            continue
        print("%-10s %d site(s)%s" % (m, len(res),
              "" if not res else ":"))
        for site, fn, slot, cap in res:
            print("     copy @%#x (fn %#x) length <- slot [sp+%#x]%s, callee-written, "
                  "NOT re-validated" % (site, fn, slot,
                  (" (init %d)" % cap) if cap is not None else " (struct-held)"))
    return 0


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