#!/usr/bin/env python3
"""Audit every derived function name against the firmware's own log strings.

The TAs log `AMD_TA_<X>: <FunctionName> ENTER` on entry to most command handlers. That
string is ground truth for what AMD calls the function, so it is a mechanical check on
names that were derived by hand -- and it catches drift that no amount of re-reading the
disassembly would.

Reports three classes:
  MISMATCH  the handler logs a name that our symbol does not contain
  OK        our symbol contains the logged name
  NOSTRING  no ENTER-style string in the function (nothing to check against)

Usage:  ./ta_name_audit.py [--prefix gcp] [--fix-json]
"""
from __future__ import annotations

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

import capstone as cs

HERE = Path(__file__).resolve().parent
BASE = 0x400100


def load(prefix: str):
    meta = json.loads((HERE / "syms" / ("%s.meta.json" % prefix)).read_text())
    body = Path(meta["path"])
    if not body.is_absolute():
        body = HERE.parent / body
    return json.loads((HERE / "libdumps" / ("%s.json" % prefix)).read_text()), body.read_bytes()


def cstr(b: bytes, addr: int) -> str | None:
    o = addr - BASE
    if not (0 <= o < len(b)):
        return None
    e = b.find(b"\x00", o)
    if e < 0 or e - o > 200:
        return None
    try:
        s = b[o:e].decode("ascii")
    except Exception:  # noqa: BLE001
        return None
    return s if s.isprintable() or "\n" in s else None


def strings_in(b: bytes, ea: int, size: int) -> list[str]:
    md = cs.Cs(cs.CS_ARCH_ARM, cs.CS_MODE_THUMB)
    out = []
    off = ea - BASE
    if off < 0 or size <= 0 or off + size > len(b):
        return out
    for i in md.disasm(b[off:off + size], ea):
        addr = None
        if i.mnemonic == "adr":
            m = re.search(r"#(-?0x[0-9a-f]+|-?\d+)", i.op_str)
            if m:
                v = m.group(1)
                addr = ((i.address + 4) & ~3) + (int(v, 16) if v.startswith(("0x", "-0x")) else int(v))
        elif i.mnemonic.startswith("ldr") and "pc" in i.op_str:
            m = re.search(r"#(0x[0-9a-f]+)", i.op_str)
            if m:
                a = ((i.address + 4) & ~3) + int(m.group(1), 16)
                try:
                    addr = struct.unpack_from("<I", b, a - BASE)[0]
                except Exception:  # noqa: BLE001
                    addr = None
        if addr is None:
            continue
        s = cstr(b, addr)
        if s and len(s) > 4:
            out.append(s)
    return out


NAME_RE = re.compile(r"AMD_TA_[A-Za-z0-9_]+:\s*([A-Za-z_][A-Za-z0-9_]*)")


def logged_name(strs: list[str]) -> str | None:
    """Prefer the ENTER string; fall back to the most common logged identifier."""
    for s in strs:
        m = NAME_RE.search(s)
        if m and "ENTER" in s:
            return m.group(1)
    counts: dict[str, int] = {}
    for s in strs:
        m = NAME_RE.search(s)
        if m:
            counts[m.group(1)] = counts.get(m.group(1), 0) + 1
    if not counts:
        return None
    return max(counts, key=lambda k: counts[k])


def norm(s: str) -> str:
    return re.sub(r"[^a-z0-9]", "", s.lower())


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--prefix")
    ap.add_argument("--json", default=None, help="write findings to this path")
    args = ap.parse_args()

    prefixes = [args.prefix] if args.prefix else sorted(
        p.name.split(".")[0] for p in (HERE / "libdumps").glob("*.json")
    )
    allrows = []
    for prefix in prefixes:
        try:
            fns, b = load(prefix)
        except Exception as exc:  # noqa: BLE001
            print("%-10s skipped (%s)" % (prefix, exc))
            continue
        mism, ok, nostr = [], 0, 0
        for f in fns:
            nm, ea, sz = f["name"], f["ea"], f.get("size", 0)
            if not sz:
                nostr += 1
                continue
            ln = logged_name(strings_in(b, ea, sz))
            if not ln:
                nostr += 1
                continue
            if norm(ln) in norm(nm):
                ok += 1
            else:
                mism.append({"prefix": prefix, "ea": hex(ea), "name": nm, "logged": ln})
        print("%-10s checked-against-strings: %3d OK, %3d MISMATCH, %3d no string"
              % (prefix, ok, len(mism), nostr))
        for m in mism:
            print("    MISMATCH %-12s %-46s logs %r" % (m["ea"], m["name"], m["logged"]))
        allrows += mism
    if args.json:
        Path(args.json).write_text(json.dumps(allrows, indent=2))
        print("\nwrote %s (%d mismatches)" % (args.json, len(allrows)))
    return 0


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