#!/usr/bin/env python3
"""Create OpenPGP public keys and detached signatures backed by Google KMS.

This intentionally implements only the narrow RSA/SHA-256 subset Burrow needs
for Linux package repository metadata. The private key lives in Google Cloud KMS;
this script builds OpenPGP packets locally and asks KMS to produce the RSA
signature value.
"""

from __future__ import annotations

import argparse
import base64
import binascii
import hashlib
import os
import struct
import subprocess
import sys
import tempfile
import time
from pathlib import Path


RSA_PUBLIC_KEY_ALGO = 1
SHA256_HASH_ALGO = 8


def parse_der_length(data: bytes, offset: int) -> tuple[int, int]:
    first = data[offset]
    offset += 1
    if first < 0x80:
        return first, offset
    count = first & 0x7F
    if count == 0 or count > 4:
        raise ValueError("unsupported DER length")
    value = int.from_bytes(data[offset : offset + count], "big")
    return value, offset + count


def parse_der_tlv(data: bytes, offset: int, expected_tag: int | None = None) -> tuple[int, bytes, int]:
    tag = data[offset]
    offset += 1
    length, offset = parse_der_length(data, offset)
    value = data[offset : offset + length]
    if len(value) != length:
        raise ValueError("truncated DER value")
    if expected_tag is not None and tag != expected_tag:
        raise ValueError(f"expected DER tag {expected_tag:#x}, got {tag:#x}")
    return tag, value, offset + length


def pem_to_der(pem: bytes) -> bytes:
    lines = [
        line.strip()
        for line in pem.splitlines()
        if line and not line.startswith(b"-----")
    ]
    return base64.b64decode(b"".join(lines))


def parse_rsa_public_key_from_pem(path: Path) -> tuple[int, int]:
    der = pem_to_der(path.read_bytes())
    _, spki, end = parse_der_tlv(der, 0, 0x30)
    if end != len(der):
        raise ValueError("trailing data after SubjectPublicKeyInfo")
    _, _alg, offset = parse_der_tlv(spki, 0, 0x30)
    _, bit_string, offset = parse_der_tlv(spki, offset, 0x03)
    if offset != len(spki):
        raise ValueError("trailing data after SPKI bit string")
    if not bit_string or bit_string[0] != 0:
        raise ValueError("unsupported SPKI bit string")
    rsa_der = bit_string[1:]
    _, rsa_seq, end = parse_der_tlv(rsa_der, 0, 0x30)
    if end != len(rsa_der):
        raise ValueError("trailing data after RSA public key")
    _, n_bytes, offset = parse_der_tlv(rsa_seq, 0, 0x02)
    _, e_bytes, offset = parse_der_tlv(rsa_seq, offset, 0x02)
    if offset != len(rsa_seq):
        raise ValueError("trailing data after RSA integers")
    return int.from_bytes(n_bytes, "big"), int.from_bytes(e_bytes, "big")


def packet_header(tag: int, length: int) -> bytes:
    if length < 192:
        return bytes([0xC0 | tag, length])
    if length < 8384:
        length -= 192
        return bytes([0xC0 | tag, (length >> 8) + 192, length & 0xFF])
    return bytes([0xC0 | tag, 0xFF]) + struct.pack(">I", length)


def old_packet_header(tag: int, length: int) -> bytes:
    if length > 0xFFFF:
        raise ValueError("old-format two-octet packet length exceeded")
    return bytes([(tag << 2) | 1]) + struct.pack(">H", length)


