#!/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 struct
import sys

IDBL_OFFSET = 0x8000
RKNS_MAGIC  = 0x534E4B52   # "RKNS" LE — rkspi boot header
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):
    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}")

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

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print(__doc__)
        sys.exit(1)
    main(sys.argv[1])