#!/usr/bin/env python3
"""Wrap carved TA bodies as ELF32/ARM images IDA can load, for the whole corpus.

The twelve reference ELFs in `psp_out/versions_win/elf/` were produced ad hoc and only ever
covered one build per TA type, which is why cross-version diffing was never possible: there
was nothing to diff against. This generates one ELF per corpus build, in the *same* format,
so the existing symbol/porting tooling applies unchanged.

Format, read back off the existing ELFs rather than assumed (see --verify):

    ELF32, little endian, EXEC, machine ARM, flags 0x5000000 (EABI v5)
    entry    = 0x400101           (Thumb bit set)
    one PT_LOAD: offset 0x1000, vaddr/paddr 0x400100, filesz = memsz = len(body), RWE
    file     = 0x1000-byte zero-padded header area, then the body verbatim

`psp_out/versions_win/ta_*.bin` are already *bodies* (container + 0x100 stripped): body
offset 0 maps to address 0x400100, which is what makes the recorded function hashes in
`fw/libdumps/*.json` line up. Do not feed this script a container.

Naming: `<name>_<major>.<minor>.<patch>.<build>_<hash8>.elf`. Fifteen corpus files carry a
stale `ta_0.7.0.1_<hash>` name locally while the card's copy has the resolved version, so the
type/version comes from `--namemap` (hash8 -> "ta_<maj>-<min>-<pat>-<bld>") when supplied,
and from the filename otherwise.

usage:
  ./mkelf_ta.py --verify                       # regenerate the 12 known ELFs, byte-compare
  ./mkelf_ta.py --all [--outdir DIR]           # build every corpus body
"""
from __future__ import annotations

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

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

VADDR = 0x00400100
ENTRY = 0x00400101          # Thumb
PHOFF = 52
HDRPAD = 0x1000
EF_ARM = 0x5000000

# TA type number -> component name. Sources: fw/syms/index.json for the twelve Windows
# types, and AUDIT-WINDOWS-DRIVER-CORPUS.md for the four linux-firmware ones
# (RAS 27, HDCP 23, DTM 18, ASD 33).
TYPE_NAME = {
    7: "rap", 16: "auc", 18: "dtm", 19: "fp", 22: "gcp", 23: "hdcp", 24: "merit",
    25: "opm", 27: "ras", 29: "tmz", 30: "wspe", 32: "xgmi", 34: "powerplay",
    42: "csu", 50: "oc",
}


def build_elf(body: bytes, vaddr: int = VADDR, entry: int = ENTRY) -> bytes:
    """Produce the ELF exactly as the existing reference images are laid out."""
    n = len(body)
    VADDR_, ENTRY_ = vaddr, entry
    eh = struct.pack(
        "<16sHHIIIIIHHHHHH",
        b"\x7fELF\x01\x01\x01" + b"\x00" * 9,
        2,            # ET_EXEC
        40,           # EM_ARM
        1,            # version
        ENTRY_,
        PHOFF,        # e_phoff
        0,            # e_shoff
        EF_ARM,
        52, 32, 1,    # ehsize, phentsize, phnum
        0, 0, 0,      # shentsize, shnum, shstrndx
    )
    ph = struct.pack(
        "<IIIIIIII",
        1,            # PT_LOAD
        HDRPAD,       # p_offset
        VADDR_, VADDR_,
        n, n,         # filesz, memsz
        7,            # PF_R|PF_W|PF_X
        0x1000,       # align
    )
    head = eh + ph
    assert len(head) <= HDRPAD, "header area overflow"
    return head + b"\x00" * (HDRPAD - len(head)) + body


