#!/usr/bin/env python3
"""Recover helper-function arity by dataflow, so the remaining functions get signatures.

179 functions got prototypes from an established role. The other ~450 are helpers whose
semantics are unknown, and guessing their argument count would repeat the mistake that
produced `v6[23].decode_handle`: output that reads well and says something false.

Arity, however, is not a guess. On AAPCS, r0-r3 carry the first four arguments, so a
register that is READ BEFORE IT IS WRITTEN anywhere in the function is an incoming
parameter. That is a dataflow fact recoverable from the disassembly.

The analysis is block-aware: instructions are grouped into basic blocks at branch targets,
and blocks are walked from the entry in reverse-postorder so a register written on one path
is not treated as defined on another. A register is an argument if ANY path reads it before
writing it.

VALIDATION IS THE POINT. The 179 role-derived prototypes have known arity -- command
handlers take 2, the ASD RPC thunks take 1, zero_fill takes 2, memcmp takes 3. The analyzer
must reproduce those before its output is trusted anywhere else. `--selftest` does exactly
that and reports disagreements rather than hiding them.

Usage:
  ./ta_arity.py --prefix gcp --selftest
  ./ta_arity.py --prefix gcp --apply
"""
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"))
BASE = 0x400100
ARGREGS = ["r0", "r1", "r2", "r3"]

# Known arity for the role-derived prototypes: the positive control.
EXPECT_ARITY = [
    (re.compile(r"zero_fill"), 2),
    (re.compile(r"memcmp"), 3),
]

def control_set() -> dict[int, int]:
    """Addresses with PROVEN arity: every handler the dispatch table points at is called
    as `mov r0,r5 ; mov r1,r4 ; bl handler`, i.e. exactly two arguments."""
    m = json.loads((HERE / "cmdmaps" / "gcp_ea.json").read_text())
    return {int(v, 16): 2 for v in m.values()}


def blocks(md, code: bytes, ea: int):
    """Split a function into basic blocks keyed by start address."""
    insns = list(md.disasm(code, ea))
    if not insns:
        return {}, []
    starts = {ea}
    by_addr = {i.address: i for i in insns}
    for i in insns:
        if i.group(cs.CS_GRP_JUMP) or i.mnemonic.startswith(("b", "cb")):
            m = re.search(r"#(0x[0-9a-f]+)", i.op_str)
            if m:
                t = int(m.group(1), 16)
                if t in by_addr:
                    starts.add(t)
            nxt = i.address + i.size
            if nxt in by_addr:
                starts.add(nxt)
    starts = sorted(s for s in starts if s in by_addr)
    blks = {}
    for n, s in enumerate(starts):
        end = starts[n + 1] if n + 1 < len(starts) else insns[-1].address + insns[-1].size
        blks[s] = [i for i in insns if s <= i.address < end]
    return blks, insns


def succs(blk, blks):
    """Successors of a block: fallthrough and/or direct branch target."""
    if not blk:
        return []
    last = blk[-1]
    out = []
    m = re.search(r"#(0x[0-9a-f]+)", last.op_str)
    is_branch = last.mnemonic.startswith(("b", "cb")) and not last.mnemonic.startswith("bl")
    if is_branch and m and int(m.group(1), 16) in blks:
        out.append(int(m.group(1), 16))
    unconditional = last.mnemonic in ("b", "b.w") or (
        last.mnemonic.startswith("pop") and "pc" in last.op_str)
    if not unconditional:
        nxt = last.address + last.size
        if nxt in blks:
            out.append(nxt)
    return out


def reads_writes(i):
    """(registers read, registers written) restricted to r0-r3, from the operand text."""
    ops = i.op_str
    # a call clobbers r0-r3 and consumes its own arguments; treat it as defining them so a
    # post-call read is not mistaken for an incoming parameter
    if i.mnemonic in ("bl", "blx"):
        return set(), set(ARGREGS)
    # push/pop is stack save/restore, not argument dataflow: `push {r3,r4,r5,lr}` is the
    # standard 8-byte-alignment idiom and READS r3, which made every such function look
    # like it took four arguments (the selftest caught this as 15 disagreements).
    if i.mnemonic.startswith(("push", "pop")):
        return set(), set()
    # stm/ldm against a general register IS argument dataflow, and excluding it was an
    # over-correction: the ASD RPC thunks pass their arguments with exactly this idiom
    # (`stm.w ip, {r0,r1,r2,r3}` marshals four arguments into the request buffer), so
    # dropping it reported those thunks as taking none.
    if i.mnemonic.startswith("stm"):
        return set(re.findall(r"\br[0-3]\b", ops)), set()
    if i.mnemonic.startswith("ldm"):
        lst = ops.split("{", 1)[1] if "{" in ops else ""
        return set(), set(re.findall(r"\br[0-3]\b", lst))
    toks = re.findall(r"\br([0-3])\b", ops)
    regs = {"r" + t for t in toks}
    if not regs:
        return set(), set()
    dst = re.match(r"\s*(r[0-3])\s*,", ops)
    writes = set()
    reads = set(regs)
    # store instructions read their first operand rather than writing it
    if dst and not i.mnemonic.startswith(("str", "stm", "cmp", "cmn", "tst", "teq", "push")):
        writes = {dst.group(1)}
        reads = regs - writes
        # read-modify-write forms still read the destination
        if i.mnemonic.startswith(("add", "sub", "orr", "and", "eor", "mul", "lsl", "lsr",
                                  "asr", "bic", "rsb", "adc", "sbc")) and ops.count(",") < 2:
            reads |= writes
    return reads, writes


