#!/usr/bin/env python3
"""Field-level length audit, v2 — decoding ONLY inside IDA function boundaries.

v1 (`ta_fieldflow_audit.py`) is retracted: its positive control passed because a
resynchronising linear sweep decoded literal pools and ARM regions as Thumb, and the store
regex matched inside that garbage. The "writers" it reported for HDCP's field `+0x38` were a
`movw`, a `mov`, an `add` and a `pop`.

v2 removes the cause rather than patching symptoms:

  * function boundaries come from IDA (`list_funcs` on the module's database), so every
    instruction decoded starts at a real function entry and runs to that function's end;
  * every candidate writer is RE-VERIFIED: the address must disassemble, inside its own
    function, to an instruction whose mnemonic really begins with `str`, and its operand form
    must really be `[base, #imm]`. Anything else is dropped and counted as a rejected match,
    which is reported so the rejection rate is visible rather than silent.

Same question as v1: for a copy whose length is `LOAD [base+off]`, is that field ever written
from an unbounded value?

CONTROLS (both must hold, and the positive must now hold for the RIGHT reason):
  POSITIVE  hdcp `0x402cde` reported, and its writers must be verified real stores
  NEGATIVE  sys_drv `sub_E0D6B4` not reported (its writer bounds the value)

Usage:
  ./ta_fieldflow2.py --cache            # pull function boundaries from IDA once
  ./ta_fieldflow2.py --selftest
  ./ta_fieldflow2.py --module sys_drv|asd|hdcp
"""
from __future__ import annotations

import argparse
import json
import re
import sys
from pathlib import Path

import capstone as cs

