#!/usr/bin/env python3
"""Field-level length audit: is the STRUCT FIELD a length comes from ever written unbounded?

Passes 1-8 kept failing the same way. Register-level analysis cannot answer the question that
decides these cases, because the length is loaded from a struct field at the copy site and the
bound lives wherever that field was *written* -- possibly in another function, possibly from a
different field. Concretely:

  * HDCP `0x402cde`   length = `ldrb [r5,#0x38]`; the field is written from a request byte and
                      never bounded  -> REAL BUG (confirmed on hardware)
  * SYS_DRV `+0x44`   length field written from another struct's `+0x120`, and THAT value is
                      bounded `> 0x30 -> error` in the writing function -> NOT a bug

So this tool works at field granularity:

  1. find copy sites whose length is `LOAD [base + off]`
  2. find every store in the module to the same field offset, `str rS, [rX + off]`
  3. classify each writer's source: immediate, a value compared/masked in that function
     (bounded), or an unbounded load / unbounded parameter
  4. report the copy site only if some writer supplies an unbounded value

It is an offset-keyed approximation of field tracking -- two different structs sharing an
offset are conflated -- which makes it over-report rather than under-report, the safe
direction for a hunting tool.

!! DO NOT TRUST THIS TOOL'S OUTPUT -- ITS POSITIVE CONTROL PASSED SPURIOUSLY !!

  The control appeared to pass: hdcp `0x402cde` was reported, with four "unbounded writers"
  of field `+0x38`. Inspecting those four addresses shows they are not stores at all:

      0x4013f6  movw r1, #0xffff
      0x40169c  mov  r6, r0
      0x4016dc  add  r0, sp, #8
      0x4016f6  pop  {r4, pc}

  They are garbage produced by the resynchronising linear sweep, which decodes literal pools
  and ARM-mode regions as Thumb; the store regex then matched inside that garbage. So the
  tool never demonstrated it can trace a real field-writer chain, and its negatives (0 sites
  for the TAs and ASD) are NOT evidence.

  Root cause: regex matching over a resynchronised linear sweep. The fix is to decode only
  within known function boundaries (as fw/ta_copylen_slice.py does via fw/libdumps), and to
  verify that every "writer" address really disassembles to a store before counting it.

  NEGATIVE control (sys_drv sub_E0D6B4 not reported) does still hold, but a detector with
  only a negative control cannot support a negative result.

Usage:
  ./ta_fieldflow_audit.py --selftest
  ./ta_fieldflow_audit.py --module sys_drv|asd|hdcp|<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),
}


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


def fn_starts(ins):
    return [i.address for i in ins if i.mnemonic.startswith("push")]


def owner_idx(starts, ins, n):
    a = ins[n].address
    lo = 0
    for k, i in enumerate(ins):
        if i.address in starts and i.address <= a:
            lo = k
        elif i.address > a:
            break
    return lo


def bounded_here(ins, lo, hi, reg):
    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


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


def field_of_load(ins, lo, at, reg):
    """If reg came from `ldr rD,[rB,#off]`, return off."""
    cur = reg
    for _ in range(5):
        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
            if i.mnemonic.startswith("ldr"):
                mm = LOADF.match(i.op_str or "")
                return int(mm.group(3), 0) if mm else None
            rhs = m.group(2).strip()
            m2 = re.match(r"(r\d+|sb|sl|fp|ip)$", rhs)
            if m2:
                cur, at = m2.group(1), k
                break
            m3 = re.match(r"(r\d+|sb|sl|fp|ip)\b", rhs)
            if m3 and i.mnemonic.startswith(("add", "sub", "mul", "lsl", "orr", "mla")):
                cur, at = m3.group(1), k
                break
            return None
        else:
            return None
    return None


def base_origin(ins, lo, at, reg):
    """Where the BASE pointer came from: a literal-pool global, a parameter, or the stack.

    Keying a field only by its offset conflates every struct sharing it -- SYS_DRV's `+0x4`
    reported 375 writers, which is noise rather than signal. Keying by (base provenance,
    offset) separates distinct objects well enough to make the list triageable.
    """
    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") and "pc" in rhs:
                return "global"          # from a literal pool: a module global
            if i.mnemonic.startswith("ldr"):
                return "deref"
            if "sp" in rhs:
                return "stack"
            m2 = re.match(r"(r\d+|sb|sl|fp|ip)$", rhs)
            if m2:
                cur, at = m2.group(1), k
                break
            return "expr"
        else:
            return "param:%s" % cur if cur in ("r0", "r1", "r2", "r3") else "unknown"
    return "unknown"


def writers_of(ins, starts, off, want_scope=None):
    """Every `str rS,[rX,#off]` in the module, with a verdict on rS.

    When want_scope is given, only writers whose BASE pointer has the same provenance are
    counted -- that is what removes the offset-only conflation.
    """
    out = []
    for n, i in enumerate(ins):
        if not i.mnemonic.startswith("str"):
            continue
        m = STOREF.match(i.op_str or "")
        if not m or int(m.group(3), 0) != off:
            continue
        src = m.group(1)
        # A store through sp/fp is a STACK SLOT, not a struct field. Counting them conflated
        # the stack-canary save `str r0,[sp,#0x88]` with writes to a field at +0x88, which is
        # how the top SYS_DRV candidate turned out to be __stack_chk_guard.
        # sp/fp only. Excluding r7 as well broke the positive control: in Thumb r7 is often a
        # frame pointer, but HDCP's field writers use it as an ordinary struct pointer, so
        # excluding it discarded the one bug we can check against.
        if m.group(2) in ("sp", "fp"):
            continue
        if want_scope is not None:
            wlo = owner_idx(starts, ins, n)
            if base_origin(ins, wlo, n, m.group(2)) != want_scope:
                continue
        lo = owner_idx(starts, ins, n)
        if bounded_here(ins, lo, n, src):
            out.append((i.address, "bounded in writer"))
            continue
        f = field_of_load(ins, lo, n, src)
        if f is not None:
            out.append((i.address, "UNBOUNDED from field +%#x" % f))
        else:
            out.append((i.address, "UNBOUNDED (%s)" % src))
    return out


def copy_helpers(b, ins, base):
    out = set()
    for i in ins:
        if i.mnemonic == "blx" and (i.op_str or "").startswith("#"):
            t = int(i.op_str[1:], 16)
            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(path, base):
    b = (REPO / path).read_bytes() if not Path(path).is_absolute() else Path(path).read_bytes()
    ins = sweep(b, base)
    starts = set(fn_starts(ins))
    helpers = copy_helpers(b, ins, base)
    wcache = {}
    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
        lo = owner_idx(starts, ins, n)
        off = field_of_load(ins, lo, n, "r2")
        if off is None:
            continue
        # Keep readers whose base provenance is vague: HDCP's confirmed overflow is one of
        # them, so dropping them fails the positive control. Only the WRITER-side sp/fp
        # exclusion above is a sound restriction.
        scope = base_origin(ins, lo, n, "r2")     # provenance of the READER's base
        key = (off, scope)
        if key not in wcache:
            wcache[key] = writers_of(ins, starts, off, want_scope=scope)
        bad = [w for w in wcache[key] if "UNBOUNDED" in w[1]]
        if bad:
            findings.append((hex(i.address), hex(ins[lo].address), off,
                             len(wcache[key]), bad[:4], scope))
    return findings


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

    if a.selftest:
        hd = analyse(*MODULES["hdcp"])
        pos = [f for f in hd if f[0] == "0x402cde"]
        sd = analyse(*MODULES["sys_drv"])
        neg = [f for f in sd if f[0] in ("0xe0d6c4", "0xe0d6d8")]
        print("SELFTEST (field-level)")
        print("   POSITIVE hdcp 0x402cde reported:           %s" % ("yes" if pos else "NO"))
        if pos:
            print("      %s" % (pos[0][4],))
        print("   NEGATIVE sys_drv sub_E0D6B4 not reported:  %s"
              % ("yes" if not neg else "NO -- %s" % neg))
        ok = bool(pos) and 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())
        path, base = meta["path"], 0x400100
    res = analyse(path, base)
    print("=== %s: %d copy site(s) whose length field has an unbounded writer"
          % (a.module, len(res)))
    for site, fn, off, nw, bad, scope in res:
        print("   copy @%s (fn %s) len <- field +%#x [base:%s] ; %d writer(s)"
              % (site, fn, off, scope, nw))
        for wa, v in bad:
            print("        writer @%#x  %s" % (wa, v))
    return 0


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