#!/usr/bin/env python3
"""Harvest struct-field evidence for a component's context and params pointers.

Only seven PSP structs exist in this tree (fw/ta_ida.py), all for GCP and MERIT, so thirteen
of sixteen components decompile with every pointer typed as `int`. Defining those structs by
hand needs evidence, and the evidence is the set of offsets the code actually touches through
each pointer, with the WIDTH of each access.

Rather than re-implement register dataflow, this reads Hex-Rays output: the decompiler has
already resolved which register held an argument, and it writes accesses in forms that state
the offset and the width directly:

    *(_DWORD *)(a1 + 0x300)     -> offset 0x300, width 4
    *((_BYTE *)a2 + 36)         -> offset 0x24,  width 1
    a1[12]                      -> offset 0x30,  width 4 (pointer arithmetic, scaled)
    ctx->field                  -> already typed; reported so it is not double-counted

WIDTH matters as much as offset: it is what distinguishes a byte field (max 255) from a dword,
which is the same distinction the length audits turn on.

## The control

Run with --validate on GCP. Its five structs are already proven
(fw/AUDIT-NEW-TAS-RE.md: gcp_ctx size 0x1e10, drm[32] @+0x0, auth[16] @+0x300,
crypto[64] @+0x800), so the harvester must rediscover those three array bases from the code
alone. If it cannot, its output for the undocumented components is not worth reading.

usage:
  ./ta_struct_re.py --type gcp --validate
  ./ta_struct_re.py --type asd --arg 2 [--min-hits 2] [--json out.json]
"""
from __future__ import annotations

import argparse
import glob
import json
import os
import re
import struct
import sys
from collections import defaultdict
from pathlib import Path

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

ELF_DIR = REPO / "psp_out" / "versions_win" / "elf_all"

WIDTH = {"_BYTE": 1, "_WORD": 2, "_DWORD": 4, "_QWORD": 8,
         "char": 1, "__int16": 2, "int": 4, "__int64": 8,
         "unsigned __int8": 1, "unsigned __int16": 2,
         "unsigned __int32": 4, "unsigned __int64": 8}


def access_patterns(var: str):
    """Regexes that expose (offset, width) for accesses through `var`."""
    v = re.escape(var)
    return [
        # *(_DWORD *)(a1 + 0x300)   /   *(_DWORD *)(a1 + 12)
        (re.compile(r"\*\(\s*(\w[\w ]*?)\s*\*\s*\)\s*\(\s*%s\s*\+\s*(0x[0-9A-Fa-f]+|\d+)\s*\)" % v), "cast"),
        # *((_BYTE *)a1 + 36)
        (re.compile(r"\*\(\(\s*(\w[\w ]*?)\s*\*\s*\)\s*%s\s*\+\s*(0x[0-9A-Fa-f]+|\d+)\s*\)" % v), "cast2"),
        # a1[12]  -> scaled by the declared element width; recorded separately
        (re.compile(r"\b%s\s*\[\s*(0x[0-9A-Fa-f]+|\d+)\s*\]" % v), "index"),
    ]


ALIAS_RE = re.compile(r"^\s*(\w+)\s*=\s*([^;]*)\s*;\s*$", re.M)


def aliases(code: str, var: str):
    """locals derived from `var` -> (additive_offset, was_scaled).

    Array bases are NOT written against the argument. GCP computes
    'v6 = a1 + 80 * HIBYTE(v4)' and then reads '*(_DWORD *)(v6 + 768)', so the base 0x300
    only ever appears as an offset on the derived pointer. Scanning accesses through the
    argument alone therefore cannot see any array base -- which is precisely why the first
    two control runs failed to rediscover gcp_ctx's auth[] and crypto[] arrays.

    A scaled derivation (one containing a multiply or index) means the literal added later is
    an ARRAY BASE; an unscaled one means it is just a field offset carried through.
    """
    v = re.escape(var)
    out: dict[str, tuple[int, bool]] = {}
    for m in ALIAS_RE.finditer(code):
        lhs, rhs = m.group(1), m.group(2)
        if lhs == var or not re.search(r"\b%s\b" % v, rhs):
            continue
        scaled = bool(re.search(r"[*]|\[", rhs))
        lits = [int(x, 0) for x in re.findall(r"0x[0-9A-Fa-f]+|\b\d+\b", rhs)]
        add = 0 if scaled else sum(lits)
        out[lhs] = (add, scaled)
    return out