def parse_name(fname: str, namemap: dict[str, str] | None):
    """-> (name, version_string, hash8) or None when the version cannot be resolved."""
    m = re.match(r"^ta_(.+?)_([0-9a-f]{8})\.bin$", fname)
    if not m:
        return None
    vpart, h8 = m.group(1), m.group(2)
    if namemap and h8 in namemap:
        vpart = namemap[h8]
    nums = re.findall(r"\d+", vpart)
    if len(nums) != 4:
        return None                      # unresolved (e.g. the stale "0.7.0.1" names)
    major = int(nums[0])
    name = TYPE_NAME.get(major)
    if name is None:
        return None
    return name, ".".join(nums), h8


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--verify", action="store_true",
                    help="regenerate the existing reference ELFs and byte-compare")
    ap.add_argument("--all", action="store_true", help="build every corpus body")
    ap.add_argument("--bodies", default=str(REPO / "psp_out" / "versions_win"))
    ap.add_argument("--outdir", default=str(REPO / "psp_out" / "versions_win" / "elf_all"))
    ap.add_argument("--namemap", default=str(HERE / "ta_namemap.json"),
                    help="hash8 -> ta_<maj>-<min>-<pat>-<bld>, for stale local names")
    a = ap.parse_args()

    namemap = None
    if os.path.exists(a.namemap):
        namemap = json.loads(Path(a.namemap).read_text())

    if a.verify:
        refs = sorted(glob.glob(str(REPO / "psp_out" / "versions_win" / "elf" / "*.elf")))
        ok = bad = skip = 0
        for elf in refs:
            raw = Path(elf).read_bytes()
            ph = struct.unpack("<I", raw[28:32])[0]
            off, va, _pa, fsz = struct.unpack("<IIII", raw[ph + 4:ph + 20])
            body = raw[off:off + fsz]
            if va != VADDR:
                print("  SKIP %-42s vaddr %#x" % (os.path.basename(elf), va)); skip += 1
                continue
            if build_elf(body) == raw:
                ok += 1
            else:
                bad += 1
                print("  MISMATCH %s" % os.path.basename(elf))
        print("verify: %d byte-identical, %d mismatched, %d skipped" % (ok, bad, skip))
        return 1 if bad else 0

    if not a.all:
        ap.error("one of --verify / --all is required")

    out = Path(a.outdir)
    out.mkdir(parents=True, exist_ok=True)
    built, unresolved = [], []
    for p in sorted(glob.glob(os.path.join(a.bodies, "ta_*.bin"))):
        info = parse_name(os.path.basename(p), namemap)
        if info is None:
            unresolved.append(os.path.basename(p)); continue
        name, ver, h8 = info
        body = Path(p).read_bytes()
        dst = out / f"{name}_{ver}_{h8}.elf"
        dst.write_bytes(build_elf(body))
        built.append((name, ver, h8, len(body)))

    # ASD is driver-class rather than a TA: different load base, and no TA entry trampoline.
    # Base 0xD00100 matches psp_out/handover/elf/asd2.elf, whose LOAD segment contains
    # 1562/1562 of the recorded ASD symbols, so the convention is the project's existing one.
    for p in sorted(glob.glob(os.path.join(a.bodies, "asd_*.bin"))):
        fn = os.path.basename(p)
        m = re.match(r"^asd_(.+?)_([0-9a-f]{8})\.bin$", fn)
        if not m:
            unresolved.append(fn); continue
        vpart, h8 = m.group(1), m.group(2)
        if namemap and h8 in namemap:
            vpart = namemap[h8]
        nums = re.findall(r"\d+", vpart)
        if len(nums) != 4 or nums[0] != "33":
            unresolved.append(fn); continue
        ver = ".".join(nums)
        body = Path(p).read_bytes()
        dst = out / f"asd_{ver}_{h8}.elf"
        dst.write_bytes(build_elf(body, vaddr=0x00D00100, entry=0x00D00101))
        built.append(("asd", ver, h8, len(body)))

    bytype: dict[str, int] = {}
    for name, *_ in built:
        bytype[name] = bytype.get(name, 0) + 1
    print("built %d ELFs in %s" % (len(built), out))
    for k in sorted(bytype):
        print("   %-10s %d" % (k, bytype[k]))
    if unresolved:
        print("UNRESOLVED (%d) -- version/type not determinable, no ELF written:" % len(unresolved))
        for u in unresolved:
            print("   ", u)
    return 0


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