#!/usr/bin/env python3
"""Restore function signatures and structure types in the TA databases.

Naming alone leaves Hex-Rays emitting `*(_DWORD *)(a1 + 2048)` everywhere, which is
unreadable for exactly the code we care about (session lifetime handling). This applies:

  1. local struct types, built ONLY from layout facts already established in
     fw/AUDIT-NEW-TAS-RE.md and the finding docs (every offset here is one we proved
     from disassembly or hardware, never a guess);
  2. function prototypes, so handlers take named, typed parameters;
  3. a before/after pseudocode capture, so the improvement is measured rather than
     asserted.

The readability metric is deliberately mechanical: count the raw pointer-arithmetic casts
(`*(_DWORD *)(a1 + N)` and friends) and the anonymous `a1/a2/v3` identifiers, before and
after. A type that does not reduce those is not helping and should not be claimed.

Usage:
  IDATUI_RPC_SOCK=... ./ta_types.py --prefix gcp [--apply] [--sample N]
"""
from __future__ import annotations

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

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

sys.path.insert(0, str(Path.home() / "dev" / "ida-tui-maybe"))

# --------------------------------------------------------------------------------------
# Structure layouts. Every offset below is cited to where it was established.
# --------------------------------------------------------------------------------------

# GCP context. The three session arrays tile the context exactly, which is itself the
# check that the layout is right:
#     32 * 24 = 0x300   -> auth array starts at 0x300
#     16 * 80 = 0x500   -> 0x300 + 0x500 = 0x800, crypto array starts at 0x800
#     64 * 88 = 0x1600  -> 0x800 + 0x1600 = 0x1E00, the refcount, exactly where
#                          gcp_cmdFE0C computes it (add.w r6, r6, #0x1e00)
# Crypto slot fields are from the FE0C stores disassembled at 0x402fc8..0x402fee and the
# FE0F comparison at slot+0x18 (FINDING-GCP-FE0F-STALE-SLOT-LOOKUP.md).
GCP_TYPES = [
    """
    struct gcp_drm_session          // 24 bytes; stride proved by cmd02: add.w r0,r6,r6,lsl#1 ; lsl #3
    {
      unsigned int handle;          // +0x00
      unsigned int state;           // +0x04  cmd02 requires == 1
      unsigned char body[16];       // +0x08
    };
    """,
    """
    struct gcp_auth_session         // 80 bytes; index = handle >> 28
    {
      unsigned int handle;          // +0x00
      unsigned int state;           // +0x04
      unsigned char body[72];       // +0x08
    };
    """,
    """
    struct gcp_crypto_session       // 88 bytes; index = handle >> 30, slot idx in bits 24+
    {
      unsigned int handle;          // +0x00  str r0,[r5]      FE0C
      unsigned int unk04;           // +0x04
      unsigned long long key;       // +0x08  strd from params+0x10
      unsigned long long field10;   // +0x10  strd from params+0x18
      unsigned long long decode_handle; // +0x18  FE0F memcmp target -- NEVER WRITTEN by FE0C
      unsigned int state;           // +0x20  str r0(=1),[r5,#0x20]
      unsigned char rest[52];       // +0x24
    };
    """,
    """
    struct gcp_ctx
    {
      struct gcp_drm_session drm[32];        // +0x0000
      struct gcp_auth_session auth[16];      // +0x0300
      struct gcp_crypto_session crypto[64];  // +0x0800
      unsigned int asd_refcount;             // +0x1E00  incremented before the acquire
      unsigned int unk1E04;                  // +0x1E04
      unsigned int handle_counter;           // +0x1E08  monotonic, rebinds reused slots
    };
    """,
    # Host request page. Offsets are the ones handlers actually touch; the rest stays an
    # opaque tail rather than being invented. NOTE: if this struct is not declared, IDA
    # forward-declares an EMPTY ta_params when a prototype references it, and pseudocode
    # keeps showing *((_QWORD *)params + 2) -- acceptance of the prototype is therefore
    # NOT proof that the type has members.
    """
    struct ta_params
    {
      unsigned int cmd_id;          // +0x0000
      unsigned int unk04;
      unsigned int unk08;
      unsigned int unk0C;
      unsigned int arg10;           // +0x0010  session index / handle / query word
      unsigned char data14[0x1926]; // +0x0014  cmd02 certificate source, constant length
      unsigned char tail[0x6C];
      unsigned int out1948;         // +0x1948  cmd02 writes the copied length here
      unsigned int unk194C;
      unsigned int out1950;         // +0x1950  FE0C writes the new handle here
    };
    """,
]

