#!/usr/bin/env python3
"""Inter-procedural length/index audit: propagate caller-established invariants.

Eight passes of intra-procedural scanning produced the same failure three times -- a callee
whose arithmetic is locally unguarded, bounded by its caller:

  * `OPM 0x402106`      re-loaded host index, bounded by the callee it is passed to
  * `SYS_DRV sub_E0D6B4` unsigned `size - len` with no compare, bounded by `sub_E08EBC`
  * `powerplay 0x109a`   unchecked `12 * index`, emitter passes a constant

So an intra-procedural detector's output is a work queue, not a finding list. This one asks
the question that actually matters:

    a value reaches a length or a scaled index unguarded in function F.
    Is it bounded by EVERY path that can supply it?

Method
  1. resynchronising Thumb sweep -> instructions, function starts (prologues), call graph
  2. per function, find risk sites: a length/index that traces back to a PARAMETER (r0-r3)
     with no compare/mask in that function
  3. for each caller, look at what it passes in that argument position: an immediate is safe,
     a bounded local is safe, an unbounded memory load is NOT -- and if the caller itself got
     it from a parameter, recurse (bounded depth)
  4. report only sites where some path supplies an unbounded value

SCOPE LIMITATION -- READ THIS BEFORE TRUSTING ANY OUTPUT
  This tool handles only one class: a length that is an INCOMING PARAMETER of the function
  performing the copy, so that callers can be asked what they supply. That was chosen to make
  caller propagation tractable, and it excludes the other class -- a length loaded from a
  STRUCT FIELD -- by construction.

  HDCP's confirmed overflow is the excluded class: its length is `ldrb [r5,#0x38]`, a field
  load, reached directly at the copy site with no wrapper function. So the positive control
  FAILS, and it fails for a structural reason rather than a fixable bug:

      ./ta_interproc_audit.py --selftest
      POSITIVE hdcp 0x402cde reported:  NO
      NEGATIVE sys_drv sub_E0D6B4:      correctly not reported

  The negative control passes -- it does correctly dismiss the bounded-by-caller case that
  motivated the tool. But with no passing positive control, an empty result here means
  nothing, and per this tree's rule (every detector must first rediscover a known bug) the
  output is NOT reported as a negative.

  Making it sound requires field-level value tracking, not register tracking: the SYS_DRV
  case showed the bound can sit on a DIFFERENT field (`+0x120`) from the one consumed
  (`+0x44`), so a module-wide "is this offset ever compared" heuristic would re-introduce the
  false positive it was built to remove. That is a dataflow engine, not an afternoon.

Usage:
  ./ta_interproc_audit.py --selftest
  ./ta_interproc_audit.py --module sys_drv|asd|<ta prefix>
"""
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
THUMB = cs.Cs(cs.CS_ARCH_ARM, cs.CS_MODE_THUMB)
ARM = cs.Cs(cs.CS_ARCH_ARM, cs.CS_MODE_ARM)

MODULES = {
    "sys_drv": ("psp_out/sys_drv_code.bin", 0xE00100),
    "asd": ("psp_out/asd_code.bin", 0xD00100),
    "hdcp": ("psp_out/hdcp_code.bin", 0x400100),
}
ARGREGS = ("r0", "r1", "r2", "r3")


def sweep(b: bytes, base: int):
    ins, off = [], 0
    while off < len(b) - 1:
        got = 0
        for i in THUMB.disasm(b[off:], base + off):
            ins.append(i)
            got += i.size
            if got > 0x4000:
                break
        off += got if got else 2
    return ins, {i.address: n for n, i in enumerate(ins)}


def functions(ins: list) -> list[int]:
    return [i.address for i in ins if i.mnemonic.startswith("push")]


def owner(starts: list[int], addr: int) -> int | None:
    lo = None
    for s in starts:
        if s <= addr:
            lo = s
        else:
            break
    return lo


def callgraph(ins: list, starts: list[int]):
    cg = defaultdict(list)          # callee -> [(caller_fn, call_addr, insn_index)]
    for n, i in enumerate(ins):
        if i.mnemonic in ("bl", "blx") and (i.op_str or "").startswith("#"):
            t = int(i.op_str[1:], 16)
            f = owner(starts, i.address)
            if f is not None:
                cg[t].append((f, i.address, n))
    return cg


