#!/usr/bin/env python3
"""
spi_parse.py - parse / extract an AMD Navi3x (Zen4-era PSP, "MP0 13.0.0") dGPU
SPI flash image.

Written against `current-spi.bin`, a 2 MiB dump of the SPI ROM of an XFX Radeon
RX 7900 XTX (Navi 31, PCI 1002:744C).  The layout is the AMD Embedded Firmware
Structure (EFS) + PSP directory family, in the A/B "image slot" arrangement that
Zen4-era parts use, with GPU-specific directory entry types.

    python3 spi_parse.py current-spi.bin [-o spi_out] [--no-extract]

Produces:
    <out>/raw/      every directory entry, verbatim (header+body+signature)
    <out>/payload/  the signed/decompressed body of every entry we can unwrap
    <out>/keys/     every public key found (AMD root token + KDB), as PEM
    <out>/summary.json

Verification performed:
    * Fletcher-32 over every directory / image-slot header
    * RSA-PSS over every firmware image (MGF1, salt = digest length;
      SHA-256 for 2048-bit keys, SHA-384 for 4096-bit keys)
    * RSA-PSS over the whole-flash "$SGN" signature table and the type-0x41
      entry, and the coverage of each such signature is searched for.

Structures, for reference (all little-endian):

  embedded_firmware (EFS) @ 0x00000, magic 0x55AA55AA
      +0x14 u32  -> L1 PSP directory, slot A          (0x02000 here)
      +0x2C u32  -> L1 PSP directory, slot B          (0x03000 here)
      +0x4C u32  -> "$ROM" board/ASIC config, slot A  (0x10000 here)
      +0x50 u32  -> "$ROM" board/ASIC config, slot B  (0x11000 here)

  psp_directory  magic "$PSP"
      +0x00 char[4] cookie       +0x04 u32 fletcher32
      +0x08 u32 num_entries      +0x0C u32 reserved
      then num_entries * psp_directory_entry:
          u8 type; u8 subprog; u16 rsvd; u32 size; u64 loc
          loc[63:62] = address mode: 0 = absolute flash offset,
                                     3 = relative to the containing image slot

  image_slot_header (ISH), 0x20 bytes, pointed at by L1 entry types 0x3C/0x3D
      u32 fletcher32; u32 boot_priority; u32 update_retry_count;
      u32 counter;    u32 image_slot_addr; u32 psp_id (0xBC0D0500 = Navi3x);
      u32 slot_size;  u32 second_ptr (VBIOS base inside the slot)

  PSP_FW_IMAGE_HEADER, 0x100 bytes.  Field names below are AMD's own, from
  psp_fw_image.h in the Milan (Zen3) ASP bootloader source; the struct is
  unchanged on this Navi 31 part.
      0x00 Nonce[16]          0x10 HeaderVersion ("$PS1" or 0)
      0x14 SizeFWSigned       0x18 EncOption      0x1C IkekType (0 = iKEK-AMD)
      0x20 EncParameters[16]  0x30 SigOption      0x34 SigAlgID
      0x38 SigParameters[16] <- 16-byte id of the signing key
      0x48 CompOption         0x4C SecPatchLevel  0x50 UnCompImageSize
      0x54 CompImageSize      0x58 CompParameters[8]
      0x60 ImageVersion       0x64 APUFamilyID    0x68 FirmwareLoadAddr
      0x6C SizeImage          0x70 SizeFWUnSigned
      0x74 FirmwareSplitAddr  ("joining point of combined FWs, Nwd/Swd split")
      0x7C FwType             ("must match Directory Table Entry Type"; Renoir
                               names the same field FwTypeLegacy, and here it
                               carries the *older* GPU type numbering while the
                               L2 directory uses a renumbered space)
      0x7D FwSubType          0x80 EncKey[16] (wrapped MEK)
      0x90 SigningInfo[16]    0xA0 FwSpecificData[32]
      0xC0 DebugEncKey[16]    0xD0 Padd[48]
      image = header || body(SizeFWSigned) || signature(modulus bytes)

  AMD key token (directory types 0x39/0x32/0x3A)
      u32 version; u8 key_id[16]; u8 certifying_key_id[16]; u32 key_usage;
      u32 rsvd[4]; u32 pubexp_bits; u32 modulus_bits;
      pubexp[mb/8] (LE) ; modulus[mb/8] (LE) ; signature (LE, PSS)

  $KDB key database (payload of directory type 0x4E)
      u32 size; u32 version; char[4] "$KDB"; ... entries from +0x50:
      u32 DataSize; u32 f1; u32 key_index; u32 PubExp;
      u8 KeyId[16]; u32 KeySize(bits); ... +0x50 Modulus[KeySize/8] (LE)
  AMD calls these KEY_DB_HEADER / KEY_DB_PUB_KEY, and defines *two* databases:
  KEY_DATBASE_PSP_BL (0x50) and KEY_DATBASE_PSP_TOS (0x51).  That is why this
  SPI KDB and the one in linux-firmware's psp_13_0_0_sos.bin share only 6 of
  their 17 / 15 keys.

  Offsets this firmware is known to write (traced in sys_drv2.elf):
      0x007000   12-byte marker, "$AMDVBFL" followed by 4 NUL bytes, written by
                 vbios_flash_verify() when vbios_find_flash_header() (which
                 maps VA 0x0D107000 = SPI window base + 0x7000) misses it
      0x1EF000   0x48-byte "SPIS" record: checksum16 over flash[0:0x1E0000]
                 plus an HMAC-SHA256; written, read back and compared
  Both are blank here, and both lie outside every signature.
"""

import argparse
import hashlib
import json
import os
import re
import struct
import sys
import zlib

# --------------------------------------------------------------------------- #
# names
# --------------------------------------------------------------------------- #