# Per-TA context struct name; TAs other than GCP get a generic opaque context, because
# their internal layouts were not established to the same standard. Claiming a detailed
# struct for them would be invention.
CTX_TYPE = {"gcp": "struct gcp_ctx *"}

GENERIC_TYPES = [
    """
    struct ta_ctx_opaque { unsigned char raw[1]; };
    """,
]

# MERIT's slot array IS established: 65 slots of stride 257, base off_401CE0, and
# TA_CreateEntryPoint zero-fills 0x4141 = 257 * 65 (fw/AUDIT-NEW-TAS-RE.md).
MERIT_TYPES = [
    """
    struct merit_slot { unsigned char body[257]; };
    """,
    """
    struct merit_table { struct merit_slot slot[65]; };
    """,
]


def proto_for(name: str, ctx: str) -> str | None:
    """Pick a prototype from the function's established role. Returns None to skip."""
    # Command handlers: dispatched with r0 = context, r1 = host params. Verified on
    # gcp_cmd02 (mov r4,r0 / mov r7,r1) and gcp_cmdFE0C (ctx+0x800, r4 = params).
    if re.search(r"_cmd[0-9A-Fa-f]{2,4}(_|$)", name):
        return "int %s(%sctx, struct ta_params *params)" % (name, ctx)
    # ASD RPC thunks take a single word argument (movs r0,#1 ; bl thunk).
    m = re.match(r"^(\w+)_asd_rpc_([0-9a-fA-F]+)$", name)
    if m:
        return "int %s(unsigned int arg)" % name
    if name.endswith("_log") or re.search(r"_log$", name):
        return "int %s(int level, const char *fmt)" % name
    if "zero_fill" in name:
        return "void *%s(void *dst, unsigned int len)" % name
    if "memcmp" in name:
        return "int %s(const void *a, const void *b, unsigned int len)" % name
    if re.search(r"_mac_over_host_data$", name):
        return "int %s(%sctx, struct ta_params *params)" % (name, ctx)
    if name in ("TA_CreateEntryPoint", "TA_DestroyEntryPoint"):
        return "int %s(void)" % name
    if name == "TA_OpenSessionEntryPoint":
        return "int %s(unsigned int argc, void *argv)" % name
    if name == "TA_InvokeCommandEntryPoint":
        return ("int %s(void *sess, unsigned int cmd_id, unsigned int argc, "
                "struct ta_params *params)" % name)
    return None


def pseudo(c, target: str) -> str:
    """Fetch decompiled text. The RPC method is `pseudocode` (remote_ops.decompile is an
    internal worker op and is NOT exposed as an RPC method -- calling it raises
    "unknown method")."""
    try:
        r = c.call("pseudocode", target=target)
    except Exception as exc:  # noqa: BLE001
        return "<pseudocode failed: %s>" % exc
    if isinstance(r, str):
        return r
    if isinstance(r, dict):
        for k in ("text", "pseudocode", "body", "lines"):
            v = r.get(k)
            if isinstance(v, str):
                return v
            if isinstance(v, list):
                return "\n".join(
                    x if isinstance(x, str) else str(x.get("text", x)) for x in v
                )
    if isinstance(r, list):
        return "\n".join(x if isinstance(x, str) else str(x.get("text", x)) for x in r)
    return json.dumps(r)[:4000]


CAST_RE = re.compile(r"\*\((?:_[A-Z]+|unsigned \w+|\w+) \*\)\((?:a\d|v\d)\s*\+")
ANON_RE = re.compile(r"\b(?:a[1-9]\d?)\b")