HERE = Path(__file__).resolve().parent
REPO = HERE.parent
sys.path.insert(0, str(Path.home() / "dev" / "ida-tui-maybe"))
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, "psp_out/hdcp2.elf"),
    "sys_drv": ("psp_out/sys_drv_code.bin", 0xE00100, "psp_out/sys_drv2.elf"),
    "asd":     ("psp_out/asd_code.bin",     0xD00100, "psp_out/asd2.elf"),
}

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 pull_bounds():
    from idatui import remote_ops
    from idatui.nexus_client import NexusClient
    out = {}
    for name, (_b, _base, elf) in MODULES.items():
        c = NexusClient(str(REPO / elf))
        c.connect(timeout=900)
        raw, off = [], 0
        while True:
            rr = c.call(remote_ops.list_funcs, queries=[{"count": 500, "offset": off}])
            blk = (rr.get("result") or [{}])[0]
            raw += blk.get("data") or []
            total = int(blk.get("total") or 0)
            off = int(blk.get("next_offset") or 0)
            if not off or len(raw) >= total:
                break
        out[name] = [[int(str(i["addr"]), 16), int(i.get("size") or 0)] for i in raw
                     if int(i.get("size") or 0) > 0]
        print("   %-8s %d functions" % (name, len(out[name])))
        c.close()
    CACHE.parent.mkdir(exist_ok=True)
    CACHE.write_text(json.dumps(out))
    print("cached -> %s" % CACHE)


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

    This is a FALSE-NEGATIVE fix, not cosmetics: `bounded_here` scans from the function start,
    so in a merged extent a compare belonging to the PRECEDING function can mark a value
    "bounded" that nothing in its own function bounds -- masking a real site.
    """
    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 decode(b: bytes, base: int, fns: list) -> tuple[list, dict]:
    """Instructions decoded strictly within function extents, plus an address index."""
    ins = []
    for ea, size in fns:
        off = ea - base
        if off < 0 or size <= 0 or off + size > len(b):
            continue
        ins.extend(THUMB.disasm(b[off:off + size], ea))
    ins.sort(key=lambda i: i.address)
    return ins, {i.address: n for n, i in enumerate(ins)}


def fn_range(fns: list, addr: int):
    for ea, size in fns:
        if ea <= addr < ea + size:
            return ea, size
    return None


def bounded_here(ins, lo, hi, reg) -> bool:
    """Bounded = compared, masked, OR assigned a constant.

    The third case was missing in v1 and produced false "unbounded writers": SYS_DRV's
    `movs r2,#0 ; str.w r2,[r1,#0x88]` stores a constant zero, which is as bounded as a value
    can be. Walk backwards for the nearest definition of reg first.
    """
    for k in range(hi - 1, lo - 1, -1):
        i = ins[k]
        m = re.match(r"(%s),\s*(.+)$" % re.escape(reg), i.op_str or "")
        if not m or i.mnemonic.startswith(("cmp", "str", "push", "tst", "b")):
            continue
        if re.fullmatch(r"#(0x[0-9a-f]+|\d+)", m.group(2).strip()) and \
           i.mnemonic.startswith(("mov", "movs", "movw")):
            return True                      # defined by an immediate
        break                                # nearest definition is not an immediate
    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
        # Backward scan: if the tracked register was DEFINED by `mov rTracked, rSrc`, then
        # rSrc is what we must keep following. v1 added group(1) when group(2) was tracked,
        # i.e. it propagated FORWARD while scanning backward, so a bound on the source was
        # missed -- that is how SYS_DRV's `mov r7,r2 ; cmp r2,#0x1f9c` looked unbounded.
        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 field_of(ins, lo, at, reg):
    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 helpers_of(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(name: str, verbose=False):
    if name in MODULES:
        bpath, base, _elf = MODULES[name]
        b = (REPO / bpath).read_bytes()
        fns = json.loads(CACHE.read_text())[name]
    else:
        # a TA prefix: boundaries already exist in fw/libdumps (from IDA), body via syms meta
        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]
    fns = split_merged(b, base, fns)
    ins, byaddr = decode(b, base, fns)
    helpers = helpers_of(b, ins, base)
    starts = {ea for ea, _ in fns}

    def lo_of(n):
        a = ins[n].address
        r = fn_range(fns, a)
        return byaddr.get(r[0], 0) if r else 0

    # index every VERIFIED field store: address must decode to a real str [base,#imm]
    stores = {}
    rejected = 0
    for n, i in enumerate(ins):
        if not i.mnemonic.startswith("str"):
            continue
        m = STOREF.match(i.op_str or "")
        if not m:
            rejected += 1
            continue
        if m.group(2) in ("sp", "fp"):
            continue                       # stack slot, not a struct field
        stores.setdefault(int(m.group(3), 0), []).append((n, m.group(1)))
    if verbose:
        print("   verified field stores: %d offsets; %d str-form rejects"
              % (len(stores), rejected))

    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 = lo_of(n)
        off = field_of(ins, lo, n, "r2")
        if off is None:
            continue
        bad = []
        for (wn, src) in stores.get(off, []):
            wlo = lo_of(wn)
            if not bounded_here(ins, wlo, wn, src):
                bad.append((ins[wn].address, src))
        if bad:
            findings.append((i.address, ins[lo].address, off, len(stores.get(off, [])), bad[:4]))
    return findings


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

    if a.cache:
        pull_bounds()
        return 0
    if not CACHE.exists():
        print("run --cache first"); return 2

    if a.selftest:
        hd = analyse("hdcp", verbose=True)
        pos = [f for f in hd if f[0] == 0x402CDE]
        print("SELFTEST v2 (boundary-based, writers verified)")
        print("   POSITIVE hdcp 0x402cde reported:  %s" % ("yes" if pos else "NO"))
        if pos:
            print("      field +%#x, writers: %s"
                  % (pos[0][2], [(hex(w), s) for w, s in pos[0][4]]))
        sd = analyse("sys_drv")
        neg = [f for f in sd if f[0] in (0xE0D6C4, 0xE0D6D8)]
        print("   NEGATIVE sys_drv sub_E0D6B4 clean: %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

    res = analyse(a.module, verbose=True)
    print("=== %s: %d site(s)" % (a.module, len(res)))
    for site, fn, off, nw, bad in res:
        print("   copy @%#x (fn %#x) len <- field +%#x ; %d verified writer(s)"
              % (site, fn, off, nw))
        for wa, src in bad:
            print("        unbounded writer @%#x (src %s)" % (wa, src))
    return 0


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