#!/usr/bin/env python3
"""Name TA functions from their own log strings.

GCP and WSPE log through strings of the form

    AMD_TA_GCP: AllocateDRMSession ENTER
    AMD_TA_WSPE: Request DOPP Active/Deactive

The leading identifier is the vendor's own name for the enclosing function, so a string
reference is direct evidence of that function's identity -- stronger than anything derived
by similarity or call-graph position. This pass finds the function that references each
such string and proposes that name.

Only GCP (72 distinct) and WSPE (7) carry these strings; the other 14 modules have none,
which was established by scanning every build binary for "AMD?TA" and finding nothing. ASD
in particular contains just 41 printable runs, all property keys like "gpd.ta.appID", so
there is nothing to harvest there.

Report mode by default. --apply writes names into fw/cgnames/<ta>.json where ta_re.py
picks them up as the IDA tier, so they survive database loss.
"""
import argparse, json, re, struct, sys
from pathlib import Path

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(Path.home() / "dev" / "ida-tui-maybe"))

# message -> function identity. The ENTER/EXIT suffix and any trailing detail are dropped;
# what remains is the leading CamelCase identifier.
LOGPAT = re.compile(r"^AMD[_ ]?TA[_ ]?([A-Za-z0-9]+)\s*:\s*(.*)$")
# Second form, with no AMD_TA prefix: POWERPLAY logs "CreateEntry: failed status = 0x%x" and
# "InvokeCommand: status = 0x%x". The identifier before the colon is the function name. Found
# only after re-checking the ceiling claim, which had rested entirely on the AMD_TA regex and
# so reported POWERPLAY as having no log strings.
BAREPAT = re.compile(r"^([A-Z][A-Za-z0-9_]{5,30})\s*:\s+(?=\S)(.*)$")
IDENT = re.compile(r"^([A-Z][A-Za-z0-9_]*(?:[A-Z][A-Za-z0-9_]*)*)")
ADDRY = re.compile(r"^(j_)?sub_[0-9A-Fa-f]+$")

# Messages whose leading word is a verb of the log itself, not a function name.
NOT_A_NAME = {"Request", "Error", "Fail", "Failed", "Warning", "Info", "Debug", "ERROR",
              "Received", "Running", "Entering", "Leaving", "Invalid", "Unknown"}

# A log name that is a short form of a standard entry point resolves to the canonical name.
# POWERPLAY logs "CreateEntry: status = 0x%x" from the function that five sibling builds name
# TA_CreateEntryPoint at the same address, so taking the log spelling verbatim would name the
# same function two different things across builds. The log confirms the identity; the
# canonical spelling states it.
ALIAS = {
    "CreateEntry": "TA_CreateEntryPoint",
    "InvokeCommand": "TA_InvokeCommandEntryPoint",
    "OpenSession": "TA_OpenSessionEntryPoint",
    "CloseSession": "TA_CloseSessionEntryPoint",
    "DestroyEntry": "TA_DestroyEntryPoint",
}

# Names IDA or the libdump already establishes authoritatively; a log line never overrides these.
NEVER_OVERRIDE = ("TA_CreateEntryPoint", "TA_DestroyEntryPoint", "TA_OpenSessionEntryPoint",
                  "TA_CloseSessionEntryPoint", "TA_InvokeCommandEntryPoint")


def base_name(msg: str) -> str | None:
    m = IDENT.match(msg.strip())
    if not m:
        return None
    ident = m.group(1)
    if ident in NOT_A_NAME or len(ident) < 6:
        return None
    # Require a multi-word CamelCase identifier: two or more capital-initiated segments.
    # Vendor function names all satisfy this (AllocateDRMSession, GenerateOMAC,
    # QueryOutputIDCount), while prose openers do not. Without this, WSPE's
    # "AMD_TA_WSPE: Received command..." yielded the name "Received" and would have
    # clobbered the correct TA_InvokeCommandEntryPoint -- a log line is evidence about a
    # function only when it actually contains a function name.
    if len(re.findall(r"[A-Z]", ident)) < 2 or ident.isupper():
        return None
    return ALIAS.get(ident, ident)


