#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-2.0-or-later

"""Create an isolated A6 IMG3 lab identity and a derived SecureROM image."""

import argparse
import hashlib
import json
import os
from pathlib import Path
import shutil
import stat
import subprocess
import sys


ROM_SIZE = 0x10000
ROOT_CERT_OFFSET = 0xBFD0
ROOT_CERT_SLOT_SIZE = 1215
ROOT_SUBJECT = (
    "/C=US/O=Apple Inc./OU=Apple Certification Authority/CN=Apple Root CA"
)
INTERMEDIATE_SUBJECT = (
    "/C=ZZ/O=QEMU A6 Lab/OU=Secure Boot Research/"
    "CN=Apple Secure Boot Certification Authority"
)
IMG3_LEAF_SUBJECT = (
    "/C=ZZ/O=QEMU A6 Lab/OU=Secure Boot Research/"
    "CN=A6-Darwin-Prod-CEPO10-SDOM3-Lab"
)
TICKET_LEAF_SUBJECT = (
    "/C=ZZ/O=QEMU A6 Lab/OU=Secure Boot Research/"
    "CN=H5P-Darwin-Prod-CEPO1-Ticket-DataCenter"
)


class LabIdentityError(RuntimeError):
    pass


def run(command, cwd):
    try:
        subprocess.run(command, cwd=cwd, check=True)
    except FileNotFoundError as error:
        raise LabIdentityError(
            f"required program not found: {command[0]}"
        ) from error
    except subprocess.CalledProcessError as error:
        raise LabIdentityError(
            f"command failed with exit status {error.returncode}: "
            + " ".join(command)
        ) from error


def openssl(*arguments, cwd):
    run(["openssl", *arguments], cwd)


def sha256(data):
    return hashlib.sha256(data).hexdigest()


def der_object_size(data, offset):
    if offset >= len(data) or data[offset] != 0x30:
        raise LabIdentityError(
            f"no DER SEQUENCE at SecureROM offset 0x{offset:x}"
        )
    if offset + 2 > len(data):
        raise LabIdentityError("truncated DER length")
    first = data[offset + 1]
    if first < 0x80:
        return 2 + first
    length_bytes = first & 0x7F
    if not 1 <= length_bytes <= 4 or offset + 2 + length_bytes > len(data):
        raise LabIdentityError("invalid DER length")
    payload_size = int.from_bytes(
        data[offset + 2:offset + 2 + length_bytes], "big"
    )
    return 2 + length_bytes + payload_size


def ensure_clean_output(path):
    if path.exists():
        if not path.is_dir():
            raise LabIdentityError(
                f"output exists and is not a directory: {path}"
            )
        if any(path.iterdir()):
            raise LabIdentityError(
                f"output directory is not empty: {path}; choose a new directory"
            )
    else:
        path.mkdir(parents=True, mode=0o700)