def indexed_bases(code: str, var: str):
    """Array bases reached through a computed index: *(_DWORD *)(a1 + 80 * i + 768).

    The plain 'a1 + <literal>' patterns cannot see these, and they are exactly where the
    interesting structure lives -- an array base is the thing a struct definition has to get
    right. For each dereference through `var`, every integer literal in the offset expression
    is collected and the LARGEST is taken as the base, the smaller ones being element strides
    or field offsets within the element. Heuristic, and reported as kind='indexed' so it is
    never confused with a directly observed offset.
    """
    v = re.escape(var)
    rx = re.compile(r"\*\(\s*(\w[\w ]*?)\s*\*\s*\)\s*\(\s*%s\s*\+\s*([^()]+?)\s*\)" % v)
    out = []
    for m in rx.finditer(code):
        w = WIDTH.get(m.group(1).strip(), 0)
        expr = m.group(2)
        if not w or not re.search(r"[*\[]", expr):
            continue        # no scaling present -> the plain patterns already handled it
        lits = [int(x, 0) for x in re.findall(r"0x[0-9A-Fa-f]+|\b\d+\b", expr)]
        if not lits:
            continue
        base = max(lits)
        if base >= 0x10:    # a bare stride like 80 is not a base
            out.append((base, w))
    return out


def decl_width(code: str, var: str) -> int:
    """Element width implied by the variable's own declaration, for a1[i] indexing."""
    m = re.search(r"^\s*([\w ]+?)\s*\*+\s*%s\s*[;,)]" % re.escape(var), code, re.M)
    if not m:
        m = re.search(r"\(\s*([\w ]+?)\s*\*+\s*%s\s*[,)]" % re.escape(var), code)
    if m:
        return WIDTH.get(m.group(1).strip(), 4)
    return 4


def harvest(client, remote_ops, funcs, argno: int, min_size: int = 0):
    """-> {offset: {'widths': {w: n}, 'sites': [names], 'kinds': set}}"""
    ev: dict[int, dict] = defaultdict(
        lambda: {"widths": defaultdict(int), "sites": [], "kinds": set()})
    typed_hits = 0
    for fn in funcs:
        try:
            r = client.call(remote_ops.decompile, addr=fn["addr"],
                            include_addresses=False)
            code = (r or {}).get("code") or ""
        except Exception:
            continue
        if not code:
            continue
        head = code.split("\n", 1)[0]
        # positional argument name as Hex-Rays wrote it: a1, a2, ... or a named param
        params = re.findall(r"[\w \*]*?\*?\s*(\w+)\s*(?:,|\))", head[head.find("(") + 1:]) \
            if "(" in head else []
        if len(params) < argno:
            continue
        var = params[argno - 1]
        if not re.match(r"^[A-Za-z_]\w*$", var):
            continue
        if "->" in code and re.search(r"\b%s\s*->" % re.escape(var), code):
            # Already a struct pointer. Skipping is right for harvesting NEW layout, but it
            # also means a database that already has the structs applied cannot validate the
            # harvester -- the first GCP control run was circular for exactly this reason and
            # has to be done on a freshly analysed, untyped build.
            typed_hits += 1
            continue
        ew = decl_width(code, var)
        for rx, kind in access_patterns(var):
            for m in rx.finditer(code):
                if kind == "index":
                    off = int(m.group(1), 0) * ew
                    w = ew
                else:
                    w = WIDTH.get(m.group(1).strip(), 0)
                    off = int(m.group(2), 0)
                    if not w:
                        continue
                e = ev[off]
                e["widths"][w] += 1
                e["kinds"].add(kind)
                if fn["name"] not in e["sites"]:
                    e["sites"].append(fn["name"])
        for base, w in indexed_bases(code, var):
            e = ev[base]
            e["widths"][w] += 1
            e["kinds"].add("indexed")
            if fn["name"] not in e["sites"]:
                e["sites"].append(fn["name"])
        # one level of pointer derivation, which is where array bases actually show up
        for alias, (add, scaled) in aliases(code, var).items():
            for rx, kind in access_patterns(alias):
                for m in rx.finditer(code):
                    if kind == "index":
                        continue
                    w = WIDTH.get(m.group(1).strip(), 0)
                    if not w:
                        continue
                    off = int(m.group(2), 0) + add
                    e = ev[off]
                    e["widths"][w] += 1
                    e["kinds"].add("array_base" if scaled else "alias")
                    if fn["name"] not in e["sites"]:
                        e["sites"].append(fn["name"])
    return ev, typed_hits


