#!/usr/bin/env python3
"""Compare a copy's MAXIMUM POSSIBLE length against the destination's capacity.

Twelve passes established that "is the length compared?" is the wrong question here. Most
lengths come from byte or halfword fields, so they are type-bounded (255 / 65535) whether or
not anyone compares them, and what decides safety is:

        max_possible_length x scaling   vs   destination capacity

Both sides are recoverable for stack destinations:

  * max length: trace the length register back; `ldrb` -> 255, `ldrh` -> 65535,
    an explicit `cmp rX,#C` -> C, and accumulate scaling from `lsl #n` and the
    `add rD, rA, rA, lsl #n` idiom (which is x(1+2^n), e.g. x5)
  * capacity: destination is `sp` or `sp+k` and the frame is `sub sp,#F`, so capacity = F - k

This is exactly the arithmetic that makes the two known cases differ, which is why they serve
as controls in opposite directions:

  POSITIVE  hdcp `0x402cde`  byte(255) x5 = 1275 into a 232-byte stack buffer  -> MUST fire
  NEGATIVE  xgmi `0x402d4c`  count(<=64) x16 = 1024 into a 1024-byte buffer    -> MUST NOT fire

Usage:
  ./ta_capacity_audit.py --selftest
  ./ta_capacity_audit.py --all
"""
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)}
TAS = ["gcp", "opm", "xgmi", "merit", "fp", "auc", "powerplay", "wspe", "rap", "oc",
       "csu", "tmz"]


def load(name):
    if name in MODULES:
        bp, base = MODULES[name]
        return (REPO / bp).read_bytes(), base, json.loads(CACHE.read_text())[name]
    meta = json.loads((HERE / "syms" / ("%s.meta.json" % name)).read_text())
    p = Path(meta["path"])
    b = (p if p.is_absolute() else REPO / p).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, 0x400100, fns


def helpers(b, base, fns):
    tg = set()
    for ea, sz in fns:
        o = ea - base
        if o < 0 or o + sz > len(b):
            continue
        for i in THUMB.disasm(b[o:o + sz], ea):
            if i.mnemonic == "blx" and (i.op_str or "").startswith("#"):
                tg.add(int(i.op_str[1:], 16))
    out = set()
    for t in tg:
        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 max_len(ins, lo, at, reg):
    """(max value, description) for reg at index `at`, following scaling."""
    cur, scale, desc = reg, 1, []
    for _ in range(8):
        # A compare is only a BOUND if it constrains the magnitude: `cmp reg,#0` is a null
        # test, and taking it as a bound returned max=0 and silently suppressed HDCP's
        # overflow. So look for the definition first and use a compare only to tighten.
        cap_from_cmp = None
        for k in range(at - 1, lo - 1, -1):
            i = ins[k]
            if i.mnemonic.startswith("cmp"):
                m = re.match(r"(%s),\s*#(0x[0-9a-f]+|\d+)$" % re.escape(cur), i.op_str or "")
                if m and int(m.group(2), 0) > 0:
                    cap_from_cmp = int(m.group(2), 0)
                    break
        for k in range(at - 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).strip()
            if i.mnemonic.startswith("ldrb"):
                v = min(0xFF, cap_from_cmp) if cap_from_cmp else 0xFF
                return v * scale, "byte(max %d) x%d" % (v, scale)
            if i.mnemonic.startswith("ldrh"):
                v = min(0xFFFF, cap_from_cmp) if cap_from_cmp else 0xFFFF
                return v * scale, "half(max %d) x%d" % (v, scale)
            if i.mnemonic.startswith("ldr"):
                return None, "word x%d" % scale      # unknown width: not decidable here
            mi = re.fullmatch(r"#(0x[0-9a-f]+|\d+)", rhs)
            if mi:
                v = int(mi.group(1), 0)
                # a literal that looks like a code/data address is not a length: the SYS_DRV
                # false positive reported "max 14779680" (= 0xE1A0E0, an address)
                if v >= 0x100000:
                    return None, "address-like literal"
                return v * scale, "imm x%d" % scale
            # x(1+2^n): add rD, rA, rA, lsl #n
            m5 = re.match(r"(r\d+|sb|sl|ip),\s*(r\d+|sb|sl|ip),\s*lsl\s*#(\d+)$", rhs)
            if m5 and i.mnemonic.startswith("add") and m5.group(1) == m5.group(2):
                scale *= (1 + (1 << int(m5.group(3))))
                cur, at = m5.group(1), k
                break
            msh = re.match(r"(r\d+|sb|sl|ip),\s*#(\d+)$", rhs)
            if msh and i.mnemonic.startswith("lsl"):
                scale *= (1 << int(msh.group(2)))
                cur, at = msh.group(1), k
                break
            mmv = re.fullmatch(r"(r\d+|sb|sl|fp|ip)", rhs)
            if mmv:
                cur, at = mmv.group(1), k
                break
            return None, "expr x%d" % scale
        else:
            return None, "unresolved"
    return None, "deep"


