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

"""Replace an iBoot payload's embedded IMG3 trust anchor with the lab root."""

import argparse
import hashlib
import json
from pathlib import Path
import sys
import tempfile


ROM_SIZE = 0x10000
ROOT_CERT_OFFSET = 0xBFD0
ROOT_CERT_SLOT_SIZE = 1215


class RootPatchError(RuntimeError):
    pass


def der_object_size(data):
    if len(data) < 2 or data[0] != 0x30:
        raise RootPatchError("lab root is not a DER SEQUENCE")
    first = data[1]
    if first < 0x80:
        return 2 + first
    length_bytes = first & 0x7F
    if not 1 <= length_bytes <= 4 or 2 + length_bytes > len(data):
        raise RootPatchError("lab root has an invalid DER length")
    payload_size = int.from_bytes(data[2:2 + length_bytes], "big")
    return 2 + length_bytes + payload_size


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source-rom", required=True, type=Path)
    parser.add_argument("--identity", required=True, type=Path)
    parser.add_argument("--input", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    args = parser.parse_args()

    if args.output.exists():
        raise RootPatchError(f"refusing to overwrite output: {args.output}")
    source_rom = args.source_rom.read_bytes()
    if len(source_rom) != ROM_SIZE:
        raise RootPatchError(f"source SecureROM must be {ROM_SIZE} bytes")
    source_anchor = source_rom[
        ROOT_CERT_OFFSET:ROOT_CERT_OFFSET + ROOT_CERT_SLOT_SIZE
    ]
    lab_root = (args.identity / "root.der").read_bytes()
    if der_object_size(lab_root) != len(lab_root):
        raise RootPatchError("lab root contains trailing data")
    if len(lab_root) > ROOT_CERT_SLOT_SIZE:
        raise RootPatchError(
            "lab root does not fit the iBoot trust-anchor slot"
        )
    lab_anchor = lab_root + bytes(ROOT_CERT_SLOT_SIZE - len(lab_root))

    payload = args.input.read_bytes()
    offsets = []
    start = 0
    while True:
        offset = payload.find(source_anchor, start)
        if offset < 0:
            break
        offsets.append(offset)
        start = offset + 1
    if len(offsets) != 1:
        raise RootPatchError(
            f"expected exactly one embedded source root, found {len(offsets)}"
        )

    offset = offsets[0]
    patched = (
        payload[:offset]
        + lab_anchor
        + payload[offset + ROOT_CERT_SLOT_SIZE:]
    )
    args.output.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile(
        prefix=args.output.name + ".", dir=args.output.parent, delete=False
    ) as temporary:
        temporary.write(patched)
        temporary_path = Path(temporary.name)
    temporary_path.replace(args.output)
    metadata = {
        "format": 1,
        "input": str(args.input.resolve()),
        "output": args.output.name,
        "trust_anchor_offset": offset,
        "trust_anchor_slot_size": ROOT_CERT_SLOT_SIZE,
        "input_sha256": hashlib.sha256(payload).hexdigest(),
        "output_sha256": hashlib.sha256(patched).hexdigest(),
        "lab_root_sha256": hashlib.sha256(lab_root).hexdigest(),
    }
    args.output.with_name(args.output.name + ".json").write_text(
        json.dumps(metadata, indent=2) + "\n", encoding="utf-8"
    )
    print(
        f"Patched {args.output} ({len(patched)} bytes), "
        f"trust anchor at 0x{offset:x}"
    )


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