def report(ev, min_hits: int, limit: int = 60):
    rows = []
    for off, e in sorted(ev.items()):
        n = sum(e["widths"].values())
        if n < min_hits:
            continue
        rows.append((off, n, dict(e["widths"]), e["sites"][:3]))
    print("  %-10s %-6s %-18s %s" % ("offset", "hits", "widths", "example sites"))
    for off, n, w, sites in rows[:limit]:
        print("  %-10s %-6d %-18s %s"
              % (hex(off), n, ",".join("%dB:%d" % (k, v) for k, v in sorted(w.items())),
                 ", ".join(s[:26] for s in sites)))
    if len(rows) > limit:
        print("  ... %d more offsets" % (len(rows) - limit))
    return rows


# Control, from the proofs in fw/AUDIT-NEW-TAS-RE.md ("gcp_ctx size 0x1e10  drm[32] @+0x0
# auth[16] @+0x300  crypto[64] @+0x800") and the refcount finding (ctx+0x1E00).
#
# Demanding the literal base offsets was the wrong test and failed a working harvester twice.
# The index is folded into the derived pointer ('v6 = a1 + 80 * i' then '*(v6 + 0x318)'), so
# what is observable is base + FIELD, and the bare base only shows up when the field offset
# happens to be zero. The honest requirement is that each known array region is reached, and
# that every offset seen inside it lies within one element -- which is what folding the index
# into the alias implies.
GCP_REGIONS = [("drm[32]", 0x0, 0x18), ("auth[16]", 0x300, 80),
               ("crypto[64]", 0x800, 88)]
GCP_EXACT = [0x1E00]


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--type", required=True)
    ap.add_argument("--build", help="version substring; default the reference build")
    ap.add_argument("--arg", type=int, default=1, help="1 = ctx, 2 = params")
    ap.add_argument("--min-hits", type=int, default=1)
    ap.add_argument("--validate", action="store_true",
                    help="GCP only: require the known array bases to be rediscovered")
    ap.add_argument("--json")
    a = ap.parse_args()

    from idatui import remote_ops
    from idatui.nexus_client import NexusClient

    pat = f"{a.type}_*{a.build}*.elf" if a.build else f"{a.type}_*.elf"
    cands = sorted(glob.glob(str(ELF_DIR / pat)))
    if not cands:
        raise SystemExit("no ELF matching %s" % pat)
    path = cands[0]
    print("harvesting %s arg%d from %s" % (a.type, a.arg, os.path.basename(path)))
    client = NexusClient(path)
    client.connect()
    funcs, off = [], 0
    while True:
        r = client.call(remote_ops.list_funcs, queries=[{"count": 500, "offset": off}])
        blk = r["result"][0]
        funcs.extend(blk["data"])
        off = blk["next_offset"]
        if off >= blk["total"]:
            break
    ev, typed = harvest(client, remote_ops, funcs, a.arg)
    print("  %d functions scanned, %d already struct-typed (skipped), %d distinct offsets"
          % (len(funcs), typed, len(ev)))
    rows = report(ev, a.min_hits)

    rc = 0
    if a.validate:
        found = sorted(off for off, *_ in rows)
        print()
        bad = []
        for label, base, stride in GCP_REGIONS:
            if label == "drm[32]":
                continue    # base 0 is indistinguishable from plain field access
            hi = base + stride * 64
            inside = [o for o in found if base <= o < hi]
            within = [o for o in inside if (o - base) < stride]
            ok = bool(within)
            print("  CONTROL %-12s base %#06x stride %3d: %d offsets in region, "
                  "%d within one element -> %s"
                  % (label, base, stride, len(inside), len(within),
                     "OK" if ok else "NOT REACHED"))
            if not ok:
                bad.append(label)
        for e in GCP_EXACT:
            ok = e in found
            print("  CONTROL exact %#06x (asd_refcount): %s" % (e, "OK" if ok else "MISSING"))
            if not ok:
                bad.append(hex(e))
        print("  -> %s" % ("harvester trustworthy" if not bad
                            else "NOT trustworthy: %s" % bad))
        rc = 1 if bad else 0
    if a.json:
        out = {hex(off): {"hits": n, "widths": {str(k): v for k, v in w.items()},
                          "sites": s} for off, n, w, s in rows}
        Path(a.json).write_text(json.dumps(out, indent=1, sort_keys=True))
        print("  wrote %s" % a.json)
    try:
        client.close()
    except Exception:
        pass
    return rc


if __name__ == "__main__":
    sys.exit(main())