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

"""Build a signed, optionally GID-wrapped IMG3 for the A6 lab identity."""

import argparse
import hashlib
import json
import os
from pathlib import Path
import struct
import subprocess
import sys
import tempfile


IMG3_HEADER = struct.Struct("<4s4I")
TAG_HEADER = struct.Struct("<4s2I")
IMG3_MAGIC = b"3gmI"
SHSH_SIZE = 128
DEFAULT_ECID = 0x200000
DEFAULT_SECURITY_DOMAIN = 3
DEFAULT_PRODUCTION_MODE = 1
DEFAULT_BOARD_ID = 0
DEFAULT_CHIP_EPOCH = 0x10


class Img3BuildError(RuntimeError):
    pass


def raw_tag(name):
    return name.encode("ascii")[::-1]


def display_tag(value):
    return value[::-1].decode("ascii", errors="replace")


def parse_template(data):
    if len(data) < IMG3_HEADER.size:
        raise Img3BuildError("template is shorter than the IMG3 header")
    magic, full_size, data_size, _, image_type = IMG3_HEADER.unpack_from(data)
    if magic != IMG3_MAGIC:
        raise Img3BuildError("template does not have IMG3 magic")
    if full_size > len(data) or data_size > full_size - IMG3_HEADER.size:
        raise Img3BuildError("template IMG3 sizes are out of bounds")

    tags = []
    offset = IMG3_HEADER.size
    end = IMG3_HEADER.size + data_size
    while offset < end:
        if end - offset < TAG_HEADER.size:
            raise Img3BuildError(f"truncated tag at 0x{offset:x}")
        tag, total_size, payload_size = TAG_HEADER.unpack_from(data, offset)
        if (
            total_size < TAG_HEADER.size
            or payload_size > total_size - TAG_HEADER.size
            or total_size > end - offset
        ):
            raise Img3BuildError(f"invalid tag at 0x{offset:x}")
        tags.append({
            "name": display_tag(tag),
            "payload": data[
                offset + TAG_HEADER.size:offset + TAG_HEADER.size + payload_size
            ],
            "padding": total_size - TAG_HEADER.size - payload_size,
        })
        offset += total_size
    if offset != end:
        raise Img3BuildError("template tag sizes do not cover data_size")
    return image_type, tags


def make_tag(name, payload, padding=0):
    total_size = TAG_HEADER.size + len(payload) + padding
    return (
        TAG_HEADER.pack(raw_tag(name), total_size, len(payload))
        + payload
        + bytes(padding)
    )


def openssl_filter(arguments, payload):
    try:
        result = subprocess.run(
            ["openssl", *arguments], input=payload, stdout=subprocess.PIPE,
            stderr=subprocess.PIPE, check=True,
        )
    except FileNotFoundError as error:
        raise Img3BuildError("openssl is required") from error
    except subprocess.CalledProcessError as error:
        detail = error.stderr.decode("utf-8", errors="replace").strip()
        raise Img3BuildError(f"openssl failed: {detail}") from error
    return result.stdout


def aes256_cbc(payload, key, iv, decrypt=False):
    if len(payload) % 16:
        raise Img3BuildError("AES-CBC input must be a multiple of 16 bytes")
    arguments = [
        "enc", "-aes-256-cbc", "-nopad", "-K", key.hex(), "-iv", iv.hex()
    ]
    if decrypt:
        arguments.append("-d")
    return openssl_filter(arguments, payload)


def sign_sha1(payload, private_key):
    return openssl_filter(
        ["dgst", "-sha1", "-sign", str(private_key)], payload
    )