# Directory entry types.  This is the GPU type space, which only partly overlaps
# the CPU/APU one that psptool / amdfwtool document.  "!" marks a name that is
# established from the content of this image rather than from documentation.
DIR_TYPES = {
    # NB: there is no SYS_DRV entry in this image -- neither the GPU id 0x13 nor
    # the CPU id 0x28.  It is supplied by the host driver over the bootloader
    # mailbox on every cold boot.
    (0x01, 0x00): 'PSP_BOOTLOADER',        # ARM, encrypted, load 0x17500
    (0x02, 0x00): 'PSP_SOS',               # ARM, encrypted, load 0x00100
    (0x0B, 0x00): 'PSP_DRV_0B',            # ARM module, encrypted, load 0x80100
    (0x31, 0x01): 'FW_STUB_31',
    (0x32, 0x01): 'PUBKEY_USAGE3',         # ! RSA-2048 token, certified by root
    (0x33, 0x01): 'TABLE_33',
    (0x37, 0x01): 'VBIOS_IMAGE',           # ! legacy x86 ATOMBIOS option ROM
    (0x39, 0x01): 'AMD_ROOT_PUBKEY',       # ! RSA-4096 self-signed token
    (0x3A, 0x01): 'PUBKEY_USAGE8',         # ! image-signing key (see $SGN)
    (0x3A, 0x20): 'PSP_TOC',               # ! same signing key as PSP_TOC in
                                           #   linux-firmware psp_13_0_0_sos.bin
    (0x3B, 0x01): 'DIGEST_3B',             # ! bare 48-byte SHA-384 digest
    (0x40, 0x00): 'PSP_DRV_40',            # ARM module, load 0x80100
    (0x41, 0x01): 'IMAGE_SIGNATURE',       # ! RSA-2048 PSS sig by PUBKEY_USAGE8
    (0x42, 0x01): 'FLAG_42',
    (0x44, 0x01): 'NV_RECORD',             # ! 64 KiB record area; sys_drv writes the
                                           #   SPIS record at +0xF000, rest unused
    (0x45, 0x01): 'ASIC_TABLE_45',
    (0x46, 0x01): 'ASIC_TABLE_46',
    (0x47, 0x01): 'VBL_DATA_TABLE',        # ! ATOM "VBL_DATA_TABLE" blob
    (0x4B, 0x00): 'FW_4B',
    (0x4E, 0x00): 'PSP_KDB',               # ! payload carries the "$KDB" magic
    (0x55, 0x00): 'PSP_SPL',               # = SPL_TABLE_BL in AMD's own enum;
                                           #   the only id matching a documented one
    (0x59, 0x00): 'PSP_DRV_59',            # ARM module, encrypted, load 0x80100
    (0x5B, 0x00): 'INFO_5B',
    (0x5C, 0x00): 'PSP_DRV_5C',            # ARM module, load 0x80100
    (0x5E, 0x00): 'PSP_DRV_5E',            # ARM module, load 0x80100
    # subprog 0x10 = GPU IP microcontroller firmware
    (0x01, 0x10): 'SMU_PMFW',              # ! payload starts with SMU version
    (0x05, 0x10): 'IPFW_05',
    (0x07, 0x10): 'IP_DISCOVERY',          # ! binary_signature 0x28211407
    # 0x09/0x10 carries FwType 0x11 = PSP_DIR_ENTRY_TYPE_MC_FW in AMD's GPU
    # directory enum.  Together with 0x0A/0x10 (same version, same key, adjacent,
    # loading at 0x0 and 0xC000) it is the UMC memory-controller sequencer
    # micro-program that the PSP streams into each memory controller through
    # UMCCTRL_SUP_CNTL/UMCCTRL_SUP_PGM.  Microcode, not CPU code.
    (0x09, 0x10): 'UMC_MC_FW',
    (0x0A, 0x10): 'UMC_MC_FW_PT2',
    (0x1B, 0x10): 'MPIO',                  # ! "release ci build for MPIO"
    (0x1D, 0x10): 'IPFW_1D',
    # L1 directory
    (0x3C, 0x01): 'IMAGE_SLOT_HDR_A',
    (0x3D, 0x01): 'IMAGE_SLOT_HDR_B',
}

DISCOVERY_HW = {
    1: 'MP1', 2: 'MP2', 3: 'THM', 4: 'SMUIO', 5: 'FUSE', 6: 'CLKA', 10: 'PWR',
    11: 'GC', 12: 'UVD/VCN', 13: 'AUDIO_AZ', 14: 'ACP', 15: 'DCI', 16: 'DCO',
    17: 'XDMA', 18: 'DCEAZ', 19: 'SDPMUX', 20: 'NTB', 24: 'IOHC', 28: 'L2IMU',
    32: 'VCE', 34: 'MMHUB', 35: 'ATHUB', 36: 'DBGU_NBIO', 37: 'DFX', 38: 'DBGU0',
    39: 'DBGU1', 40: 'OSSSYS', 41: 'HDP', 42: 'SDMA0', 43: 'SDMA1', 44: 'ISP',
    45: 'DBGU_IO', 46: 'DF', 47: 'CLKB', 48: 'FCH', 49: 'DFX_DAP', 70: 'PCIE',
    80: 'PCS', 89: 'DDCL', 90: 'SST', 91: 'LSDMA', 100: 'IOAGR', 108: 'NBIF',
    124: 'IOAPIC', 128: 'SYSTEMHUB', 144: 'NTBCCP', 150: 'UMC', 168: 'SATA',
    170: 'USB', 176: 'CCXSEC', 200: 'XGMI', 216: 'XGBE', 255: 'MP0',
    256: 'SDMA2', 257: 'SDMA3', 258: 'VPE', 271: 'DMU', 272: 'DIO', 273: 'DAZ',
    274: 'DAZ', 282: 'MALL', 286: 'VCN_UNIFIED',
}

EFS_MAGIC = 0x55AA55AA
PSP_COOKIE = b'$PSP'
ROM_COOKIE = b'MOR$'          # "$ROM" stored as a big-endian-ASCII u32
SGN_COOKIE = b'$SGN'


# --------------------------------------------------------------------------- #
# helpers
# --------------------------------------------------------------------------- #

