#!/usr/bin/env python3
"""Rename GCP functions to the names the firmware itself logs, and fix command tags.

The TAs log `AMD_TA_GCP: <Name> ENTER`, which is ground truth for AMD's own function name.
Hand-derived names drifted from it in 18 places, and the command tags attached to them
drifted too, because the earlier cmdmap recorded names rather than addresses and collapsed
several dispatch arms.

This re-derives everything from two authoritative sources:

  * cmd -> handler address, by decoding the dispatcher's two `tbb` tables directly
    (tier 1 at 0x404002 for 0xFE02..0xFE11, tier 2 at 0x404026 for cmd 1..5);
  * function -> AMD name, from the logged ENTER string inside each function.

Findings this corrects, all of which the hardware sweep independently agrees with:
  * 0xFE03..0xFE06 are UNIMPLEMENTED -- all four table arms point at the error path
    0x4040ce. Hardware returned an identical 0x102 for all four.
  * 0xFE07 is CreateAuthenticatedSession (0x402e94), previously tagged FE03. Hardware
    returned 0x0 with a handle for FE07 only.
  * 0xFE0C/0xFE0D are Create/DestroyCryptoSession, previously tagged ...AuthenticatedSession.
  * 0xFE0A is QueryAuthenticatedSession; "QueryAccessibilityAttributes" was invented.
  * cmd 3 is 0x4025f8, which IDA never defined as a function.

Usage:  IDATUI_RPC_SOCK=... ./ta_rename_from_logstrings.py [--apply]
"""
from __future__ import annotations

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

import capstone as cs

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(Path.home() / "dev" / "ida-tui-maybe"))
BASE = 0x400100
BODY = HERE.parent / "psp_out" / "versions_win" / "ta_22-0-4-50_81b082b4.bin"

TIER = [(0x404002, 16, 0xFE02), (0x404026, 5, 1)]
ERROR_ARM = 0x4040CE


def dispatch_map(b: bytes) -> dict[int, int | None]:
    md = cs.Cs(cs.CS_ARCH_ARM, cs.CS_MODE_THUMB)
    out: dict[int, int | None] = {}
    for base, n, first in TIER:
        for k in range(n):
            tgt = base + 2 * b[base + k - BASE]
            if tgt == ERROR_ARM:
                out[first + k] = None  # unimplemented: jumps to the shared error path
                continue
            h = None
            for i in md.disasm(b[tgt - BASE:tgt - BASE + 0x10], tgt):
                if i.mnemonic == "bl":
                    h = int(i.op_str.replace("#", ""), 16)
                    break
            out[first + k] = h
    return out


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

    b = BODY.read_bytes()
    cmds = dispatch_map(b)
    logged = {
        int(r["ea"], 16): r["logged"]
        for r in json.loads((HERE / "audits" / "name_vs_logstring.json").read_text())
        if r["prefix"] == "gcp"
    }
    fns = {f["ea"]: f["name"] for f in json.loads((HERE / "libdumps" / "gcp.json").read_text())}

    cmd_of = {ea: c for c, ea in cmds.items() if ea}
    renames: dict[int, str] = {}

    # every function with a logged name gets that name, tagged with its command if the
    # dispatch table says it is a command handler
    for ea, name in sorted(logged.items()):
        # the logged string may sit a few bytes into the function; map it to the owner
        owner = max((f for f in fns if f <= ea), default=None)
        if owner is None:
            continue
        cmd = cmd_of.get(owner)
        new = ("gcp_cmd%02X_%s" % (cmd, name)) if cmd and cmd < 0x100 else (
            "gcp_cmd%04X_%s" % (cmd, name) if cmd else "gcp_%s" % name)
        renames[owner] = new

    # command handlers whose existing name already matches their logged name still need
    # their TAG corrected when the dispatch table disagrees with it
    for cmd, ea in sorted(cmds.items()):
        if not ea or ea in renames:
            continue
        cur = fns.get(ea)
        if not cur:
            continue
        m = re.match(r"^gcp_cmd([0-9A-Fa-f]{2,4})_(.*)$", cur)
        if not m:
            continue
        tag = int(m.group(1), 16)
        if tag != cmd:
            new = ("gcp_cmd%02X_%s" % (cmd, m.group(2))) if cmd < 0x100 else (
                "gcp_cmd%04X_%s" % (cmd, m.group(2)))
            renames[ea] = new

    print("dispatch table (authoritative):")
    for cmd, ea in sorted(cmds.items()):
        tag = "UNIMPLEMENTED (error arm)" if ea is None else hex(ea)
        print("   cmd %#06x -> %-12s %s" % (cmd, tag, renames.get(ea or 0, fns.get(ea or 0, ""))))

    print("\n%d renames:" % len(renames))
    for ea, new in sorted(renames.items()):
        print("   %#x  %-46s -> %s" % (ea, fns.get(ea, "<undefined>"), new))

    if not args.apply:
        return 0

    from idatui.rpcclient import RpcClient

    sock = os.environ.get("IDATUI_RPC_SOCK")
    if not sock:
        print("IDATUI_RPC_SOCK not set", file=sys.stderr)
        return 2
    c = RpcClient(sock, timeout=180)
    c.connect()

    # cmd 3's handler was never defined as a function by IDA
    missing = [ea for ea in cmds.values() if ea and ea not in fns]
    for ea in missing:
        try:
            r = c.call("define_func", target=hex(ea))
            print("define_func %#x -> %s" % (ea, json.dumps(r)[:120]))
        except Exception as exc:  # noqa: BLE001
            print("define_func %#x FAILED: %s" % (ea, exc))

    edits = [{"addr": hex(ea), "name": nm} for ea, nm in sorted(renames.items())]
    for ea in missing:
        nm = "gcp_cmd%02X_%s" % (cmd_of.get(ea, 0), logged.get(ea, "handler_%X" % ea))
        edits.append({"addr": hex(ea), "name": nm})
    # rename_many takes items=[{addr,name}] and renames BY ADDRESS. The single-token
    # `rename` op resolves against the cursor/current view instead, which silently renames
    # the wrong thing -- always use rename_many here.
    r = c.call("rename_many", items=edits, batch={"allow_overwrite": True})
    summary = r.get("rename_many", r) if isinstance(r, dict) else r
    print("\nrename_many: %s" % json.dumps(summary)[:600])
    try:
        c.call("save")
    except Exception as exc:  # noqa: BLE001
        print("save failed: %s" % exc)
    return 0


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