def bounded_in(ins: list, lo: int, hi: int, reg: str) -> bool:
    """Is reg (or a mov-chain source) compared or masked anywhere in [lo,hi)?"""
    chain = {reg}
    for k in range(lo, hi):
        i = ins[k]
        ops = i.op_str or ""
        if i.mnemonic.startswith(("cmp", "cmn", "tst", "cbz", "cbnz", "subs")):
            if any(re.search(r"\b%s\b" % r, ops) for r in chain):
                return True
        if i.mnemonic.startswith(("and", "ubfx", "uxtb", "uxth", "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(2) in chain and i.mnemonic.startswith("mov"):
            chain.add(m.group(1))
    return False


def traces_to_param(ins: list, lo: int, at: int, reg: str) -> str | None:
    """If reg's value at `at` came from an incoming parameter, return that parameter."""
    cur = reg
    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 i.mnemonic.startswith("ldr"):
                return None                        # from memory, not a parameter
            m2 = re.match(r"(r\d+|sb|sl|fp|ip)$", rhs)
            if m2:
                cur, at = m2.group(1), k
                break
            return None
        else:
            return cur if cur in ARGREGS else None
    return None


def arg_at_call(ins: list, lo: int, call_idx: int, reg: str):
    """Classify what a caller puts in `reg` before a call: immediate / bounded / UNBOUNDED."""
    cur = reg
    for _ in range(6):
        for k in range(call_idx - 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 re.fullmatch(r"#(0x[0-9a-f]+|\d+)", rhs):
                return "immediate %s" % rhs
            if i.mnemonic.startswith("ldr"):
                return ("bounded-load" if bounded_in(ins, lo, call_idx, cur)
                        else "UNBOUNDED-load %s" % rhs)
            m2 = re.match(r"(r\d+|sb|sl|fp|ip)$", rhs)
            if m2:
                cur = m2.group(1)
                call_idx = k
                break
            return "bounded-expr" if bounded_in(ins, lo, call_idx, cur) else "expr %s" % rhs
        else:
            return "parameter %s" % cur if cur in ARGREGS else "unknown"
    return "unknown"


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 analyse(path: str, base: int):
    b = (REPO / path).read_bytes()
    ins, byaddr = sweep(b, base)
    starts = functions(ins)
    cg = callgraph(ins, starts)
    helpers = copy_helpers(b, ins, base)
    findings = []

    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
        f = owner(starts, i.address)
        if f is None:
            continue
        lo = byaddr[f]
        if bounded_in(ins, lo, n, "r2"):
            continue                       # bounded locally: not interesting
        param = traces_to_param(ins, lo, n, "r2")
        if not param:
            continue                       # from memory: the intra-procedural passes covered it
        # the length is an incoming parameter -- ask every caller what it supplies
        verdicts = []
        for (cf, ca, ci) in cg.get(f, []):
            clo = byaddr.get(owner(starts, ca) or -1)
            if clo is None:
                continue
            verdicts.append((hex(ca), arg_at_call(ins, clo, ci, param)))
        if not verdicts:
            verdicts = [("<no caller found>", "unknown")]
        if any("UNBOUNDED" in v or v.startswith("parameter") for _, v in verdicts):
            findings.append((hex(i.address), hex(f), param, verdicts))
    return findings


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

    if a.selftest:
        ok = True
        hd = analyse(*MODULES["hdcp"])
        pos = [f for f in hd if f[0] == "0x402cde"]
        print("SELFTEST")
        print("   POSITIVE hdcp 0x402cde reported:              %s"
              % ("yes" if pos else "NO -- untrustworthy"))
        ok &= bool(pos)
        sd = analyse(*MODULES["sys_drv"])
        neg = [f for f in sd if f[0] in ("0xe0d6c4", "0xe0d6d8")]
        print("   NEGATIVE sys_drv sub_E0D6B4 NOT reported:     %s"
              % ("yes" if not neg else "NO -- false positive %s" % neg))
        ok &= not neg
        print("\nverdict: %s" % ("TRUSTWORTHY" if ok else "NOT trustworthy"))
        return 0 if ok else 1

    path, base = MODULES.get(a.module, (None, None))
    if path is None:
        meta = json.loads((HERE / "syms" / ("%s.meta.json" % a.module)).read_text())
        p = Path(meta["path"])
        path = str(p if p.is_absolute() else p)
        base = 0x400100
    res = analyse(path, base)
    print("=== %s: %d site(s) where a caller supplies an unbounded length" % (a.module, len(res)))
    for site, fn, param, verdicts in res:
        print("   copy @%s in fn %s, length = incoming %s" % (site, fn, param))
        for ca, v in verdicts[:6]:
            print("        caller @%-12s supplies %s" % (ca, v[:60]))
    return 0


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