def score(text: str) -> dict:
    return {
        "raw_casts": len(CAST_RE.findall(text)),
        "anon_args": len(set(ANON_RE.findall(text))),
        "lines": len(text.splitlines()),
    }


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--prefix", required=True)
    ap.add_argument("--apply", action="store_true")
    ap.add_argument("--sample", type=int, default=6)
    args = ap.parse_args()

    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()

    dump = json.loads((HERE / "libdumps" / ("%s.json" % args.prefix)).read_text())
    fns = {f["name"]: f for f in dump}

    ctx = CTX_TYPE.get(args.prefix, "struct ta_ctx_opaque *")
    decls = list(GENERIC_TYPES)
    if args.prefix == "gcp":
        decls += GCP_TYPES
    if args.prefix == "merit":
        decls += MERIT_TYPES

    # sample the most interesting functions for the before/after measurement
    interesting = [n for n in fns if re.search(r"_cmd[0-9A-Fa-f]{2,4}", n)][: args.sample]
    before = {}
    for n in interesting:
        try:
            before[n] = pseudo(c, hex(fns[n]["ea"]))
        except Exception as exc:  # noqa: BLE001
            before[n] = "<decompile failed: %s>" % exc

    if args.apply:
        # `declare_type` is NOT an exposed RPC method -- it lives only behind the TUI's
        # struct editor. Drive that: structs -> ctrl+n -> type the decl -> ctrl+s, then
        # close the modal (it swallows keystrokes for every other verb while open).
        c.call("structs")
        for d in decls:
            one = " ".join(d.split())
            c.call("keys", keys=["ctrl+n"])
            c.call("text", text=one)
            c.call("keys", keys=["ctrl+s"])
        c.call("keys", keys=["escape"])
        try:
            c.call("close")
        except Exception:  # noqa: BLE001
            pass
        print("declared %d struct types via the struct editor" % len(decls))

    applied, failed, skipped = [], [], 0
    for name, f in sorted(fns.items(), key=lambda kv: kv[1]["ea"]):
        p = proto_for(name, ctx)
        if not p:
            skipped += 1
            continue
        if not args.apply:
            applied.append((name, p))
            continue
        # `retype` resolves the token in the CURRENT view, so navigate by NAME first
        # (goto by address leaves the cursor on the address column and retype then fails
        # with "token not found in the current view").
        try:
            c.call("goto", target=name, delay_ms=0)
            c.call("retype", proto=p, delay_ms=0)
            applied.append((name, p))
        except Exception as exc:  # noqa: BLE001
            failed.append((name, str(exc)))

    print("\nprototypes: %d applied, %d failed, %d skipped (no established role)"
          % (len(applied), len(failed), skipped))
    for n, e in failed[:10]:
        print("   FAILED %-42s %s" % (n, e))

    if args.apply:
        tot_b = {"raw_casts": 0, "anon_args": 0}
        tot_a = {"raw_casts": 0, "anon_args": 0}
        print("\nbefore/after readability on %d sampled handlers:" % len(before))
        for n in before:
            try:
                after = pseudo(c, hex(fns[n]["ea"]))
            except Exception as exc:  # noqa: BLE001
                after = "<failed: %s>" % exc
            sb, sa = score(before[n]), score(after)
            for k in tot_b:
                tot_b[k] += sb[k]
                tot_a[k] += sa[k]
            print("   %-40s casts %2d -> %2d   anon args %d -> %d"
                  % (n, sb["raw_casts"], sa["raw_casts"], sb["anon_args"], sa["anon_args"]))
            outdir = REPO / "fw" / "pseudo"
            outdir.mkdir(exist_ok=True)
            (outdir / ("%s.%s.before.c" % (args.prefix, n))).write_text(before[n])
            (outdir / ("%s.%s.after.c" % (args.prefix, n))).write_text(after)
        print("   TOTAL raw casts %d -> %d ; anon args %d -> %d"
              % (tot_b["raw_casts"], tot_a["raw_casts"],
                 tot_b["anon_args"], tot_a["anon_args"]))
        try:
            c.call("save")
        except Exception as exc:  # noqa: BLE001
            print("save failed: %s" % exc)
    return 0


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