#!/usr/bin/env python3
# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2026 Open-IOV contributors
#
# Generate the kernel bindata headers for the real-GSP boot stages from the
# public open kernel-modules generated bindata archives.
#
# The archives are C files holding raw-deflate compressed byte arrays. This
# tool extracts the production booter-load image (runs on SEC2 to load,
# verify and start the GSP-RM RISC-V image) and the GSP-RM boot image (the
# "FMC" 4 KB bootloader), applies the production signature patch the driver
# applies at runtime, verifies sizes/headers, and cross-checks the booter
# image against the distribution firmware blob when present.
#
# Usage: python3 tools/gen_gsp_bindata.py [--check]

import argparse
import hashlib
import re
import struct
import sys
import zlib
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
GENERATED = ROOT / "nvidia-open-gpu-kernel-modules" / "src" / "nvidia" / "generated"
KERNEL = ROOT / "kernel"

BOOTER_ARCHIVE = GENERATED / "g_bindata_kgspGetBinArchiveBooterLoadUcode_TU102.c"
FMC_ARCHIVE = GENERATED / "g_bindata_kgspGetBinArchiveGspRmBoot_TU102.c"
DEFAULT_FW_BOOTER = Path("/lib/firmware/nvidia/tu102/gsp/booter_load-570.144.bin")

# Booter-load production image expectations (reference ABI, Turing).
BOOTER_IMAGE_SIZE = 59136
BOOTER_SIG_OFFSET = 0x8700
BOOTER_SIG = bytes.fromhex("3260436798897b06e85f1aa871435dcb")
BOOTER_HEADER = (0x0, 0x100, 0x8500, 0x6200, 0x1, 0x100, 0x8400, 0x100, 0x0)
BOOTER_PATCH_META = (0x0, 0x1, 0xd)

# GSP-RM boot image expectations (FMC).
FMC_IMAGE_SIZE = 4096
FMC_DESC_SIZE = 76
FMC_DESC_VERSION = 4
FMC_DESC_WORDS = [FMC_DESC_VERSION, 0, 0x488, 0x488, 0x10] + [0] * 14

EXPECTED_BOOTER_SHA256 = "1de31591beae01d3a59778d9066532c6e808029258927d270c31193265c3305f"
EXPECTED_FMC_SHA256 = "8765694cd38beaa4e2b8b1b83a31fb7aee6cb9d8b7c6bc7a86e6d931a3a056c2"

FIELD_RE = re.compile(r'// FUNCTION: .*\("([A-Z0-9_]+)"\)')
COMPRESSION_RE = re.compile(r"// COMPRESSION: (YES|NO)")


def parse_archive(path):
    """Return {label: (compressed, bytes)} from a generated bindata C file."""
    entries = {}
    label = None
    compressed = None
    data = bytearray()
    with open(path) as fh:
        for line in fh:
            stripped = line.strip()
            m = FIELD_RE.match(stripped)
            if m:
                if label is not None:
                    entries[label] = (compressed, bytes(data))
                label, compressed, data = m.group(1), None, bytearray()
                continue
            if label is None:
                continue
            m = COMPRESSION_RE.match(stripped)
            if m:
                compressed = m.group(1) == "YES"
                continue
            if stripped.startswith("//") or stripped == "":
                continue
            if stripped == "};":
                entries[label] = (compressed, bytes(data))
                label = None
                continue
            if stripped.startswith("{") or stripped.startswith("static") or \
                    stripped.startswith("BINDATA") or stripped.startswith("#"):
                continue
            for token in re.findall(r"0x[0-9a-fA-F]{1,2}", line):
                data.append(int(token, 16))
    if label is not None:
        entries[label] = (compressed, bytes(data))
    return entries


def extract(entries, label, want_compressed=None):
    if label not in entries:
        raise SystemExit(f"missing bindata label {label}")
    compressed, raw = entries[label]
    if want_compressed is not None and compressed != want_compressed:
        raise SystemExit(
            f"{label}: compression flag {compressed} != expected {want_compressed}"
        )
    if compressed:
        return zlib.decompress(raw, -15)
    return raw


