#!/usr/bin/env python3
"""Find scaled-index dereferences whose index is not bounded.

Third shape after length and destination: `base + K*index`, where the index comes from
memory (a request field) and nothing compares it against a limit. This is the
`powerplay 0x109a` shape (`dword_E11A9C + 12 * *a1`, unchecked) but searched inside the TAs.

Detection is at the disassembly layer, because that is where the scaling is visible:
  * `mla rD, rIdx, rK, rBase`         index * constant + base
  * `add rD, rBase, rIdx, lsl #N`     index << N + base
  * `lsls rT, rIdx, #N` then `add`    the two-instruction form
then check whether rIdx is compared against anything before the use, anywhere in the
function, and whether its value came from a load.

CONTROL: SYS_DRV's `sub_E11A78` (opcode 0x109a) must be reported as unbounded, and a
function with a real bound (`GCP FE0D`, which constrains the index via `arg10 >> 30`) must
not be. Without both, the detector is not trustworthy.

Usage:
  ./ta_index_audit.py --selftest
  ./ta_index_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)

SCALE_MLA = re.compile(r"^(\w+),\s*(\w+),\s*(\w+),\s*(\w+)$")
SCALE_SHIFT = re.compile(r"^(\w+),\s*(\w+),\s*(\w+),\s*lsl\s*#(\d+)$")


def bounded(ins: list, reg: str, upto: int) -> bool:
    """Is `reg` (or anything it was copied from) compared against a constant, or masked?"""
    chain = {reg}
    for k in range(upto - 1, -1, -1):
        i = ins[k]
        ops = i.op_str or ""
        # a compare against an immediate, on any register in the chain
        if i.mnemonic.startswith(("cmp", "cmn", "tst", "subs", "cbz", "cbnz")):
            for r in chain:
                if re.search(r"\b%s\b" % re.escape(r), ops):
                    return True
        # masking or shifting the value down also bounds it (e.g. `and r0,#0x3f`,
        # or GCP's trick of requiring `arg10 >> 30 == 0` and then using byte 3)
        if i.mnemonic.startswith(("and", "ubfx", "uxtb", "lsr", "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(1) in chain and i.mnemonic.startswith("mov"):
            chain.add(m.group(2))
    return False


def from_load(ins: list, reg: str, upto: int) -> str | None:
    for k in range(upto - 1, -1, -1):
        i = ins[k]
        ops = i.op_str or ""
        m = re.match(r"(%s),\s*(.+)$" % re.escape(reg), ops)
        if not m:
            continue
        if i.mnemonic.startswith("ldr"):
            return "%s %s" % (i.mnemonic, m.group(2))
        if i.mnemonic.startswith("mov"):
            m2 = re.match(r"(r\d+|sb|sl|fp|ip)$", m.group(2).strip())
            if m2:
                return from_load(ins, m2.group(1), k)
        return None
    return None


def scan_body(b: bytes, fns: list, base: int):
    rows = []
    for f in fns:
        sz = f.get("size", 0)
        if not sz:
            continue
        off = f["ea"] - base
        if off < 0 or off + sz > len(b):
            continue
        ins = list(THUMB.disasm(b[off:off + sz], f["ea"]))
        for n, i in enumerate(ins):
            idx = None
            scale = None
            if i.mnemonic == "mla":
                m = SCALE_MLA.match(i.op_str or "")
                if m:
                    idx, scale = m.group(2), "mla*%s" % m.group(3)
            elif i.mnemonic.startswith("add"):
                m = SCALE_SHIFT.match(i.op_str or "")
                if m:
                    idx, scale = m.group(3), "lsl#%s" % m.group(4)
            if not idx:
                continue
            src = from_load(ins, idx, n)
            if src and not bounded(ins, idx, n):
                rows.append((f["name"], i.address, scale, idx, src))
    return rows


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

    if args.selftest:
        # positive: SYS_DRV 0x109a handler; negative: GCP FE0D (bounded via >>30 + byte)
        sd = (REPO / "psp_out" / "sys_drv_code.bin").read_bytes()
        pos = scan_body(sd, [{"ea": 0xE11A78, "size": 0x40, "name": "sub_E11A78"}], 0xE00100)
        gcp = json.loads((HERE / "syms" / "gcp.meta.json").read_text())
        gb = Path(gcp["path"])
        gb = gb if gb.is_absolute() else REPO / gb
        neg = scan_body(gb.read_bytes(),
                        [{"ea": 0x403170, "size": 100, "name": "gcp_cmdFE0D"}], 0x400100)
        print("SELFTEST")
        print("   positive  sub_E11A78 (SYS_DRV 0x109a, unchecked 12*index): %s"
              % ("FOUND" if pos else "MISSED"))
        print("   negative  gcp_cmdFE0D (index bounded by >>30 then byte):  %s"
              % ("clean" if not neg else "FALSE POSITIVE %s" % neg))
        ok = bool(pos) and not neg
        print("\nverdict: %s" % ("TRUSTWORTHY" if ok else "NOT trustworthy"))
        return 0 if ok else 1

    idx = json.loads((HERE / "syms" / "index.json").read_text())
    for ta in (sorted(idx) if args.all else [args.prefix]):
        meta = json.loads((HERE / "syms" / ("%s.meta.json" % ta)).read_text())
        body = Path(meta["path"])
        body = body if body.is_absolute() else REPO / body
        fns = json.loads((HERE / "libdumps" / ("%s.json" % ta)).read_text())
        rows = scan_body(body.read_bytes(), fns, 0x400100)
        print("\n=== %s: %d unbounded scaled-index site(s)" % (ta, len(rows)))
        for nm, ea, scale, reg, src in rows:
            print("   %-32s @%#x  %-10s idx=%-4s <- %s" % (nm[:32], ea, scale, reg, src[:40]))
    return 0


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