def create_identity(source_rom, output, config):
    source = source_rom.read_bytes()
    if len(source) != ROM_SIZE:
        raise LabIdentityError(
            f"SecureROM must be exactly {ROM_SIZE} bytes, got {len(source)}"
        )
    embedded_size = der_object_size(source, ROOT_CERT_OFFSET)
    if embedded_size != ROOT_CERT_SLOT_SIZE:
        raise LabIdentityError(
            "unexpected embedded trust anchor size at 0xbfd0: "
            f"{embedded_size}, expected {ROOT_CERT_SLOT_SIZE}"
        )

    ensure_clean_output(output)
    os.chmod(output, 0o700)

    openssl(
        "genpkey", "-algorithm", "RSA", "-pkeyopt", "rsa_keygen_bits:2048",
        "-out", "root-key.pem", cwd=output,
    )
    openssl(
        "req", "-new", "-x509", "-sha1", "-days", "3650",
        "-set_serial", "0x02", "-key", "root-key.pem",
        # iBSS parses the complete fixed 1215-byte root slot.  Its public
        # subject layout plus the standard policy extension in the config
        # preserves that DER size while all trust material remains lab-owned.
        "-subj", ROOT_SUBJECT,
        "-config", str(config), "-extensions", "v3_root",
        "-out", "root.pem", cwd=output,
    )

    openssl(
        "genpkey", "-algorithm", "RSA", "-pkeyopt", "rsa_keygen_bits:2048",
        "-out", "intermediate-key.pem", cwd=output,
    )
    openssl(
        "req", "-new", "-sha256", "-key", "intermediate-key.pem",
        "-subj", INTERMEDIATE_SUBJECT,
        "-config", str(config),
        "-out", "intermediate.csr", cwd=output,
    )
    openssl(
        "x509", "-req", "-sha1", "-days", "3650",
        "-set_serial", "0xA60010", "-in", "intermediate.csr",
        "-CA", "root.pem", "-CAkey", "root-key.pem",
        "-extfile", str(config), "-extensions", "v3_intermediate",
        "-out", "intermediate.pem", cwd=output,
    )

    openssl(
        "genpkey", "-algorithm", "RSA", "-pkeyopt", "rsa_keygen_bits:1024",
        "-out", "img3-leaf-key.pem", cwd=output,
    )
    openssl(
        "req", "-new", "-sha256", "-key", "img3-leaf-key.pem",
        "-subj", IMG3_LEAF_SUBJECT,
        "-config", str(config),
        "-out", "img3-leaf.csr", cwd=output,
    )
    openssl(
        "x509", "-req", "-sha1", "-days", "3650",
        "-set_serial", "0xA60105", "-in", "img3-leaf.csr",
        "-CA", "intermediate.pem", "-CAkey", "intermediate-key.pem",
        "-extfile", str(config), "-extensions", "v3_img3_leaf",
        "-out", "img3-leaf.pem", cwd=output,
    )

    openssl(
        "genpkey", "-algorithm", "RSA", "-pkeyopt", "rsa_keygen_bits:1024",
        "-out", "ticket-leaf-key.pem", cwd=output,
    )
    openssl(
        "req", "-new", "-sha256", "-key", "ticket-leaf-key.pem",
        # iBSS compares this public ticket-class name while certificate trust
        # still comes exclusively from the generated lab root and keys.
        "-subj", TICKET_LEAF_SUBJECT,
        "-config", str(config),
        "-out", "ticket-leaf.csr", cwd=output,
    )
    openssl(
        "x509", "-req", "-sha1", "-days", "3650",
        "-set_serial", "0xA6010B", "-in", "ticket-leaf.csr",
        "-CA", "intermediate.pem", "-CAkey", "intermediate-key.pem",
        "-extfile", str(config), "-extensions", "v3_ticket_leaf",
        "-out", "ticket-leaf.pem", cwd=output,
    )

    for name in ("root", "intermediate", "img3-leaf", "ticket-leaf"):
        openssl(
            "x509", "-in", f"{name}.pem", "-outform", "DER",
            "-out", f"{name}.der", cwd=output,
        )

    openssl(
        "verify", "-auth_level", "0", "-ignore_critical",
        "-CAfile", "root.pem",
        "-untrusted", "intermediate.pem", "img3-leaf.pem", cwd=output,
    )
    openssl(
        "verify", "-auth_level", "0", "-ignore_critical",
        "-CAfile", "root.pem",
        "-untrusted", "intermediate.pem", "ticket-leaf.pem", cwd=output,
    )

    root_der = (output / "root.der").read_bytes()
    intermediate_der = (output / "intermediate.der").read_bytes()
    leaf_der = (output / "img3-leaf.der").read_bytes()
    ticket_leaf_der = (output / "ticket-leaf.der").read_bytes()
    if len(root_der) > ROOT_CERT_SLOT_SIZE:
        raise LabIdentityError(
            f"lab root DER is {len(root_der)} bytes and does not fit the "
            f"{ROOT_CERT_SLOT_SIZE}-byte SecureROM slot"
        )

    (output / "cert-chain.der").write_bytes(intermediate_der + leaf_der)
    (output / "ticket-cert-chain.der").write_bytes(
        intermediate_der + ticket_leaf_der
    )
    gid_key = os.urandom(32)
    (output / "gid-key.bin").write_bytes(gid_key)

    patched = bytearray(source)
    patched[ROOT_CERT_OFFSET:ROOT_CERT_OFFSET + ROOT_CERT_SLOT_SIZE] = (
        root_der + bytes(ROOT_CERT_SLOT_SIZE - len(root_der))
    )
    derived_rom = output / "s5l8950x-secure-rom-lab.bin"
    derived_rom.write_bytes(patched)

    for name in (
        "root-key.pem", "intermediate-key.pem", "img3-leaf-key.pem",
        "ticket-leaf-key.pem", "gid-key.bin",
    ):
        os.chmod(output / name, stat.S_IRUSR | stat.S_IWUSR)
    for name in ("intermediate.csr", "img3-leaf.csr", "ticket-leaf.csr"):
        (output / name).unlink()

    manifest = {
        "format": 1,
        "source_rom": str(source_rom.resolve()),
        "source_rom_sha256": sha256(source),
        "derived_rom": derived_rom.name,
        "derived_rom_sha256": sha256(patched),
        "root_certificate_offset": ROOT_CERT_OFFSET,
        "root_certificate_slot_size": ROOT_CERT_SLOT_SIZE,
        "root_certificate_der_size": len(root_der),
        "root_certificate_sha256": sha256(root_der),
        "certificate_chain_der_size": len(intermediate_der) + len(leaf_der),
        "ticket_certificate_chain_der_size": (
            len(intermediate_der) + len(ticket_leaf_der)
        ),
        "gid_key_bits": 256,
        "gid_key_sha256": sha256(gid_key),
    }
    (output / "identity.json").write_text(
        json.dumps(manifest, indent=2) + "\n", encoding="utf-8"
    )
    return manifest


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("source_rom", type=Path)
    parser.add_argument("output_directory", type=Path)
    args = parser.parse_args()

    script_directory = Path(__file__).resolve().parent
    config = script_directory / "a6-lab-openssl.cnf"
    if not config.is_file():
        raise LabIdentityError(f"OpenSSL configuration not found: {config}")
    if not args.source_rom.is_file():
        raise LabIdentityError(f"SecureROM not found: {args.source_rom}")
    if shutil.which("openssl") is None:
        raise LabIdentityError("openssl is required")

    manifest = create_identity(
        args.source_rom.resolve(), args.output_directory.resolve(), config
    )
    print(f"Lab identity created in {args.output_directory.resolve()}")
    print(
        "Derived SecureROM SHA-256: " + manifest["derived_rom_sha256"]
    )
    print(
        "GID key: 256-bit private lab value (SHA-256 fingerprint "
        + manifest["gid_key_sha256"] + ")"
    )


if __name__ == "__main__":
    try:
        main()
    except (OSError, LabIdentityError) as error:
        print(f"error: {error}", file=sys.stderr)
        sys.exit(2)