def capacity(ins, lo, at, frame):
    """If the destination (r0) is sp or sp+k, capacity = frame - k."""
    cur = "r0"
    for _ in range(6):
        for k in range(at - 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).strip()
            if rhs == "sp":
                return frame, 0
            ms = re.match(r"sp,\s*#(0x[0-9a-f]+|\d+)$", rhs)
            if ms and i.mnemonic.startswith("add"):
                k2 = int(ms.group(1), 0)
                return frame - k2, k2
            mm = re.fullmatch(r"(r\d+|sb|sl|fp|ip)", rhs)
            if mm:
                cur, at = mm.group(1), k
                break
            return None, None
        else:
            return None, None
    return None, None


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

    IDA merged OPM's cmd-8 handler (0x401320) into 0x4012c0, so a copy in the second function
    was measured against the FIRST function's frame. That produced a false positive here, but
    the same error in the other direction -- picking a larger frame than the real one -- would
    MASK a genuine overflow. So subdivide: an interior `push {...,lr}` followed within a few
    instructions by `sub sp` is a real function entry.
    """
    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
            nxt = ins[n + 1:n + 4]
            if any(x.mnemonic in ("sub", "subw", "sub.w")
                   and re.match(r"sp,\s*(?:sp,\s*)?#", x.op_str or "") for x in nxt):
                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 analyse(name):
    b, base, fns = load(name)
    hp = helpers(b, base, fns)
    fns = split_merged(b, base, fns)
    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))
        frame = None
        for i in ins[:8]:
            # `sub sp,#imm` AND `sub.w sp,sp,#imm` -- missing the second form left frame=None
            # for HDCP's `sub.w sp, sp, #0x214`, so the positive control never fired.
            m = re.match(r"sp,\s*(?:sp,\s*)?#(0x[0-9a-f]+|\d+)$", i.op_str or "")
            if m and i.mnemonic in ("sub", "subw", "sub.w"):
                frame = int(m.group(1), 0)
                break
        if frame is None:
            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 hp:
                continue
            cap, koff = capacity(ins, 0, n, frame)
            if cap is None or cap <= 0:
                continue
            mx, how = max_len(ins, 0, n, "r2")
            if mx is None:
                continue
            if mx > cap:
                out.append((i.address, ea, mx, how, cap, koff))
    return out


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

    if a.selftest:
        hd = analyse("hdcp")
        pos = [f for f in hd if f[0] == 0x402CDE]
        xg = analyse("xgmi")
        neg = [f for f in xg if f[0] == 0x402D4C]
        print("SELFTEST (max length vs destination capacity)")
        print("   POSITIVE hdcp 0x402cde fires:        %s" % ("yes" if pos else "NO"))
        if pos:
            print("      max %d (%s) into capacity %d" % (pos[0][2], pos[0][3], pos[0][4]))
        print("   NEGATIVE xgmi 0x402d4c does not fire: %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

    for m in (TAS + ["asd", "sys_drv"] if a.all else [a.module]):
        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 res else ""))
        for site, fn, mx, how, cap, koff in res:
            print("     copy @%#x (fn %#x) max %d (%s) > capacity %d (sp+%#x)"
                  % (site, fn, mx, how, cap, koff))
    return 0


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