def mpi(value: int) -> bytes:
    if value < 0:
        raise ValueError("MPI value must be non-negative")
    raw = value.to_bytes(max(1, (value.bit_length() + 7) // 8), "big")
    return struct.pack(">H", value.bit_length()) + raw.lstrip(b"\x00")


def subpacket(kind: int, body: bytes) -> bytes:
    payload = bytes([kind]) + body
    length = len(payload)
    if length < 192:
        return bytes([length]) + payload
    if length < 8384:
        length -= 192
        return bytes([(length >> 8) + 192, length & 0xFF]) + payload
    return b"\xff" + struct.pack(">I", length) + payload


def public_key_body(n: int, e: int, created_at: int) -> bytes:
    return (
        b"\x04"
        + struct.pack(">I", created_at)
        + bytes([RSA_PUBLIC_KEY_ALGO])
        + mpi(n)
        + mpi(e)
    )


def fingerprint(public_body: bytes) -> bytes:
    return hashlib.sha1(old_packet_header(6, len(public_body)) + public_body).digest()


def armor(kind: str, body: bytes) -> str:
    b64 = base64.b64encode(body).decode("ascii")
    crc = crc24(body)
    checksum = base64.b64encode(crc.to_bytes(3, "big")).decode("ascii")
    lines = [f"-----BEGIN PGP {kind}-----", ""]
    lines.extend(b64[i : i + 64] for i in range(0, len(b64), 64))
    lines.append(f"={checksum}")
    lines.append(f"-----END PGP {kind}-----")
    return "\n".join(lines) + "\n"


def crc24(data: bytes) -> int:
    crc = 0xB704CE
    for octet in data:
        crc ^= octet << 16
        for _ in range(8):
            crc <<= 1
            if crc & 0x1000000:
                crc ^= 0x1864CFB
    return crc & 0xFFFFFF


def kms_sign(payload: bytes, args: argparse.Namespace) -> bytes:
    with tempfile.TemporaryDirectory(prefix="burrow-kms-openpgp.") as tmp:
        input_path = Path(tmp) / "payload"
        signature_path = Path(tmp) / "signature.b64"
        input_path.write_bytes(payload)
        command = [
            "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(input_path),
            "--signature-file",
            str(signature_path),
        ]
        if args.gcloud_project:
            command.extend(["--project", args.gcloud_project])
        subprocess.run(command, check=True)
        return base64.b64decode(signature_path.read_text(encoding="utf-8").strip())


def signature_packet(
    signed_payload: bytes,
    public_body: bytes,
    sig_type: int,
    created_at: int,
    args: argparse.Namespace,
) -> bytes:
    fpr = fingerprint(public_body)
    key_id = fpr[-8:]
    hashed = b"".join(
        [
            subpacket(2, struct.pack(">I", created_at)),
            subpacket(33, b"\x04" + fpr),
        ]
    )
    unhashed = subpacket(16, key_id)
    sig_data = (
        b"\x04"
        + bytes([sig_type, RSA_PUBLIC_KEY_ALGO, SHA256_HASH_ALGO])
        + struct.pack(">H", len(hashed))
        + hashed
    )
    trailer = b"\x04\xff" + struct.pack(">I", len(sig_data))
    to_hash = signed_payload + sig_data + trailer
    digest = hashlib.sha256(to_hash).digest()
    signature = kms_sign(to_hash, args)
    body = (
        sig_data
        + struct.pack(">H", len(unhashed))
        + unhashed
        + digest[:2]
        + mpi(int.from_bytes(signature, "big"))
    )
    return packet_header(2, len(body)) + body


def public_key_command(args: argparse.Namespace) -> int:
    created_at = args.created_at or int(time.time())
    n, e = parse_rsa_public_key_from_pem(Path(args.public_key_pem))
    public_body = public_key_body(n, e, created_at)
    public_packet = packet_header(6, len(public_body)) + public_body
    uid = args.user_id.encode("utf-8")
    uid_packet = packet_header(13, len(uid)) + uid
    uid_signed_payload = (
        old_packet_header(6, len(public_body))
        + public_body
        + b"\xb4"
        + struct.pack(">I", len(uid))
        + uid
    )
    cert_packet = signature_packet(
        uid_signed_payload,
        public_body,
        sig_type=0x13,
        created_at=created_at,
        args=args,
    )
    key = public_packet + uid_packet + cert_packet
    output = Path(args.output)
    if args.armor:
        output.write_text(armor("PUBLIC KEY BLOCK", key), encoding="utf-8")
    else:
        output.write_bytes(key)
    return 0


def detach_sign_command(args: argparse.Namespace) -> int:
    created_at = args.created_at or int(time.time())
    n, e = parse_rsa_public_key_from_pem(Path(args.public_key_pem))
    public_body = public_key_body(n, e, args.key_created_at)
    payload = Path(args.input).read_bytes()
    sig_type = 0x01 if args.text else 0x00
    packet = signature_packet(payload, public_body, sig_type, created_at, args)
    output = Path(args.output)
    if args.armor:
        output.write_text(armor("SIGNATURE", packet), encoding="utf-8")
    else:
        output.write_bytes(packet)
    return 0


def add_kms_args(parser: argparse.ArgumentParser) -> None:
    parser.add_argument("--kms-location", default=os.environ.get("BURROW_KMS_LOCATION", "global"))
    parser.add_argument("--kms-keyring", default=os.environ.get("BURROW_KMS_KEYRING", "burrow-identity"))
    parser.add_argument("--kms-key", required=not os.environ.get("BURROW_KMS_KEY"), default=os.environ.get("BURROW_KMS_KEY"))
    parser.add_argument("--kms-version", default=os.environ.get("BURROW_KMS_VERSION", "1"))
    parser.add_argument("--gcloud-project", default=os.environ.get("GOOGLE_CLOUD_PROJECT"))


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    subcommands = parser.add_subparsers(dest="command", required=True)

    public_key = subcommands.add_parser("public-key")
    public_key.add_argument("--public-key-pem", required=True)
    public_key.add_argument("--user-id", default="Burrow Package Repository <packages@burrow.net>")
    public_key.add_argument("--output", required=True)
    public_key.add_argument("--armor", action="store_true")
    public_key.add_argument("--created-at", type=int)
    add_kms_args(public_key)
    public_key.set_defaults(func=public_key_command)

    detach = subcommands.add_parser("detach-sign")
    detach.add_argument("--public-key-pem", required=True)
    detach.add_argument("--key-created-at", type=int, default=1704067200)
    detach.add_argument("--input", required=True)
    detach.add_argument("--output", required=True)
    detach.add_argument("--armor", action="store_true")
    detach.add_argument("--text", action="store_true")
    detach.add_argument("--created-at", type=int)
    add_kms_args(detach)
    detach.set_defaults(func=detach_sign_command)

    return parser


def main() -> int:
    parser = build_parser()
    args = parser.parse_args()
    try:
        return args.func(args)
    except subprocess.CalledProcessError as exc:
        print(f"gcloud signing failed: {exc}", file=sys.stderr)
        return exc.returncode or 1
    except Exception as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())