def build_image(template, payload, identity, ecid, encrypted, ticketed,
                security_domain, production_mode, board_id, chip_epoch):
    image_type, template_tags = parse_template(template)
    chain = None
    leaf_key = None
    if not ticketed:
        chain = (identity / "cert-chain.der").read_bytes()
        leaf_key = identity / "img3-leaf-key.pem"
        if not leaf_key.is_file():
            raise Img3BuildError(f"leaf private key not found: {leaf_key}")

    if encrypted and len(payload) % 16:
        raise Img3BuildError(
            "encrypted DATA length must be a multiple of the AES block size"
        )

    content_iv = os.urandom(16) if encrypted else None
    content_key = os.urandom(32) if encrypted else None
    image_data = (
        aes256_cbc(payload, content_key, content_iv) if encrypted else payload
    )

    elements = []
    saw_type = False
    saw_data = False
    for tag in template_tags:
        name = tag["name"]
        if name in {"KBAG", "ECID", "SHSH", "CERT"}:
            continue
        if name == "TYPE":
            saw_type = True
        if name == "DATA":
            saw_data = True
            padding = (-(TAG_HEADER.size + len(image_data))) % 4
            elements.append(make_tag("DATA", image_data, padding))
        else:
            elements.append(make_tag(name, tag["payload"], tag["padding"]))
    if not saw_type or not saw_data:
        raise Img3BuildError("template must contain TYPE and DATA tags")

    if not ticketed:
        # The public IPSW template is unpersonalized and therefore omits the
        # device-bound scalar tags that SecureROM requires after SHSH
        # validation.  The iBEC path carries these values in its external
        # APTicket instead, matching the original iOS 10 boot flow.
        existing_tags = {tag["name"] for tag in template_tags}
        for name, value in (
            ("SDOM", security_domain),
            ("PROD", production_mode),
            ("CEPO", chip_epoch),
            ("BORD", board_id),
        ):
            if name not in existing_tags:
                elements.append(make_tag(name, struct.pack("<I", value)))

    if encrypted:
        gid_key_path = identity / "gid-key.bin"
        gid_key = gid_key_path.read_bytes()
        if len(gid_key) != 32:
            raise Img3BuildError("lab GID key must be exactly 32 bytes")
        clear_keybag = content_iv + content_key
        wrapped_keybag = aes256_cbc(clear_keybag, gid_key, bytes(16))
        unwrapped_keybag = aes256_cbc(
            wrapped_keybag, gid_key, bytes(16), decrypt=True
        )
        if unwrapped_keybag != clear_keybag:
            raise Img3BuildError("internal KBAG AES round-trip failed")
        elements.append(
            make_tag("KBAG", struct.pack("<II", 1, 256) + wrapped_keybag)
        )

    before_shsh = b"".join(elements)
    if ticketed:
        data_size = len(before_shsh)
        full_size = IMG3_HEADER.size + data_size
        header = IMG3_HEADER.pack(
            IMG3_MAGIC, full_size, data_size, data_size, image_type
        )
        result = header + before_shsh
        if encrypted:
            decrypted = aes256_cbc(
                image_data, content_key, content_iv, decrypt=True
            )
            if decrypted != payload:
                raise Img3BuildError("internal DATA AES round-trip failed")
        metadata = {
            "format": 1,
            "image_type": display_tag(struct.pack("<I", image_type)),
            "full_size": full_size,
            "data_size": data_size,
            "shsh_offset": data_size,
            "authentication": "external-apticket",
            "encrypted": encrypted,
            "kbag_key_modifier": 1 if encrypted else None,
            "kbag_key_bits": 256 if encrypted else None,
            "apticket_component_range": "0x0c..EOF",
            "apticket_component_sha1": hashlib.sha1(result[12:]).hexdigest(),
            "img3_sha256": hashlib.sha256(result).hexdigest(),
        }
        return result, metadata

    # Apple's A6 TSS response uses a 64-byte ECID element (8 data + 44 pad).
    elements.append(make_tag("ECID", struct.pack("<Q", ecid), 44))
    before_shsh = b"".join(elements)
    cert_element_size = TAG_HEADER.size + len(chain)
    data_size = (
        len(before_shsh) + TAG_HEADER.size + SHSH_SIZE + cert_element_size
    )
    full_size = IMG3_HEADER.size + data_size
    header = IMG3_HEADER.pack(
        IMG3_MAGIC, full_size, data_size, len(before_shsh), image_type
    )

    signed_bytes = header[12:] + before_shsh
    signature = sign_sha1(signed_bytes, leaf_key)
    if len(signature) != SHSH_SIZE:
        raise Img3BuildError(
            f"IMG3 leaf produced a {len(signature)}-byte signature; "
            "expected 128"
        )

    result = (
        header + before_shsh + make_tag("SHSH", signature)
        + make_tag("CERT", chain)
    )
    if len(result) != full_size:
        raise Img3BuildError("internal IMG3 size mismatch")
    if encrypted:
        decrypted = aes256_cbc(
            image_data, content_key, content_iv, decrypt=True
        )
        if decrypted != payload:
            raise Img3BuildError("internal DATA AES round-trip failed")

    metadata = {
        "format": 1,
        "image_type": display_tag(struct.pack("<I", image_type)),
        "full_size": full_size,
        "data_size": data_size,
        "shsh_offset": len(before_shsh),
        "ecid": f"0x{ecid:016x}",
        "security_domain": security_domain,
        "production_mode": production_mode,
        "board_id": board_id,
        "chip_epoch": chip_epoch,
        "authentication": "embedded-shsh",
        "encrypted": encrypted,
        "kbag_key_modifier": 1 if encrypted else None,
        "kbag_key_bits": 256 if encrypted else None,
        "signed_range": "0x0c..SHSH",
        "signed_sha1": hashlib.sha1(signed_bytes).hexdigest(),
        "img3_sha256": hashlib.sha256(result).hexdigest(),
    }
    return result, metadata


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--template", required=True, type=Path)
    parser.add_argument("--payload", required=True, type=Path)
    parser.add_argument("--identity", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    parser.add_argument(
        "--ecid", type=lambda value: int(value, 0), default=DEFAULT_ECID
    )
    parser.add_argument(
        "--security-domain", type=lambda value: int(value, 0),
        default=DEFAULT_SECURITY_DOMAIN,
    )
    parser.add_argument(
        "--production-mode", type=lambda value: int(value, 0), choices=(0, 1),
        default=DEFAULT_PRODUCTION_MODE,
    )
    parser.add_argument(
        "--board-id", type=lambda value: int(value, 0),
        default=DEFAULT_BOARD_ID,
    )
    parser.add_argument(
        "--chip-epoch", type=lambda value: int(value, 0),
        default=DEFAULT_CHIP_EPOCH,
    )
    parser.add_argument("--encrypt", action="store_true")
    parser.add_argument(
        "--ticketed", action="store_true",
        help="omit embedded ECID/SHSH/CERT for an external APTicket (iBEC)",
    )
    args = parser.parse_args()

    for path, description in (
        (args.template, "IMG3 template"),
        (args.payload, "payload"),
    ):
        if not path.is_file():
            raise Img3BuildError(f"{description} not found: {path}")
    if not args.ticketed and not (args.identity / "cert-chain.der").is_file():
        raise Img3BuildError(
            "lab certificate chain not found: "
            f"{args.identity / 'cert-chain.der'}"
        )
    if args.output.exists():
        raise Img3BuildError(f"refusing to overwrite output: {args.output}")

    image, metadata = build_image(
        args.template.read_bytes(), args.payload.read_bytes(),
        args.identity.resolve(), args.ecid, args.encrypt, args.ticketed,
        args.security_domain, args.production_mode, args.board_id,
        args.chip_epoch,
    )
    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(image)
        temporary_path = Path(temporary.name)
    temporary_path.replace(args.output)
    manifest_path = args.output.with_name(args.output.name + ".json")
    manifest_path.write_text(
        json.dumps(metadata, indent=2) + "\n", encoding="utf-8"
    )

    print(
        f"Created {metadata['image_type']} IMG3: {args.output} "
        f"({metadata['full_size']} bytes)"
    )
    print(f"Encrypted DATA/KBAG: {'yes' if args.encrypt else 'no'}")
    if args.ticketed:
        print(
            "APTicket component SHA-1: "
            f"{metadata['apticket_component_sha1']}"
        )
    else:
        print(f"Signed SHA-1: {metadata['signed_sha1']}")


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