def harvest(elf: Path) -> dict[int, str]:
    """string address -> proposed function name, read straight from the binary."""
    raw = elf.read_bytes()
    ph = struct.unpack_from("<I", raw, 28)[0]
    off, va, _, fsz = struct.unpack_from("<IIII", raw, ph + 4)
    body = raw[off : off + fsz]
    out = {}
    # The log strings end "...ENTER\n\0", and \n is not in [ -~], so the terminator must
    # allow trailing newlines -- without this the harvest silently returns nothing.
    for m in re.finditer(rb"[ -~]{8,120}[\r\n]*\x00", body):
        s = m.group(0).rstrip(b"\x00").rstrip().decode("ascii", "replace")
        g = LOGPAT.match(s)
        if g:
            nm = base_name(g.group(2))
        else:
            g2 = BAREPAT.match(s)
            # The bare form names the function BEFORE the colon, so validate that identifier
            # directly rather than the message after it.
            nm = base_name(g2.group(1)) if g2 else None
        if nm:
            out[m.start() + va] = nm
    return out


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--type", required=True)
    ap.add_argument("--apply", action="store_true")
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--dump", help="write have->want pairs as JSON")
    args = ap.parse_args()
    ta = args.type

    from idatui.nexus_client import NexusClient
    from idatui import remote_ops

    elfs = sorted((HERE.parent / "psp_out/versions_win/elf_all").glob(f"{ta}_*.elf"))
    if args.limit:
        elfs = elfs[: args.limit]
    if not elfs:
        print("no ELFs for %s" % ta)
        return 1

    cgf = HERE / "cgnames" / f"{ta}.json"
    cg = json.loads(cgf.read_text()) if cgf.exists() else {}
    agree = conflict = fresh = nostr = 0
    conflicts = []

    for elf in elfs:
        strs = harvest(elf)
        if not strs:
            nostr += 1
            continue
        build = elf.name.replace(".elf", "")
        client = NexusClient(str(elf))
        try:
            client.connect(progress=None)
            qs = [{"addr": hex(a), "direction": "to", "xref_type": "data"} for a in strs]
            rows = client.call(remote_ops.xref_query, queries=qs)["result"]
            rows = rows if isinstance(rows, list) else []
            proposals = {}
            for q, r in zip(strs.items(), rows):
                sa, nm = q
                for x in (r.get("data") or []) if isinstance(r, dict) else []:
                    fn = x.get("fn")
                    if not fn:
                        continue
                    proposals.setdefault(fn["addr"].lower(), (nm, fn.get("name") or ""))
            for addr, (nm, cur) in sorted(proposals.items()):
                # Preserve a verified command id if the current name carries one: the id comes
                # from the dispatcher table (address-exact against fw/cmdmaps) and the log
                # string supplies the semantics, so the two are complementary.
                #
                # But NOT across a __hN suffix. That suffix means the function was treated as a
                # fragment of the preceding one, so the command id was never established for
                # this address -- it was inherited from a neighbour. The log string proves the
                # function is its own routine, which means the id is unverified for it and
                # carrying it over would manufacture a false dispatcher mapping. 0x4038d8 is
                # the case in point: labelled cmdFE0A__h9, its log string says
                # GetCryptoSessionByDecodeHandle, and FE0A belongs to a different function.
                if any(k in (cur or "") for k in NEVER_OVERRIDE):
                    agree += 1
                    continue
                cur_nohn = cur or ""
                cmd = (None if "__h" in cur_nohn
                       else re.match(r"^%s_(cmd[0-9A-Fa-f]+)" % ta, cur_nohn))
                want = "%s_%s_%s" % (ta, cmd.group(1), nm) if cmd else "%s_%s" % (ta, nm)
                if not cur or ADDRY.match(cur):
                    fresh += 1
                    if args.apply:
                        cg.setdefault(build, {})[addr] = want
                elif cur == want:
                    agree += 1
                elif nm.lower() in (cur or "").lower():
                    agree += 1
                else:
                    conflict += 1
                    if args.apply:
                        cg.setdefault(build, {})[addr] = want
                    if len(conflicts) < 500:
                        conflicts.append((build, addr, cur, want))
        finally:
            client.close()

    print("%s: %d builds, %d without log strings" % (ta, len(elfs), nostr))
    print("  log-string evidence AGREES with existing name : %d" % agree)
    print("  names a function that was UNNAMED             : %d" % fresh)
    print("  CONTRADICTS an existing name                  : %d" % conflict)
    for b, a, cur, want in conflicts:
        print("     %-28s %-10s have=%-42s log=%s" % (b, a, cur, want))
    if args.dump:
        pairs = {}
        bad = {}
        for b, a, cur, want in conflicts:
            if cur in pairs and pairs[cur] != want:
                bad.setdefault(cur, set()).add(want); bad[cur].add(pairs[cur])
            pairs[cur] = want
        Path(args.dump).write_text(json.dumps(
            {"pairs": pairs, "inconsistent": {k: sorted(v) for k, v in bad.items()}},
            indent=1, sort_keys=True))
        print("  %d distinct have->want pairs, %d inconsistent across builds" % (len(pairs), len(bad)))
    if args.apply:
        cgf.parent.mkdir(parents=True, exist_ok=True)
        cgf.write_text(json.dumps(cg, indent=1, sort_keys=True))
        print("  wrote %s" % cgf)
    return 0


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