#!/usr/bin/env python3
"""Decide each field-flow site: writer WIDTH x scaling vs destination capacity.

Pass twelve left SYS_DRV's field-flow sites "characterised as type-bounded narrow fields" --
a description, not a verdict. The deciding facts are available:

  * the field's maximum value follows from the WIDTH of the stores that write it
    (`strb` -> 255, `strh` -> 65535, `str` -> unknown/32-bit);
  * the scaling applied between the load and the copy (`lsl #n`, `add rD,rA,rA,lsl #n`);
  * the destination capacity when the copy targets a stack buffer (`sub sp,#F`, dst `sp+k`).

A site is only dangerous when `max_value x scaling > capacity`. That is exactly the arithmetic
that makes HDCP's byte field (255 x 5 = 1275 into 276 bytes) a bug while a byte field copied
into a 4 KB buffer is not.

CONTROL: hdcp `0x402cde` must be decided UNSAFE; xgmi `0x402d4c` must be decided safe.

Usage:
  ./ta_fieldcap_audit.py --selftest
  ./ta_fieldcap_audit.py --module sys_drv
"""
from __future__ import annotations

import argparse
import re
from collections import Counter

import ta_capacity_audit as CAP
import ta_fieldflow2 as FF


def store_widths(ins, fns, off):
    """Widths of every verified store to field +off: 1, 2 or 4 bytes."""
    w = Counter()
    for i in ins:
        if not i.mnemonic.startswith("str"):
            continue
        m = FF.STOREF.match(i.op_str or "")
        if not m or int(m.group(3), 0) != off or m.group(2) in ("sp", "fp"):
            continue
        if i.mnemonic.startswith("strb"):
            w[1] += 1
        elif i.mnemonic.startswith("strh"):
            w[2] += 1
        else:
            w[4] += 1
    return w


def analyse(name):
    b, base, fns = FF.load_module(name) if hasattr(FF, "load_module") else CAP.load(name)
    fns = CAP.split_merged(b, base, fns)
    hp = CAP.helpers(b, base, fns)
    rows = []
    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))
        frame = None
        for i in ins[:8]:
            m = re.match(r"sp,\s*(?:sp,\s*)?#(0x[0-9a-f]+|\d+)$", i.op_str or "")
            if m and i.mnemonic in ("sub", "subw", "sub.w"):
                frame = int(m.group(1), 0)
                break
        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 hp:
                continue
            off = FF.field_of(ins, 0, n, "r2")
            if off is None:
                continue
            mx, how = CAP.max_len(ins, 0, n, "r2")
            cap, koff = (CAP.capacity(ins, 0, n, frame) if frame else (None, None))
            rows.append((i.address, ea, off, mx, how, cap, koff))
    return rows


def widths_for(name, offs):
    b, base, fns = CAP.load(name)
    fns = CAP.split_merged(b, base, fns)
    allins = []
    for ea, sz in fns:
        o = ea - base
        if 0 <= o and o + sz <= len(b) and sz > 0:
            allins.extend(CAP.THUMB.disasm(b[o:o + sz], ea))
    return {off: store_widths(allins, fns, off) for off in offs}


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

    if a.selftest:
        hd = [r for r in analyse("hdcp") if r[0] == 0x402CDE]
        xg = [r for r in analyse("xgmi") if r[0] == 0x402D4C]
        ok = True
        print("SELFTEST (width x scaling vs capacity)")
        for tag, rows, want_unsafe in (("hdcp 0x402cde", hd, True), ("xgmi 0x402d4c", xg, False)):
            if not rows:
                print("   %-16s NOT FOUND -> untrustworthy" % tag)
                ok = False
                continue
            _s, _f, off, mx, how, cap, _k = rows[0]
            unsafe = (mx is not None and cap is not None and mx > cap)
            print("   %-16s max %s (%s) vs capacity %s -> %s  %s"
                  % (tag, mx, how, cap, "UNSAFE" if unsafe else "safe",
                     "OK" if unsafe == want_unsafe else "WRONG"))
            ok &= (unsafe == want_unsafe)
        print("\nverdict: %s" % ("TRUSTWORTHY" if ok else "NOT trustworthy"))
        return 0 if ok else 1

    rows = analyse(a.module)
    offs = sorted({r[2] for r in rows})
    wid = widths_for(a.module, offs)
    dec = Counter()
    unsafe = []
    print("=== %s: deciding %d field-flow copy site(s)" % (a.module, len(rows)))
    for site, fn, off, mx, how, cap, koff in rows:
        w = wid.get(off, Counter())
        maxw = max(w) if w else 4
        typemax = {1: 0xFF, 2: 0xFFFF, 4: None}[maxw]
        if cap is None:
            dec["no stack destination (capacity unknown)"] += 1
            continue
        if typemax is None:
            dec["32-bit field (width gives no bound)"] += 1
            unsafe.append((site, fn, off, "word", cap))
            continue
        if typemax > cap:
            dec["width %d-byte COULD exceed capacity" % maxw] += 1
            unsafe.append((site, fn, off, "%d-byte max %d" % (maxw, typemax), cap))
        else:
            dec["width %d-byte cannot exceed capacity %d" % (maxw, cap)] += 1
    for k, v in dec.most_common():
        print("   %-52s %d" % (k, v))
    print("\n   sites needing a human read: %d" % len(unsafe))
    for site, fn, off, why, cap in unsafe[:12]:
        print("      copy @%#x (fn %#x) field +%#x  %s  capacity %d" % (site, fn, off, why, cap))
    return 0


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