#!/usr/bin/env python3
"""Check whether GCP's three defects exist across every harvested GCP build.

The HDCP sweep is the precedent: the four overflow sites there turned out to be present in
all 17 builds, which converted "a bug" into "a bug AMD has never fixed". The same question
matters here, and for the same reason: a defect present across three years of builds is a
different report from a recent regression.

Handlers are located by their OWN log string (`AMD_TA_GCP: DestroyCryptoSession ENTER`),
not by address, because addresses move between builds. From the string reference the
enclosing function is recovered by walking back to its prologue and forward to its epilogue,
then hashed and compared against the analysed build (22.0.4.50).

Identical bytes mean the defect is present unchanged. Different bytes mean the handler was
touched and needs reading -- reported as CHANGED rather than guessed either way.

Usage:  ./gcp_version_sweep.py
"""
from __future__ import annotations

import hashlib
import json
import re
import struct
from pathlib import Path

import capstone as cs

HERE = Path(__file__).resolve().parent
REPO = HERE.parent
BODIES = REPO / "psp_out" / "versions_win"
BASE = 0x400100
BASELINE = "ta_22-0-4-50_81b082b4.bin"

TARGETS = {
    "FE0D DestroyCryptoSession": b"AMD_TA_GCP: DestroyCryptoSession ENTER",
    "FE0F GetCryptoSessionForDevice": b"AMD_TA_GCP: GetCryptoSessionForDevice ENTER",
    "FE0C CreateCryptoSession": b"AMD_TA_GCP: CreateCryptoSession ENTER",
}

PUSH_T1 = re.compile(rb"[\x00-\xff]\xb5")          # push {..., lr}
PUSH_T2 = b"\x2d\xe9"                                # push.w {...}


def find_str_refs(b: bytes, s: bytes) -> list[int]:
    """Addresses of instructions that materialise a pointer to string s."""
    off = b.find(s)
    if off < 0:
        return []
    va = BASE + off
    md = cs.Cs(cs.CS_ARCH_ARM, cs.CS_MODE_THUMB)
    refs = []
    # ADR-style: pc-relative add. Scan a window before the string for adr/ldr literals.
    for probe in range(max(0, off - 0x4000), off, 2):
        for i in md.disasm(b[probe:probe + 4], BASE + probe):
            o = i.op_str or ""
            if i.mnemonic == "adr":
                m = re.search(r"#(-?0x[0-9a-f]+|-?\d+)", o)
                if m:
                    v = m.group(1)
                    tgt = ((i.address + 4) & ~3) + (int(v, 16) if "0x" in v else int(v))
                    if tgt == va:
                        refs.append(i.address)
            elif i.mnemonic.startswith("ldr") and "pc" in o:
                m = re.search(r"#(0x[0-9a-f]+)", o)
                if m:
                    lit = ((i.address + 4) & ~3) + int(m.group(1), 16)
                    if 0 <= lit - BASE < len(b) - 4:
                        if struct.unpack_from("<I", b, lit - BASE)[0] == va:
                            refs.append(i.address)
            break
    return refs


def enclosing(b: bytes, ea: int) -> tuple[int, int] | None:
    """(start, size) of the function containing ea, by prologue/epilogue walk."""
    start = None
    for probe in range(ea, max(BASE, ea - 0x400) - 1, -2):
        o = probe - BASE
        if o < 0 or o + 2 > len(b):
            continue
        pair = b[o:o + 2]
        if pair[1] == 0xB5 or pair == PUSH_T2[::-1] or b[o:o + 2] == b"\x2d\xe9":
            start = probe
            break
    if start is None:
        return None
    # A handler has SEVERAL epilogues (each early return pops). Stopping at the first one
    # past the string reference measured only the leading guard -- 30 bytes for a 100-byte
    # function. Walk to the LAST epilogue before the next function's prologue instead.
    md = cs.Cs(cs.CS_ARCH_ARM, cs.CS_MODE_THUMB)
    end = None
    seen_epilogue = False
    for i in md.disasm(b[start - BASE:start - BASE + 0x800], start):
        if (i.mnemonic.startswith("pop") and "pc" in (i.op_str or "")) or i.mnemonic in ("bx",):
            end = i.address + i.size
            seen_epilogue = True
            continue
        if seen_epilogue and i.mnemonic.startswith("push"):
            break          # next function begins
    if end is None or end <= start:
        return None
    return start, end - start


def digest(b: bytes, start: int, size: int) -> str:
    """Hash the MNEMONIC SEQUENCE, not the bytes.

    Raw bytes differ between builds even when the logic is identical, because literal pool
    offsets and branch displacements move. The opcode sequence is the thing that stays
    stable under relocation, so it is what tells us whether a handler was really touched.
    """
    md = cs.Cs(cs.CS_ARCH_ARM, cs.CS_MODE_THUMB)
    seq = [i.mnemonic for i in md.disasm(b[start - BASE:start - BASE + size], start)]
    return hashlib.sha256((",".join(seq)).encode()).hexdigest()[:12]


def main() -> int:
    builds = sorted(p for p in BODIES.glob("ta_22-*.bin"))
    base_body = (BODIES / BASELINE).read_bytes()
    baseline = {}
    for tag, s in TARGETS.items():
        refs = find_str_refs(base_body, s)
        fn = enclosing(base_body, refs[0]) if refs else None
        baseline[tag] = (digest(base_body, *fn), fn[1]) if fn else None
    print("baseline %s:" % BASELINE)
    for tag, v in baseline.items():
        print("   %-32s %s" % (tag, ("hash %s size %d" % v) if v else "NOT FOUND"))

    print("\n%-28s %s" % ("build", "  ".join("%-22s" % t.split()[0] for t in TARGETS)))
    for p in builds:
        b = p.read_bytes()
        cells = []
        for tag, s in TARGETS.items():
            refs = find_str_refs(b, s)
            if not refs:
                cells.append("string absent")
                continue
            fn = enclosing(b, refs[0])
            if not fn:
                cells.append("fn unresolved")
                continue
            d, sz = digest(b, *fn), fn[1]
            want = baseline.get(tag)
            if want and d == want[0]:
                cells.append("IDENTICAL")
            else:
                cells.append("CHANGED(%d B)" % sz)
        print("%-28s %s" % (p.name.replace("ta_", "").replace(".bin", ""),
                            "  ".join("%-22s" % c for c in cells)))
    return 0


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