def fletcher32(buf):
    c0 = c1 = 0xFFFF
    i, n = 0, len(buf) // 2
    while n:
        blk = min(n, 359)
        n -= blk
        for _ in range(blk):
            c0 = (c0 + struct.unpack_from('<H', buf, i)[0]) & 0xFFFFFFFF
            i += 2
            c1 = (c1 + c0) & 0xFFFFFFFF
        c0 = (c0 & 0xFFFF) + (c0 >> 16)
        c1 = (c1 & 0xFFFF) + (c1 >> 16)
    c0 = (c0 & 0xFFFF) + (c0 >> 16)
    c1 = (c1 & 0xFFFF) + (c1 >> 16)
    return ((c1 << 16) | c0) & 0xFFFFFFFF


def b2l(b):
    return int.from_bytes(b, 'big')


def mgf1(seed, length, hfunc):
    out, c = b'', 0
    while len(out) < length:
        out += hfunc(seed + struct.pack('>I', c)).digest()
        c += 1
    return out[:length]


def pss_open(sig_be, n, e, hfunc, hlen):
    """Raw-verify a PSS signature; return (H, salt) or None."""
    k = (n.bit_length() + 7) // 8
    if len(sig_be) != k:
        return None
    em = pow(b2l(sig_be), e, n).to_bytes(k, 'big')
    if em[-1] != 0xBC:
        return None
    dblen = k - hlen - 1
    h = em[dblen:dblen + hlen]
    db = bytes(a ^ b for a, b in zip(em[:dblen], mgf1(h, dblen, hfunc)))
    db = bytes([db[0] & 0x7F]) + db[1:]
    j = 0
    while j < len(db) and db[j] == 0:
        j += 1
    if j >= len(db) or db[j] != 1:
        return None
    return h, db[j + 1:]


def pss_check(h, salt, msg, hfunc):
    mh = hfunc(msg).digest()
    return hfunc(b'\x00' * 8 + mh + salt).digest() == h


def hash_for_bits(bits):
    return (hashlib.sha384, 48) if bits > 2048 else (hashlib.sha256, 32)


def ver_str(raw4):
    """AMD stores the version as four bytes; print them in file order."""
    return '%d.%d.%d.%d' % tuple(raw4)


# --------------------------------------------------------------------------- #
# image
# --------------------------------------------------------------------------- #

