#!/usr/bin/env python3
"""Build an Apple-compatible PKCS#10 CSR backed by Google Cloud KMS.

The script intentionally implements only the RSA/SHA-256 subset Burrow needs for
Apple Developer ID and iOS Distribution certificate requests. It fetches the KMS
public key, builds CertificationRequestInfo locally, asks KMS to sign that DER
payload, and writes a standard PEM CSR.
"""

from __future__ import annotations

import argparse
import base64
import os
import subprocess
import sys
import tempfile
from pathlib import Path


DEFAULT_PROJECT = "project-88c23ce9-918a-470a-b33"
DEFAULT_LOCATION = "global"
DEFAULT_KEYRING = "burrow-identity"
DEFAULT_KEY = "apple-developer-id-application"
DEFAULT_VERSION = "1"
DEFAULT_COMMON_NAME = "Burrow Developer ID Application"
DEFAULT_EMAIL = "release@burrow.net"


def der_length(length: int) -> bytes:
    if length < 0:
        raise ValueError("negative DER length")
    if length < 0x80:
        return bytes([length])
    encoded = length.to_bytes((length.bit_length() + 7) // 8, "big")
    return bytes([0x80 | len(encoded)]) + encoded


def der(tag: int, body: bytes) -> bytes:
    return bytes([tag]) + der_length(len(body)) + body


def der_sequence(*items: bytes) -> bytes:
    return der(0x30, b"".join(items))


def der_set(*items: bytes) -> bytes:
    return der(0x31, b"".join(items))


def der_integer(value: int) -> bytes:
    if value < 0:
        raise ValueError("negative INTEGER is not supported")
    raw = value.to_bytes(max(1, (value.bit_length() + 7) // 8), "big")
    if raw[0] & 0x80:
        raw = b"\x00" + raw
    return der(0x02, raw)


def der_null() -> bytes:
    return der(0x05, b"")


def der_bit_string(value: bytes) -> bytes:
    return der(0x03, b"\x00" + value)


def der_utf8(value: str) -> bytes:
    return der(0x0C, value.encode("utf-8"))


def der_printable(value: str) -> bytes:
    return der(0x13, value.encode("ascii"))


def der_ia5(value: str) -> bytes:
    return der(0x16, value.encode("ascii"))


def der_oid(dotted: str) -> bytes:
    parts = [int(part) for part in dotted.split(".")]
    if len(parts) < 2 or parts[0] > 2 or parts[1] > 39:
        raise ValueError(f"unsupported OID: {dotted}")
    body = bytes([40 * parts[0] + parts[1]])
    for value in parts[2:]:
        if value < 0:
            raise ValueError(f"negative OID component: {dotted}")
        encoded = [value & 0x7F]
        value >>= 7
        while value:
            encoded.append(0x80 | (value & 0x7F))
            value >>= 7
        body += bytes(reversed(encoded))
    return der(0x06, body)


def name_attribute(oid: str, value_der: bytes) -> bytes:
    return der_set(der_sequence(der_oid(oid), value_der))


def subject_name(common_name: str, email_address: str, country: str | None) -> bytes:
    attrs: list[bytes] = []
    if country:
        if len(country) != 2 or not country.isalpha():
            raise ValueError("--country must be a two-letter ISO country code")
        attrs.append(name_attribute("2.5.4.6", der_printable(country.upper())))
    if email_address:
        attrs.append(name_attribute("1.2.840.113549.1.9.1", der_ia5(email_address)))
    attrs.append(name_attribute("2.5.4.3", der_utf8(common_name)))
    return der_sequence(*attrs)


def pem_to_der(pem: bytes) -> bytes:
    lines = [
        line.strip()
        for line in pem.splitlines()
        if line and not line.startswith(b"-----")
    ]
    if not lines:
        raise ValueError("PEM input did not contain base64 data")
    return base64.b64decode(b"".join(lines), validate=True)


def pem_wrap(label: str, body: bytes) -> str:
    encoded = base64.b64encode(body).decode("ascii")
    lines = [f"-----BEGIN {label}-----"]
    lines.extend(encoded[i : i + 64] for i in range(0, len(encoded), 64))
    lines.append(f"-----END {label}-----")
    return "\n".join(lines) + "\n"


def run(command: list[str]) -> None:
    subprocess.run(command, check=True)


def fetch_public_key(args: argparse.Namespace, output: Path) -> None:
    command = [
        args.gcloud,
        "kms",
        "keys",
        "versions",
        "get-public-key",
        args.kms_version,
        "--location",
        args.kms_location,
        "--keyring",
        args.kms_keyring,
        "--key",
        args.kms_key,
        "--output-file",
        str(output),
    ]
    if args.gcloud_project:
        command.extend(["--project", args.gcloud_project])
    run(command)


def kms_sign(args: argparse.Namespace, payload: bytes) -> bytes:
    with tempfile.TemporaryDirectory(prefix="burrow-kms-csr.") as tmp:
        payload_path = Path(tmp) / "certification-request-info.der"
        signature_path = Path(tmp) / "signature.b64"
        payload_path.write_bytes(payload)
        command = [
            args.gcloud,
            "kms",
            "asymmetric-sign",
            "--location",
            args.kms_location,
            "--keyring",
            args.kms_keyring,
            "--key",
            args.kms_key,
            "--version",
            args.kms_version,
            "--digest-algorithm",
            "sha256",
            "--input-file",
            str(payload_path),
            "--signature-file",
            str(signature_path),
        ]
        if args.gcloud_project:
            command.extend(["--project", args.gcloud_project])
        run(command)
        signature = signature_path.read_bytes()
        try:
            return base64.b64decode(signature.strip(), validate=True)
        except ValueError:
            return signature


def build_csr(args: argparse.Namespace, spki_der: bytes) -> bytes:
    cri = der_sequence(
        der_integer(0),
        subject_name(args.common_name, args.email_address, args.country),
        spki_der,
        b"\xA0\x00",
    )
    signature = kms_sign(args, cri)
    signature_algorithm = der_sequence(
        der_oid("1.2.840.113549.1.1.11"),
        der_null(),
    )
    return der_sequence(cri, signature_algorithm, der_bit_string(signature))


def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Generate an Apple-compatible CSR backed by a Google Cloud KMS RSA signing key.",
    )
    parser.add_argument("--gcloud-project", default=os.environ.get("GOOGLE_CLOUD_PROJECT") or os.environ.get("BURROW_GOOGLE_PROJECT_ID") or DEFAULT_PROJECT)
    parser.add_argument("--kms-location", default=os.environ.get("BURROW_KMS_LOCATION", DEFAULT_LOCATION))
    parser.add_argument("--kms-keyring", default=os.environ.get("BURROW_KMS_KEYRING", DEFAULT_KEYRING))
    parser.add_argument("--kms-key", default=os.environ.get("BURROW_APPLE_KMS_KEY") or os.environ.get("BURROW_APPLE_DEVELOPER_ID_KMS_KEY", DEFAULT_KEY))
    parser.add_argument("--kms-version", default=os.environ.get("BURROW_KMS_VERSION", DEFAULT_VERSION))
    parser.add_argument("--common-name", default=os.environ.get("BURROW_APPLE_CSR_COMMON_NAME") or os.environ.get("BURROW_APPLE_DEVELOPER_ID_CSR_COMMON_NAME", DEFAULT_COMMON_NAME))
    parser.add_argument("--email-address", default=os.environ.get("BURROW_APPLE_CSR_EMAIL") or os.environ.get("BURROW_APPLE_DEVELOPER_ID_CSR_EMAIL", DEFAULT_EMAIL))
    parser.add_argument("--country", default=os.environ.get("BURROW_APPLE_CSR_COUNTRY") or os.environ.get("BURROW_APPLE_DEVELOPER_ID_CSR_COUNTRY", "US"))
    parser.add_argument("--public-key-pem", help="Existing KMS public-key PEM path. If omitted, gcloud fetches it.")
    parser.add_argument("--public-key-output", help="Write the fetched/used public key PEM to this path.")
    parser.add_argument("--output", default="developer-id-application.certSigningRequest")
    parser.add_argument("--gcloud", default=os.environ.get("GCLOUD", "gcloud"))
    return parser.parse_args(argv)


def main(argv: list[str]) -> int:
    args = parse_args(argv)
    output = Path(args.output)

    with tempfile.TemporaryDirectory(prefix="burrow-kms-csr.") as tmp:
        public_key_path = Path(args.public_key_pem) if args.public_key_pem else Path(tmp) / "kms-public.pem"
        if not args.public_key_pem:
            fetch_public_key(args, public_key_path)
        public_key_pem = public_key_path.read_bytes()
        spki_der = pem_to_der(public_key_pem)
        csr = build_csr(args, spki_der)

    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(pem_wrap("CERTIFICATE REQUEST", csr), encoding="ascii")
    if args.public_key_output:
        public_output = Path(args.public_key_output)
        public_output.parent.mkdir(parents=True, exist_ok=True)
        public_output.write_bytes(public_key_pem)
    print(f"wrote {output}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))