def emit_header(path, guard, name, image, note_lines):
    sha = hashlib.sha256(image).hexdigest()
    out = []
    out.append("/* SPDX-License-Identifier: GPL-2.0 */")
    out.append("/*")
    for line in note_lines:
        out.append(f" * {line}" if line else " *")
    out.append(" *")
    out.append(" * Generated from the open kernel-modules bindata archives; do")
    out.append(" * not edit by hand.")
    out.append(" *")
    out.append(f" * image size: {len(image)} bytes")
    out.append(f" * sha256: {sha}")
    out.append(" */")
    out.append("")
    out.append(f"#ifndef {guard}")
    out.append(f"#define {guard}")
    out.append("")
    out.append("#include <linux/types.h>")
    out.append("")
    out.append(f"static const u8 {name}[{len(image)}] = {{")
    for i in range(0, len(image), 12):
        row = ", ".join(f"0x{b:02x}" for b in image[i : i + 12])
        out.append(f"\t{row},")
    out.append("};")
    out.append("")
    out.append(f"#endif /* {guard} */")
    path.write_text("\n".join(out) + "\n")
    return sha


def cross_check_booter(fw_path, image):
    if not fw_path.exists():
        raise SystemExit(f"cross-check firmware not found: {fw_path}")
    fw = fw_path.read_bytes()
    if len(fw) < 0x88 + len(image):
        raise SystemExit(f"cross-check firmware too small: {len(fw)} bytes")
    if len(fw) != 0x88 + len(image):
        raise SystemExit(
            f"cross-check firmware size {len(fw)} != 0x88 + {len(image)}"
        )
    for i, byte in enumerate(image):
        if BOOTER_SIG_OFFSET <= i < BOOTER_SIG_OFFSET + len(BOOTER_SIG):
            continue
        if fw[0x88 + i] != byte:
            raise SystemExit(f"cross-check mismatch at image offset 0x{i:x}")
    if fw[0x88 + BOOTER_SIG_OFFSET : 0x88 + BOOTER_SIG_OFFSET + 16] != bytes(16):
        # The distribution blob carries a zeroed patch slot; anything else
        # means this check is no longer comparing what we think it is.
        raise SystemExit("cross-check firmware has a non-zero signature patch slot")
    return True


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--check",
        action="store_true",
        help="verify everything but do not rewrite the headers",
    )
    parser.add_argument(
        "--firmware-booter",
        type=Path,
        default=DEFAULT_FW_BOOTER,
        help=f"distribution booter blob for the cross-check (default {DEFAULT_FW_BOOTER})",
    )
    args = parser.parse_args()

    if not BOOTER_ARCHIVE.exists() or not FMC_ARCHIVE.exists():
        raise SystemExit(f"generated bindata sources not found under {GENERATED}")

    booter_entries = parse_archive(BOOTER_ARCHIVE)
    fmc_entries = parse_archive(FMC_ARCHIVE)

    image = extract(booter_entries, "BINDATA_LABEL_IMAGE_PROD", want_compressed=True)
    header = extract(booter_entries, "BINDATA_LABEL_HEADER_PROD", want_compressed=True)
    sig = extract(booter_entries, "BINDATA_LABEL_SIG_PROD", want_compressed=False)
    patch_loc = extract(booter_entries, "BINDATA_LABEL_PATCH_LOC", want_compressed=False)
    patch_sig = extract(booter_entries, "BINDATA_LABEL_PATCH_SIG", want_compressed=False)
    patch_meta = extract(booter_entries, "BINDATA_LABEL_PATCH_META", want_compressed=False)
    num_sigs = extract(booter_entries, "BINDATA_LABEL_NUM_SIGS", want_compressed=False)

    if len(image) != BOOTER_IMAGE_SIZE:
        raise SystemExit(f"booter image size {len(image)} != {BOOTER_IMAGE_SIZE}")
    if len(header) != 36:
        raise SystemExit(f"booter header size {len(header)} != 36")
    fields = struct.unpack("<9I", header)
    if fields != BOOTER_HEADER:
        raise SystemExit(f"booter header {fields} != {BOOTER_HEADER}")
    if len(sig) != 16:
        raise SystemExit(f"booter signature size {len(sig)} != 16")
    if patch_loc != struct.pack("<I", BOOTER_SIG_OFFSET):
        raise SystemExit(f"booter patch location {patch_loc.hex()} unexpected")
    if patch_sig != bytes(4):
        raise SystemExit(f"booter patch signature {patch_sig.hex()} != zero")
    if struct.unpack("<3I", patch_meta) != BOOTER_PATCH_META:
        raise SystemExit(f"booter patch meta {patch_meta.hex()} unexpected")
    if struct.unpack("<I", num_sigs)[0] != 1:
        raise SystemExit("booter archive must carry exactly one signature")

    image = bytearray(image)
    if bytes(image[BOOTER_SIG_OFFSET : BOOTER_SIG_OFFSET + 16]) != bytes(16):
        raise SystemExit("booter signature slot is not empty before patching")
    image[BOOTER_SIG_OFFSET : BOOTER_SIG_OFFSET + 16] = sig
    if bytes(image[BOOTER_SIG_OFFSET : BOOTER_SIG_OFFSET + 16]) != BOOTER_SIG:
        raise SystemExit("booter signature patch did not take")
    booter_sha = hashlib.sha256(bytes(image)).hexdigest()

    fmc_image = extract(fmc_entries, "BINDATA_LABEL_UCODE_IMAGE", want_compressed=True)
    fmc_desc = extract(fmc_entries, "BINDATA_LABEL_UCODE_DESC", want_compressed=True)
    if len(fmc_image) != FMC_IMAGE_SIZE:
        raise SystemExit(f"FMC image size {len(fmc_image)} != {FMC_IMAGE_SIZE}")
    if len(fmc_desc) != FMC_DESC_SIZE:
        raise SystemExit(f"FMC descriptor size {len(fmc_desc)} != {FMC_DESC_SIZE}")
    words = list(struct.unpack("<19I", fmc_desc))
    if words != FMC_DESC_WORDS:
        raise SystemExit(f"FMC descriptor words {[hex(w) for w in words]} unexpected")
    fmc_sha = hashlib.sha256(fmc_image).hexdigest()

    if booter_sha != EXPECTED_BOOTER_SHA256:
        raise SystemExit(
            f"booter sha256 {booter_sha} != expected {EXPECTED_BOOTER_SHA256}"
        )
    if fmc_sha != EXPECTED_FMC_SHA256:
        raise SystemExit(f"FMC sha256 {fmc_sha} != expected {EXPECTED_FMC_SHA256}")

    cross_check_booter(args.firmware_booter, bytes(image))

    print(f"booter-load image: {len(image)} bytes sha256={booter_sha}")
    print(f"  header: osCode 0x100, appCode 0x8400@0x100, data 0x6200@0x8500, numApps 1")
    print(f"  signature patched at 0x{BOOTER_SIG_OFFSET:x}: {BOOTER_SIG.hex()}")
    print(f"  cross-check vs {args.firmware_booter}: OK (except the 16-byte patch slot)")
    print(f"GSP-RM boot (FMC) image: {len(fmc_image)} bytes sha256={fmc_sha}")
    print(f"  descriptor v{FMC_DESC_VERSION}: bootloader 0x488@0, param 0x10@0x488, appVersion 0")

    if args.check:
        print("--check: headers not rewritten")
        return 0

    booter_header_path = KERNEL / "gsp_phys_booter_load_img.h"
    fmc_header_path = KERNEL / "gsp_phys_gsp_rm_boot_img.h"
    emit_header(
        booter_header_path,
        "GSP_PHYS_BOOTER_LOAD_IMG_H",
        "gsp_phys_booter_load_image",
        bytes(image),
        [
            "Booter Load image (production, Turing) executed on SEC2 to load,",
            "verify and start the GSP-RM RISC-V image. The 16-byte production",
            "signature is already patched at its bindata patch location",
            "(offset 0x8700).",
        ],
    )
    emit_header(
        fmc_header_path,
        "GSP_PHYS_GSP_RM_BOOT_IMG_H",
        "gsp_phys_gsp_rm_boot_image",
        fmc_image,
        [
            "GSP-RM boot image (the 4 KB FMC bootloader, production, Turing)",
            "consumed by the Booter Load ucode through the WPR metadata",
            "sysmemAddrOfBootloader field.",
        ],
    )
    print(f"wrote {booter_header_path}")
    print(f"wrote {fmc_header_path}")
    return 0


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