#!/usr/bin/env python3
"""Detector class 9: a length/count read from a buffer that is itself a copy DESTINATION.

This generalises the one real defect this hunt produced. In `sub_E04B68` the count at
`buf+0x88` is read and used as a copy length, while `buf` is the staging buffer that
`sub_E1D6BC` fills with message payload — so the "length" is attacker-plantable data. None of
the eight earlier detector classes model this: they check whether a length is bounded, not
whether the *memory the length lives in* is itself written by a copy.

Shape hunted:

    ldr  rL, [rB, #F]        ; a field of some object B
    ...
    <copy>(dst, src, rL)     ; that field used as a length
  and elsewhere
    <copy>(rB, src2, n)      ; B is the DESTINATION of some copy

Positive control: `sub_E04B68` @0xe04bae/0xe04bc4 with `[r1+0x88]`.
Negative control: `sub_E17334`, whose length is a stacked argument, not a buffer field.

Usage:
  ./ta_bufferfield_audit.py --selftest
  ./ta_bufferfield_audit.py --module sys_drv|asd|<ta>
"""
from __future__ import annotations

import argparse
import json
import re
from pathlib import Path

import ta_capacity_audit as CAP

HERE = Path(__file__).resolve().parent


def analyse(name: str):
    b, base, fns = CAP.load(name)
    fns = CAP.split_merged(b, base, fns)
    helpers = CAP.helpers(b, base, fns)
    out = []
    for ea, sz in fns:
        o = ea - base
        if o < 0 or sz <= 0 or o + sz > len(b):
            continue
        ins = list(CAP.THUMB.disasm(b[o:o + sz], ea))
        # 1. collect  ldr rL,[rB,#F]  -> (L, B, F, index)
        loads = {}
        for n, i in enumerate(ins):
            m = re.match(r"(\w+),\s*\[(\w+),\s*#(0x[0-9a-f]+|\d+)\]$", i.op_str or "")
            if i.mnemonic.startswith("ldr") and m and m.group(2) != "sp":
                loads[m.group(1)] = (m.group(2), int(m.group(3), 0), n)
        # 2. find copies whose length register came from such a load
        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
            # which register feeds r2 (the length)?
            cur, at = "r2", n
            src = None
            for _ in range(5):
                for k in range(at - 1, -1, -1):
                    j = ins[k]
                    mm = re.match(r"(%s),\s*(.+)$" % re.escape(cur), j.op_str or "")
                    # cbz/cbnz READ the register; treating them as writes truncated the slice
                    if not mm or j.mnemonic.startswith(("cmp", "cb", "str", "push", "tst", "b")):
                        continue
                    if j.mnemonic.startswith("ldr"):
                        m2 = re.match(r"\w+,\s*\[(\w+),\s*#(0x[0-9a-f]+|\d+)\]$", j.op_str or "")
                        if m2 and m2.group(1) != "sp":
                            src = (m2.group(1), int(m2.group(2), 0), k)
                        break
                    m3 = re.fullmatch(r"(r\d+|sb|sl|fp|ip)", mm.group(2).strip())
                    if m3:
                        cur, at = m3.group(1), k
                        break
                    break
                else:
                    break
                if src:
                    break
            if not src:
                continue
            basereg, field, _ = src
            # 3. is that same base register ever a copy DESTINATION in this function?
            dest_here = False
            for k, j in enumerate(ins):
                if not (j.mnemonic in ("blx", "bl") and (j.op_str or "").startswith("#")):
                    continue
                if int(j.op_str[1:], 16) not in helpers:
                    continue
                for q in range(max(0, k - 5), k):
                    if re.fullmatch(r"r0,\s*%s" % re.escape(basereg), ins[q].op_str or "") \
                       and ins[q].mnemonic.startswith("mov"):
                        dest_here = True
            out.append((ea, i.address, basereg, field, dest_here))
    return out


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

    if a.selftest:
        rows = analyse("sys_drv")
        pos = [r for r in rows if r[1] in (0xE04BAE, 0xE04BC4) and r[3] == 0x88]
        neg = [r for r in rows if r[0] == 0xE17334]
        print("SELFTEST (length read from a buffer field)")
        print("   POSITIVE sub_E04B68 buf+0x88 used as a copy length : %s"
              % ("FOUND" if pos else "MISSED"))
        print("   NEGATIVE sub_E17334 (length is a stacked argument)  : %s"
              % ("clean" if not neg else "FALSE POSITIVE"))
        ok = bool(pos) and not neg
        print("\nverdict: %s" % ("TRUSTWORTHY" if ok else "NOT trustworthy"))
        return 0 if ok else 1

    rows = analyse(a.module)
    print("=== %s: %d copy site(s) whose length is read from an object field" % (a.module, len(rows)))
    hot = [r for r in rows if r[4]]
    print("    of which the same object is also a copy destination: %d\n" % len(hot))
    for ea, site, breg, field, dest in hot[:20]:
        print("   fn %#-9x copy @%#-9x len <- [%s+%#x]   (object is copy-filled)" % (ea, site, breg, field))
    return 0


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