class Flash:
    def __init__(self, data):
        self.d = data
        self.keys = {}      # key_id(bytes16) -> dict(n, e, bits, source)
        self.notes = []

    def u8(self, o):
        return self.d[o]

    def u16(self, o):
        return struct.unpack_from('<H', self.d, o)[0]

    def u32(self, o):
        return struct.unpack_from('<I', self.d, o)[0]

    def u64(self, o):
        return struct.unpack_from('<Q', self.d, o)[0]

    # ---------------- top level ----------------

    def parse_efs(self):
        assert self.u32(0) == EFS_MAGIC, 'no EFS signature at offset 0'
        efs = {'magic': '0x%08X' % EFS_MAGIC,
               'nonzero_fields': {'+0x%02x' % o: '0x%08x' % self.u32(o)
                                  for o in range(0, 0x80, 4) if self.u32(o)}}
        efs['l1_dir_a'] = self.u32(0x14)
        efs['l1_dir_b'] = self.u32(0x2C)
        efs['rom_cfg_a'] = self.u32(0x4C)
        efs['rom_cfg_b'] = self.u32(0x50)
        return efs

    def parse_dir(self, base):
        cookie = self.d[base:base + 4]
        if cookie != PSP_COOKIE:
            return None
        stored = self.u32(base + 4)
        n = self.u32(base + 8)
        body = self.d[base + 8: base + 0x10 + n * 0x10]
        ents = []
        for i in range(n):
            o = base + 0x10 + i * 0x10
            t, sub, rsvd = struct.unpack_from('<BBH', self.d, o)
            size = self.u32(o + 4)
            loc = self.u64(o + 8)
            mode = loc >> 62
            addr = loc & ((1 << 62) - 1)
            ents.append(dict(index=i, type=t, subprog=sub, rsvd=rsvd,
                             size=size, addr=addr, mode=mode,
                             name=DIR_TYPES.get((t, sub), 'TYPE_%02X_%02X' % (t, sub))))
        return dict(base=base, cookie='$PSP', checksum='0x%08x' % stored,
                    checksum_ok=(fletcher32(body) == stored),
                    num_entries=n, entries=ents)

    def parse_ish(self, base):
        stored = self.u32(base)
        f = struct.unpack_from('<7I', self.d, base + 4)
        return dict(offset=base, checksum='0x%08x' % stored,
                    checksum_ok=(fletcher32(self.d[base + 4:base + 0x20]) == stored),
                    boot_priority=f[0], update_retry_count=f[1], counter='0x%x' % f[2],
                    image_slot_addr='0x%x' % f[3], psp_id='0x%08x' % f[4],
                    slot_size='0x%x' % f[5], second_ptr='0x%x' % f[6])

    def parse_sgn(self, base):
        if self.d[base:base + 4] != SGN_COOKIE:
            return None
        ver, cnt, rsvd = struct.unpack_from('<III', self.d, base + 4)
        out = []
        for i in range(cnt):
            a, b, size, off = struct.unpack_from('<IIII', self.d, base + 0x10 + i * 0x10)
            out.append(dict(id=a, flags=b, size='0x%x' % size, offset='0x%x' % off,
                            offset_int=off, size_int=size))
        return dict(offset=base, version=ver, count=cnt, entries=out)

    # ---------------- keys ----------------

    def parse_key_token(self, o):
        ver = self.u32(o)
        kid = self.d[o + 4:o + 0x14]
        cid = self.d[o + 0x14:o + 0x24]
        usage = self.u32(o + 0x24)
        ebits, mbits = struct.unpack_from('<II', self.d, o + 0x38)
        if mbits not in (2048, 4096) or ebits not in (2048, 4096):
            return None
        ml = mbits // 8
        e = int.from_bytes(self.d[o + 0x40:o + 0x40 + ml][::-1], 'big')
        n = int.from_bytes(self.d[o + 0x40 + ml:o + 0x40 + 2 * ml][::-1], 'big')
        sig_off = o + 0x40 + 2 * ml
        return dict(version=ver, key_id=kid, cert_key_id=cid, usage=usage,
                    exp_bits=ebits, mod_bits=mbits, e=e, n=n, sig_off=sig_off)

    def load_kdb(self, payload):
        """payload = body of the type 0x4E entry."""
        if payload[8:12] != b'$KDB':
            return []
        out, o = [], 0x50
        while o < len(payload) - 0x10:
            ln = struct.unpack_from('<I', payload, o)[0]
            if ln == 0 or ln > 0x1000 or o + ln > len(payload):
                break
            f1, idx, exp = struct.unpack_from('<III', payload, o + 4)
            kid = payload[o + 0x10:o + 0x20]
            bits = struct.unpack_from('<I', payload, o + 0x20)[0]
            n = int.from_bytes(payload[o + 0x50:o + 0x50 + bits // 8][::-1], 'big')
            self.keys[kid] = dict(n=n, e=exp, bits=bits, source='KDB')
            out.append(dict(offset='0x%x' % o, key_id=kid.hex(), key_index=idx,
                            bits=bits, exponent='0x%x' % exp))
            o += ln
        return out


# --------------------------------------------------------------------------- #
# firmware image header
# --------------------------------------------------------------------------- #

FW_HDR_FIELDS = [
    ('nonce', 0x00, 16), ('header_version', 0x10, 4), ('size_fw_signed', 0x14, 4),
    ('enc_option', 0x18, 4), ('ikek_type', 0x1C, 4), ('enc_params', 0x20, 16),
    ('sig_option', 0x30, 4), ('sig_alg', 0x34, 4), ('sig_params', 0x38, 16),
    ('comp_option', 0x48, 4), ('sec_patch_level', 0x4C, 4), ('uncomp_size', 0x50, 4),
    ('comp_size', 0x54, 4), ('comp_params', 0x58, 8), ('image_version', 0x60, 4),
    ('apu_family_id', 0x64, 4), ('fw_load_addr', 0x68, 4), ('size_image', 0x6C, 4),
    ('size_fw_unsigned', 0x70, 4), ('fw_split_addr', 0x74, 4),
]


def find_nested_images(d, base, alloc, outer_size_image):
    """A directory entry may allocate far more than its own SizeImage and carry
    further signed images inside the slack.  FW_STUB_31 does exactly this: a
    0x210-byte outer stub whose 16-byte body is {tag, 0x210} indexes a full
    0x2C10-byte signed image at +0x210.  Nested headers have HeaderVersion 0
    rather than "$PS1", so scanning for the magic alone misses them.
    """
    out, o = [], (outer_size_image or 0x100)
    o = (o + 0xF) & ~0xF
    while o + 0x100 < alloc:
        ss = struct.unpack_from('<I', d, base + o + 0x14)[0]
        si = struct.unpack_from('<I', d, base + o + 0x6C)[0]
        sig_opt = struct.unpack_from('<I', d, base + o + 0x30)[0]
        # a real header: body plus header plus a 0x100/0x200 signature == SizeImage
        if ss and si and o + si <= alloc and si - ss - 0x100 in (0x100, 0x200) and sig_opt in (0, 1):
            out.append(o)
            o += si
            o = (o + 0xF) & ~0xF
            continue
        o += 0x10
    return out


def parse_fw_header(blob):
    if len(blob) < 0x100:
        return None
    g32 = lambda o: struct.unpack_from('<I', blob, o)[0]
    h = dict(
        header_version=g32(0x10), size_fw_signed=g32(0x14), enc_option=g32(0x18),
        ikek_type=g32(0x1C), sig_option=g32(0x30), sig_alg=g32(0x34),
        sig_key_id=blob[0x38:0x48], comp_option=g32(0x48), sec_patch_level=g32(0x4C),
        uncomp_size=g32(0x50), comp_size=g32(0x54), image_version=blob[0x60:0x64],
        apu_family_id=g32(0x64), fw_load_addr=g32(0x68), size_image=g32(0x6C),
        size_fw_unsigned=g32(0x70), fw_split_addr=g32(0x74),
        fw_type=blob[0x7C], fw_subtype=blob[0x7D], nonce=blob[0x00:0x10],
    )
    magic = blob[0x10:0x14]
    plausible = (magic == b'$PS1') or (
        0 < h['size_fw_signed'] < len(blob) and
        h['size_image'] and h['size_image'] <= len(blob) + 0x200)
    return h if plausible else None


# --------------------------------------------------------------------------- #
# VBIOS / option ROM
# --------------------------------------------------------------------------- #

def parse_option_roms(d, start, limit):
    out, off = [], start
    while off < limit and d[off:off + 2] == b'\x55\xaa':
        pcir_ptr = struct.unpack_from('<H', d, off + 0x18)[0]
        p = off + pcir_ptr
        if d[p:p + 4] != b'PCIR':
            break
        vend, dev = struct.unpack_from('<HH', d, p + 4)
        imglen = struct.unpack_from('<H', d, p + 0x10)[0] * 512
        codetype, indicator = d[p + 0x14], d[p + 0x15]
        strings = [m.group().decode() for m in
                   re.finditer(rb'[ -~]{8,}', d[off:off + min(imglen, 0x1000)])]
        out.append(dict(offset='0x%x' % off, size='0x%x' % imglen,
                        vendor='0x%04x' % vend, device='0x%04x' % dev,
                        code_type=codetype,
                        code_type_name={0: 'x86 legacy', 1: 'OpenFirmware',
                                        3: 'UEFI'}.get(codetype, '?'),
                        last=bool(indicator & 0x80), strings=strings[:24]))
        if indicator & 0x80 or imglen == 0:
            break
        off += imglen
    return out


# --------------------------------------------------------------------------- #
# IP discovery
# --------------------------------------------------------------------------- #

def parse_ip_discovery(b):
    if struct.unpack_from('<I', b, 0)[0] != 0x28211407:
        return None
    vmaj, vmin, csum, bsize = struct.unpack_from('<HHHH', b, 4)
    tnames = ['IP_DISCOVERY', 'GC', 'HARVEST_INFO', 'VCN_INFO', 'MALL_INFO', 'NPS_INFO']
    tables = []
    for i in range(6):
        off, ck, sz, _ = struct.unpack_from('<HHHH', b, 0x0C + i * 8)
        tables.append(dict(name=tnames[i], offset=off, size=sz,
                           checksum_ok=(sum(b[off:off + sz]) & 0xFFFF) == ck if sz else None))
    res = dict(version='%d.%d' % (vmaj, vmin), binary_size=bsize, tables=tables)
    off = tables[0]['offset']
    sig, dver, dsize, tid, ndies = struct.unpack_from('<IHHIH', b, off)
    res['table_id'] = '0x%08x' % tid
    res['ip_header_version'] = dver
    ips = []
    for i in range(ndies):
        die_id, die_off = struct.unpack_from('<HH', b, off + 0x0E + i * 4)
        o = die_off                       # absolute within the discovery binary
        hid, nips = struct.unpack_from('<HH', b, o)
        p = o + 4
        for _ in range(nips):
            hw, inst, nba, maj, minr, rev, sv = struct.unpack_from('<HBBBBBB', b, p)
            bases = struct.unpack_from('<%dI' % nba, b, p + 8) if nba else ()
            ips.append(dict(die=hid, hw_id=hw,
                            name=DISCOVERY_HW.get(hw, 'HWID_%d' % hw),
                            instance=inst, version='%d.%d.%d' % (maj, minr, rev),
                            base_addr=['0x%08x' % x for x in bases]))
            p += 8 + 4 * nba
    res['ip_blocks'] = ips
    go = tables[1]['offset']
    gnames = ['num_se', 'num_wgp0_per_sa', 'num_wgp1_per_sa', 'num_rb_per_se',
              'num_gl2c', 'num_gprs', 'num_max_gs_thds', 'gs_table_depth',
              'gsprim_buff_depth', 'parameter_cache_depth', 'double_offchip_lds_buffer',
              'wave_size', 'max_waves_per_simd', 'max_scratch_slots_per_cu', 'lds_size',
              'num_sc_per_se', 'num_sa_per_se', 'num_packer_per_sc', 'num_gl2a',
              'num_tcp_per_sa']
    res['gc_info'] = dict(zip(gnames, struct.unpack_from('<20I', b, go + 0x0C)))
    return res


# --------------------------------------------------------------------------- #
# driver
# --------------------------------------------------------------------------- #

def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument('image')
    ap.add_argument('-o', '--out', default='spi_out')
    ap.add_argument('--no-extract', action='store_true')
    ap.add_argument('--find-sig-coverage', action='store_true',
                    help='search the image for the bytes each $SGN signature covers '
                         '(slow-ish: a few seconds)')
    args = ap.parse_args()

    d = open(args.image, 'rb').read()
    f = Flash(d)
    out = {'file': os.path.abspath(args.image), 'size': len(d),
           'sha256': hashlib.sha256(d).hexdigest()}
    P = print

    P('=' * 78)
    P('%s  %d bytes  sha256=%s' % (args.image, len(d), out['sha256']))
    P('=' * 78)

    # ---- EFS -------------------------------------------------------------
    efs = f.parse_efs()
    out['efs'] = efs
    P('\n[EFS] @0x00000  magic 0x55AA55AA')
    for k, v in efs['nonzero_fields'].items():
        P('   %-6s = %s' % (k, v))

    # ---- L1 directories --------------------------------------------------
    out['l1_directories'] = []
    l1s = [x for x in (efs['l1_dir_a'], efs['l1_dir_b']) if x]
    for base in l1s:
        dd = f.parse_dir(base)
        if dd is None:
            continue
        out['l1_directories'].append(dd)
        P('\n[L1 $PSP dir] @0x%05x  %d entries  fletcher32 %s'
          % (base, dd['num_entries'], 'OK' if dd['checksum_ok'] else 'BAD'))
        for e in dd['entries']:
            P('   type 0x%02x/%02x %-18s size 0x%-6x -> 0x%06x (mode %d)'
              % (e['type'], e['subprog'], e['name'], e['size'], e['addr'], e['mode']))

    # ---- image slot headers ---------------------------------------------
    ish_offsets = []
    for dd in out['l1_directories']:
        for e in dd['entries']:
            if e['type'] in (0x3C, 0x3D):
                ish_offsets.append(e['addr'])
    ish_offsets = sorted(set(ish_offsets))
    out['image_slots'] = [f.parse_ish(o) for o in ish_offsets]
    P('\n[Image slot headers]')
    for s in out['image_slots']:
        P('   @0x%05x  prio=%d retry=%d ctr=%s slot=%s psp_id=%s size=%s 2nd=%s  csum %s'
          % (s['offset'], s['boot_priority'], s['update_retry_count'], s['counter'],
             s['image_slot_addr'], s['psp_id'], s['slot_size'], s['second_ptr'],
             'OK' if s['checksum_ok'] else 'BAD'))

    # ---- $ROM config pages ----------------------------------------------
    out['rom_cfg_pages'] = []
    for o in (efs['rom_cfg_a'], efs['rom_cfg_b']):
        if o and d[o:o + 4] == ROM_COOKIE:
            out['rom_cfg_pages'].append(dict(offset='0x%x' % o,
                                             sha256=hashlib.sha256(d[o:o + 0x1000]).hexdigest()))
    if out['rom_cfg_pages']:
        P('\n[$ROM board/ASIC config pages] ' +
          ', '.join(p['offset'] for p in out['rom_cfg_pages']) +
          '  (identical)' if len({p['sha256'] for p in out['rom_cfg_pages']}) == 1 else '')

    # ---- $SGN ------------------------------------------------------------
    sgn = None
    for o in range(0, len(d) - 4, 0x1000):
        if d[o:o + 4] == SGN_COOKIE:
            sgn = f.parse_sgn(o)
            break
    out['sgn'] = sgn
    if sgn:
        P('\n[$SGN image signature table] @0x%05x  %d signatures'
          % (sgn['offset'], sgn['count']))

    # ---- slots -----------------------------------------------------------
    slots = []
    for s in out['image_slots']:
        slots.append(int(s['image_slot_addr'], 16))
    if not slots:
        slots = [0x20000]
    P('\n[Image slots] ' + ', '.join('0x%06x' % s for s in slots))
    if len(slots) == 2:
        sz = int(out['image_slots'][0]['slot_size'], 16)
        same = d[slots[0]:slots[0] + sz] == d[slots[1]:slots[1] + sz]
        P('   slot A and slot B are %s (0x%x bytes)'
          % ('byte-identical' if same else 'DIFFERENT', sz))
        out['slots_identical'] = same

    # ---- L2 directory ----------------------------------------------------
    base = slots[0]
    l2 = f.parse_dir(base)
    if l2 is None:
        P('no $PSP directory at the image-slot address 0x%06x' % base)
        return 1
    out['l2_directory'] = l2
    P('\n[L2 $PSP dir] @0x%06x  %d entries  fletcher32 %s'
      % (base, l2['num_entries'], 'OK' if l2['checksum_ok'] else 'BAD'))

    # resolve + extract
    if not args.no_extract:
        for sub in ('raw', 'payload', 'keys'):
            os.makedirs(os.path.join(args.out, sub), exist_ok=True)

    # pass 1: key tokens + KDB so signature checks can run in pass 2
    for e in l2['entries']:
        e['file_offset'] = e['addr'] if e['mode'] == 0 else base + e['addr']
        blob = d[e['file_offset']:e['file_offset'] + e['size']]
        if e['type'] in (0x39, 0x32, 0x3A) and e['subprog'] == 0x01 and e['size'] >= 0x340:
            t = f.parse_key_token(e['file_offset'])
            if t:
                f.keys[t['key_id']] = dict(n=t['n'], e=t['e'], bits=t['mod_bits'],
                                           source='token type 0x%02x' % e['type'])
                e['key_token'] = dict(key_id=t['key_id'].hex(),
                                      cert_key_id=t['cert_key_id'].hex(),
                                      usage=t['usage'], bits=t['mod_bits'],
                                      self_signed=t['key_id'] == t['cert_key_id'])
        if e['type'] == 0x4E:
            h = parse_fw_header(blob)
            if h:
                e['kdb_keys'] = f.load_kdb(blob[0x100:0x100 + h['size_fw_signed']])

    # pass 2: headers, payloads, signature verification
    P('\n%-4s %-9s %-20s %-9s %-9s %-13s %-4s %-5s %s'
      % ('idx', 'type/sub', 'name', 'offset', 'size', 'version', 'spl', 'flags', 'signature'))
    for e in l2['entries']:
        fo, size = e['file_offset'], e['size']
        blob = d[fo:fo + size]
        e['sha256'] = hashlib.sha256(blob).hexdigest()
        if not args.no_extract:
            open(os.path.join(args.out, 'raw', 'e%02d_t%02x_s%02x_%s.bin'
                              % (e['index'], e['type'], e['subprog'], e['name'])), 'wb').write(blob)
        h = parse_fw_header(blob) if size >= 0x100 else None
        flags = sig = ver = ''
        spl = ''
        if h:
            e['header'] = {k: (v.hex() if isinstance(v, bytes) else v)
                           for k, v in h.items()}
            ver = ver_str(h['image_version'])
            spl = str(h['sec_patch_level'])
            flags = ('E' if h['enc_option'] else '_') + \
                    ('S' if h['sig_option'] else '_') + \
                    ('C' if h['comp_option'] else '_')
            body = blob[0x100:0x100 + (h['comp_size'] if h['comp_option'] and h['comp_size']
                                       else h['size_fw_signed'])]
            if h['comp_option'] and h['comp_size']:
                try:
                    body = zlib.decompress(body)
                except zlib.error:
                    pass
            if not args.no_extract and body:
                open(os.path.join(args.out, 'payload', 'e%02d_t%02x_s%02x_%s.bin'
                                  % (e['index'], e['type'], e['subprog'], e['name'])),
                     'wb').write(body)
            k = f.keys.get(h['sig_key_id'])
            if k:
                hf, hl = hash_for_bits(k['bits'])
                slen = k['bits'] // 8
                s = blob[0x100 + h['size_fw_signed']: 0x100 + h['size_fw_signed'] + slen]
                r = pss_open(s, k['n'], k['e'], hf, hl)
                ok = bool(r) and pss_check(r[0], r[1], blob[:0x100 + h['size_fw_signed']], hf)
                sig = '%s %s/%s' % ('OK ' if ok else 'BAD', h['sig_key_id'].hex()[:16],
                                    'sha%d' % (hl * 8))
                e['signature_valid'] = ok
                e['signing_key'] = h['sig_key_id'].hex()
            else:
                sig = '?key %s' % h['sig_key_id'].hex()[:16]
        elif not args.no_extract and 'key_token' in e:
            pass
        P('%-4d %02x/%02x     %-20s 0x%06x  0x%06x  %-13s %-4s %-5s %s'
          % (e['index'], e['type'], e['subprog'], e['name'], fo, size, ver, spl, flags, sig))

        # nested signed images hiding in the entry's unused allocation
        e['nested'] = []
        for j, noff in enumerate(find_nested_images(d, fo, size, h['size_image'] if h else 0)):
            # skip anything that is already its own directory entry -- several
            # entries physically nest inside another entry's allocation
            # (VBL_DATA_TABLE sits inside VBIOS_IMAGE, for one)
            if any(x['file_offset'] == fo + noff for x in l2['entries']):
                continue
            nb = d[fo + noff:fo + size - noff]
            nh = parse_fw_header(nb)
            if not nh:
                continue
            nname = '%s_nested%d' % (e['name'], j)
            nk = f.keys.get(nh['sig_key_id'])
            nsig = '?key %s' % nh['sig_key_id'].hex()[:16]
            nok = None
            if nk:
                hf, hl = hash_for_bits(nk['bits'])
                slen = nk['bits'] // 8
                s_ = nb[0x100 + nh['size_fw_signed']: 0x100 + nh['size_fw_signed'] + slen]
                r = pss_open(s_, nk['n'], nk['e'], hf, hl)
                nok = bool(r) and pss_check(r[0], r[1], nb[:0x100 + nh['size_fw_signed']], hf)
                nsig = '%s %s/sha%d' % ('OK ' if nok else 'BAD', nh['sig_key_id'].hex()[:16], hl * 8)
            e['nested'].append(dict(offset='+0x%x' % noff, file_offset='0x%x' % (fo + noff),
                                    size_image='0x%x' % nh['size_image'],
                                    size_fw_signed='0x%x' % nh['size_fw_signed'],
                                    fw_type='0x%02x' % nh['fw_type'],
                                    version=ver_str(nh['image_version']),
                                    signing_key=nh['sig_key_id'].hex(),
                                    signature_valid=nok))
            if not args.no_extract:
                body = nb[0x100:0x100 + nh['size_fw_signed']]
                open(os.path.join(args.out, 'payload', nname + '.bin'), 'wb').write(body)
                open(os.path.join(args.out, 'raw', nname + '.bin'), 'wb').write(nb[:nh['size_image']])
            P('  \\_ nested at +0x%-5x %-16s 0x%06x  0x%06x  %-13s      %-5s %s'
              % (noff, nname, fo + noff, nh['size_image'], ver_str(nh['image_version']),
                 ('E' if nh['enc_option'] else '-') + ('S' if nh['sig_option'] else '-')
                 + ('C' if nh['comp_option'] else '-'), nsig))

    # key summary
    P('\n[Public keys] %d total' % len(f.keys))
    for kid, k in sorted(f.keys.items(), key=lambda x: x[0].hex()):
        P('   %s  RSA-%-4d e=0x%x  (%s)' % (kid.hex(), k['bits'], k['e'], k['source']))
        if not args.no_extract:
            try:
                from Crypto.PublicKey import RSA
                pem = RSA.construct((k['n'], k['e'])).export_key('PEM')
                open(os.path.join(args.out, 'keys', '%s_%d.pem' % (kid.hex(), k['bits'])),
                     'wb').write(pem)
            except ImportError:
                pass
    out['keys'] = {kid.hex(): dict(bits=k['bits'], e=k['e'], source=k['source'])
                   for kid, k in f.keys.items()}

    # ---- $SGN / type 0x41 signatures ------------------------------------
    img_key = None
    for e in l2['entries']:
        if e.get('key_token', {}).get('usage') == 8:
            img_key = f.keys[bytes.fromhex(e['key_token']['key_id'])]
    if sgn and img_key:
        P('\n[$SGN signatures]  key %s (RSA-%d, usage 8), PSS-SHA256, stored LE'
          % ([k.hex() for k, v in f.keys.items() if v is img_key][0][:16], img_key['bits']))
        targets = [('SGN%d @0x%05x' % (s['id'], s['offset_int']), s['offset_int'])
                   for s in sgn['entries']]
        for e in l2['entries']:
            if e['type'] == 0x41:
                targets.append(('entry 0x41 @0x%05x' % e['file_offset'], e['file_offset']))
        cov = {}
        if args.find_sig_coverage:
            cands = []
            for st in list(range(0, min(len(d), 0x24000), 0x10)) + \
                      [slots[0], slots[-1], 0, 0x10000, 0x11000]:
                cands.append(st)
            lens = [0x20, 0x30, 0x40, 0x54, 0x60, 0x80, 0xC0, 0x100, 0x210,
                    0x1000, 0xE0000]
        for nm, off in targets:
            r = pss_open(d[off:off + 0x100][::-1], img_key['n'], img_key['e'],
                         hashlib.sha256, 32)
            covers = ''
            if r and args.find_sig_coverage:
                h, salt = r
                for st in cands:
                    for l in lens:
                        if st + l <= len(d) and pss_check(h, salt, d[st:st + l], hashlib.sha256):
                            covers = 'covers flash[0x%06x:0x%06x]' % (st, st + l)
                            break
                    if covers:
                        break
                if not covers:
                    covers = 'coverage not found in image'
            P('   %-22s  %s  %s' % (nm, 'PSS OK' if r else 'not PSS', covers))
            cov[nm] = covers
        out['sgn_coverage'] = cov

    # ---- VBIOS -----------------------------------------------------------
    vb = None
    for e in l2['entries']:
        if e['type'] == 0x37:
            vb = e['file_offset']
    if vb is not None:
        roms = parse_option_roms(d, vb, len(d))
        out['option_roms'] = roms
        P('\n[PCI option ROM chain] starting at 0x%06x' % vb)
        for r in roms:
            P('   %s size %s  %s:%s  %s%s' % (r['offset'], r['size'], r['vendor'],
                                              r['device'], r['code_type_name'],
                                              '  (last)' if r['last'] else ''))
            for s in r['strings'][:12]:
                P('        %s' % s)

    # ---- IP discovery ----------------------------------------------------
    for e in l2['entries']:
        if e['name'] == 'IP_DISCOVERY' and 'header' in e:
            body = d[e['file_offset'] + 0x100:
                     e['file_offset'] + 0x100 + e['header']['size_fw_signed']]
            ipd = parse_ip_discovery(body)
            if ipd:
                out['ip_discovery'] = ipd
                P('\n[IP discovery] binary v%s, table id %s, %d IP blocks'
                  % (ipd['version'], ipd['table_id'], len(ipd['ip_blocks'])))
                for ip in ipd['ip_blocks']:
                    P('   %-13s inst=%-2d v%-9s %s'
                      % (ip['name'], ip['instance'], ip['version'],
                         ' '.join(ip['base_addr'])))
                g = ipd['gc_info']
                cus = g['num_se'] * g['num_sa_per_se'] * (g['num_wgp0_per_sa'] + g['num_wgp1_per_sa']) * 2
                P('   GC: %d SE x %d SA x %d WGP  => %d CUs, %d GL2C, wave%d'
                  % (g['num_se'], g['num_sa_per_se'],
                     g['num_wgp0_per_sa'] + g['num_wgp1_per_sa'], cus,
                     g['num_gl2c'], g['wave_size']))

    # ---- unsigned / writable regions -------------------------------------
    P('\n[Unsigned regions and known firmware write targets]')
    signed_spans = []
    for nm, c in (out.get('sgn_coverage') or {}).items():
        m = re.match(r'covers flash\[0x([0-9a-f]+):0x([0-9a-f]+)\]', c or '')
        if m:
            signed_spans.append((int(m.group(1), 16), int(m.group(2), 16), nm))
    signed_spans.sort()
    for a, b2, nm in signed_spans:
        P('   signed    0x%06x..0x%06x  by %s' % (a, b2, nm))
    for off, ln, what in ((0x7000, 0xC, '"$AMDVBFL" marker, written by vbios_flash_verify'),
                          (0x1EF000, 0x48, '"SPIS" record, written by vbios_flash_verify')):
        blank = d[off:off + ln].count(0) == ln
        covered = any(a <= off and off + ln <= b2 for a, b2, _ in signed_spans)
        P('   written   0x%06x..0x%06x  %s%s%s'
          % (off, off + ln, what,
             '' if covered else '  [unsigned]',
             '  [blank]' if blank else ''))
    # Only meaningful once signature coverage has actually been resolved --
    # without --find-sig-coverage every entry would look unsigned.
    if not signed_spans:
        P('   (run with --find-sig-coverage to resolve which regions are signed)')
    for e in l2['entries'] if signed_spans else []:
        a, b2 = e['file_offset'], e['file_offset'] + e['size']
        if not any(sa <= a and b2 <= sb for sa, sb, _ in signed_spans):
            blk = d[a:b2]
            P('   unsigned  0x%06x..0x%06x  directory entry %s%s'
              % (a, b2, e['name'], '  [all 0x00]' if blk.count(0) == len(blk) else ''))
    # 0xFF-heavy sectors.  0xFF only means "erased" if it is NOT inside a span
    # that a signature covers -- a signed 0xFF run is build padding.
    ff = [o for o in range(0, len(d), 0x1000) if d[o:o + 0x1000].count(0xFF) > 0x400]
    if not ff:
        P('   0xFF fill: none anywhere in the image')
    delta = (slots[1] - slots[0]) if len(slots) == 2 else 0
    for o in ff:
        inside = [nm for a, b2, nm in signed_spans if a <= o and o + 0x1000 <= b2]
        if inside:
            why = 'inside %s -> signed build padding, not erase state' % inside[0]
        elif delta and d[o:o + 0x1000] == d[o - delta:o - delta + 0x1000] and \
                any(a <= o - delta and o - delta + 0x1000 <= b2 for a, b2, _ in signed_spans):
            why = 'mirror of signed 0x%06x -> build padding, not erase state' % (o - delta)
        elif not signed_spans:
            why = 'signature coverage not resolved (pass --find-sig-coverage)'
        else:
            why = 'NOT covered by a signature -> possible erase state'
        P('   0xFF fill  0x%06x..0x%06x  %s' % (o, o + 0x1000, why))

    # ---- coverage map ----------------------------------------------------
    P('\n[Coverage]')
    spans = [(0, 0x54, 'EFS header')]
    for dd in out['l1_directories']:
        spans.append((dd['base'], 0x10 + dd['num_entries'] * 0x10, 'L1 $PSP dir'))
    for s in out['image_slots']:
        spans.append((s['offset'], 0x20, 'image slot header'))
    for p in out['rom_cfg_pages']:
        spans.append((int(p['offset'], 16), 0x1000, '$ROM config page'))
    if sgn:
        spans.append((sgn['offset'], 0x10 + sgn['count'] * 0x10, '$SGN table'))
        for s in sgn['entries']:
            spans.append((s['offset_int'], s['size_int'], 'image signature #%d' % s['id']))
    spans.append((base, 0x10 + l2['num_entries'] * 0x10, 'L2 $PSP dir'))
    for e in l2['entries']:
        spans.append((e['file_offset'], e['size'], 'entry %d %s' % (e['index'], e['name'])))
    if len(slots) == 2:
        spans.append((slots[1], int(out['image_slots'][0]['slot_size'], 16), 'slot B (mirror)'))
    spans.sort()
    prev, gaps = 0, []
    for s, l, nm in spans:
        if s > prev:
            g = d[prev:s]
            kind = 'zero' if g.count(0) == len(g) else ('0xff' if g.count(0xff) == len(g) else 'DATA')
            gaps.append((prev, s, kind))
            P('   gap  0x%06x..0x%06x (0x%06x) %s' % (prev, s, s - prev, kind))
        prev = max(prev, s + l)
    if prev < len(d):
        g = d[prev:]
        kind = 'zero' if g.count(0) == len(g) else ('0xff' if g.count(0xff) == len(g) else 'DATA')
        P('   gap  0x%06x..0x%06x (0x%06x) %s' % (prev, len(d), len(d) - prev, kind))

    if not args.no_extract:
        with open(os.path.join(args.out, 'summary.json'), 'w') as fh:
            json.dump(out, fh, indent=2, default=str)
        P('\nwrote %s/{raw,payload,keys}/ and summary.json' % args.out)
    return 0


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