#!/usr/bin/env python3
"""RK3588 SPI image pre-flash integrity check.

Verifies a u-boot-rockchip-spi*.bin image *before* the user commits
to a maskrom-recovery-prone flash. Catches the class of silent-build
failures that brick boards:

  - missing idbloader (binman packed 0xFF where RKNS wrapper should be)
  - truncated idbloader (less than the declared payload region)
  - mkimage-rejected TPL blob that left the SPL slot empty

Does NOT execute code — pure static parsing. Fast and safe.

Exit codes:
  0  = image has a valid-looking idbloader; safe to flash
  1  = file unreadable / too small
  2  = missing RKNS wrapper at 0x8000 (the "today's bug" case)
  3  = wrapper present but surrounding region is all 0xFF (payload missing)

Usage: rk3588_spi_check.py <spi-image.bin>
"""
import argparse, hashlib, struct
import sys

IDBL_OFFSET = 0x8000
RKNS_MAGIC  = 0x534E4B52   # "RKNS" LE — rkspi boot header
TPL_OFFSET  = 0x8800       # plaintext DDR blob (TPL) starts here in rkspi layout
TPL_ENTRY   = bytes.fromhex("01000014")  # first 4 bytes of the RK3588 DDR blob entry
PAYLOAD_END = 0x60000      # stock u-boot places next section here

def die(code, msg):
    print(f"FAIL: {msg}", file=sys.stderr)
    sys.exit(code)

def main(path, blob_path=None):
    with open(path, "rb") as f:
        data = f.read()
    print(f"SPI image: {path}  size=0x{len(data):x} ({len(data)} B)")

    if len(data) < PAYLOAD_END:
        die(1, f"image too small ({len(data)} < 0x{PAYLOAD_END:x})")

    wrapper = struct.unpack_from("<I", data, IDBL_OFFSET)[0]
    if wrapper != RKNS_MAGIC:
        die(2, f"no RKNS wrapper at 0x{IDBL_OFFSET:x}: got 0x{wrapper:08x} "
               f"(expect 0x{RKNS_MAGIC:08x}). idbloader was not produced — "
               f"silently-failed mkimage during u-boot build.")
    print(f"OK  RKNS wrapper present at 0x{IDBL_OFFSET:x}")

    # Count non-0xFF bytes in the SPL/TPL region. If near-zero, the wrapper
    # might be there but the payload got padded out.
    region = data[IDBL_OFFSET:PAYLOAD_END]
    content = sum(1 for b in region if b != 0xFF)
    total = len(region)
    pct = 100.0 * content / total
    print(f"    payload region 0x{IDBL_OFFSET:x}..0x{PAYLOAD_END:x}:  "
          f"{content}/{total} non-0xFF bytes ({pct:.1f}%)")
    if content < 0x1000:
        die(3, f"idbloader region is almost entirely 0xFF "
               f"({content} bytes — payload missing or truncated)")

    # Quick look at what's in the wrapper header for the human reader
    w = data[IDBL_OFFSET:IDBL_OFFSET + 16]
    size_like = struct.unpack_from("<I", w, 8)[0]
    flag_like = struct.unpack_from("<I", w, 12)[0]
    print(f"    wrapper[+8]=0x{size_like:08x}  wrapper[+12]=0x{flag_like:08x}")

    # TPL content check — the DDR blob lives in plaintext at 0x8800
    tpl_prefix = data[TPL_OFFSET:TPL_OFFSET + 4]
    if tpl_prefix != TPL_ENTRY:
        die(4, f"TPL at 0x{TPL_OFFSET:x} does not start with expected blob entry "
               f"{TPL_ENTRY.hex()}: got {tpl_prefix.hex()}")
    print(f"OK  TPL blob entry signature at 0x{TPL_OFFSET:x}")

    if blob_path:
        with open(blob_path, "rb") as f:
            blob = f.read()
        tpl_region = data[TPL_OFFSET:TPL_OFFSET + len(blob)]
        if tpl_region != blob:
            # Find the first differing byte for diagnostics
            diffs = [i for i in range(len(blob)) if i < len(tpl_region) and blob[i] != tpl_region[i]]
            first_diff = diffs[0] if diffs else len(tpl_region)
            die(5, f"TPL in SPI does NOT match blob file.\n"
                   f"    blob sha256: {hashlib.sha256(blob).hexdigest()}\n"
                   f"    SPI[0x{TPL_OFFSET:x}:+{len(blob)}] sha256: {hashlib.sha256(tpl_region).hexdigest()}\n"
                   f"    first differing byte at blob offset 0x{first_diff:x}")
        print(f"OK  TPL bytes match {blob_path}  sha256={hashlib.sha256(blob).hexdigest()[:16]}…  ({len(blob)} B)")

    print("\nPASS: image looks structurally sound. Safe to flash.")
    sys.exit(0)

if __name__ == "__main__":
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("spi", help="u-boot-rockchip-spi-*.bin to validate")
    ap.add_argument("--blob", help="optional raw DDR blob to byte-compare against the TPL region")
    a = ap.parse_args()
    main(a.spi, a.blob)