def arity(md, body: bytes, ea: int, size: int, _depth: int = 0) -> int | None:
    """Argument count, following tail-jump veneers.

    These images are full of 4-byte veneers of the form `movs r3, r0 ; b <real>`. Typed on
    their own they look 1-argument, which silently dropped the length argument of HDCP's
    memcpy helper and made a KNOWN stack overflow invisible to pseudocode analysis. So a
    veneer's arity is the union of what it reads and what its target takes.
    """
    off = ea - BASE
    if off < 0 or size <= 0 or off + size > len(body):
        return None
    blks, insns = blocks(md, body[off:off + size], ea)
    if not blks:
        return None
    args: set[str] = set()
    # walk each path from the entry; a register read before any write on that path is an arg
    seen = set()
    stack = [(ea, frozenset())]
    while stack:
        b, defined = stack.pop()
        key = (b, defined)
        if key in seen:
            continue
        seen.add(key)
        d = set(defined)
        for i in blks.get(b, []):
            r, w = reads_writes(i)
            args |= {x for x in r if x not in d}
            d |= w
        for s in succs(blks.get(b, []), blks):
            stack.append((s, frozenset(d)))
    # veneer: a couple of instructions ending in an unconditional branch elsewhere
    if _depth < 3 and len(insns) <= 4:
        last = insns[-1]
        if last.mnemonic in ("b", "b.w"):
            m = re.search(r"#(0x[0-9a-f]+)", last.op_str or "")
            if m:
                tgt = int(m.group(1), 16)
                if tgt != ea:
                    sub = arity(md, body, tgt, 0x400, _depth + 1)
                    if sub:
                        here = max((int(a[1]) for a in args), default=-1) + 1
                        return max(sub, here)
    if not args:
        return 0
    # AAPCS: arguments are contiguous from r0, so arity is the highest used index + 1
    return max(int(a[1]) for a in args) + 1


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--prefix", required=True)
    ap.add_argument("--selftest", action="store_true")
    ap.add_argument("--apply", action="store_true")
    args = ap.parse_args()

    meta = json.loads((HERE / "syms" / ("%s.meta.json" % args.prefix)).read_text())
    body = Path(meta["path"])
    if not body.is_absolute():
        body = REPO / body
    b = body.read_bytes()
    md = cs.Cs(cs.CS_ARCH_ARM, cs.CS_MODE_THUMB)
    md.detail = True

    from idatui import remote_ops
    from idatui.nexus_client import NexusClient

    elf = json.loads((HERE / "syms" / "index.json").read_text())[args.prefix]["elf"]
    c = NexusClient(str(REPO / elf))
    c.connect()

    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
    fns = [{"ea": int(str(i["addr"]), 16), "name": i["name"], "size": i.get("size", 0)}
           for i in raw]

    if args.selftest:
        ok = bad = skip = 0
        fails = []
        ctrl = control_set() if args.prefix == "gcp" else {}
        for f in fns:
            want = ctrl.get(f["ea"])
            if want is None:
                want = next((n for rx, n in EXPECT_ARITY if rx.search(f["name"])), None)
            if want is None:
                continue
            got = arity(md, b, f["ea"], f["size"])
            if got is None:
                skip += 1
                continue
            if got == want:
                ok += 1
            else:
                bad += 1
                fails.append((f["name"], want, got))
        print("SELFTEST against known-arity functions: %d agree, %d disagree, %d skipped"
              % (ok, bad, skip))
        for n, w, g in fails[:15]:
            print("   %-46s expected %d, inferred %d" % (n, w, g))
        print("\nverdict: %s" % ("TRUSTWORTHY" if bad == 0 else
                                 "NOT trustworthy -- do not apply inferred arity"))
        c.close()
        return 0 if bad == 0 else 1

    if args.apply:
        typed = 0
        for f in fns:
            nm = f["name"]
            # Keep the role-derived prototypes that the control PROVED correct, and let
            # inference fix the ones it disproved: the ASD RPC thunks are not uniform
            # (gcp_asd_rpc_304 takes none, _107 three, _202 four) yet all were declared
            # with a single argument.
            if re.search(r"_cmd[0-9A-Fa-f]{2,4}|zero_fill|memcmp|^start$", nm):
                continue
            n = arity(md, b, f["ea"], f["size"])
            if n is None:
                continue
            params = ", ".join("int a%d" % (k + 1) for k in range(n)) or "void"
            proto = "int %s(%s);" % (nm, params)
            r = c.call(remote_ops.set_type, edits=[{"addr": hex(f["ea"]), "signature": proto}])
            rows = r.get("result", []) if isinstance(r, dict) else r
            if rows and rows[0].get("ok"):
                typed += 1
        print("applied inferred-arity prototypes to %d helpers" % typed)
        print("save: %s" % json.dumps(c.save_database())[:90])
    c.close()
    return 0


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