bongo/contrib/testing/local-test-pki.py

#!/usr/bin/env python3
# This program is free software, licensed under the terms of the GNU GPL.
# See the Bongo COPYING file for full details.
# Copyright (c) 2026 Bongo Project contributors

"""Manage the two-root local PKI used by Bongo protocol tests."""

from __future__ import annotations

import argparse
import ipaddress
import os
from pathlib import Path
import re
import shutil
import subprocess
import sys
import tempfile


DEFAULT_ROOT = Path("/var/lib/bongo-test-pki")
TRUST_ANCHOR = Path(
    "/usr/local/share/ca-certificates/bongo-local-test-root.crt")
SYSTEM_BUNDLE = Path("/etc/ssl/certs/ca-certificates.crt")
CERTTOOL = Path("/usr/bin/certtool")
UPDATE_CA_CERTIFICATES = Path("/usr/bin/update-ca-certificates")
OPENSSL = Path("/usr/bin/openssl")
SAFE_NAME = re.compile(r"^[a-z0-9][a-z0-9._-]{0,62}$")


def paths(root: Path) -> dict[str, Path]:
    return {
        "trusted_key": root / "trusted-ca.key",
        "trusted_cert": root / "trusted-ca.crt",
        "untrusted_key": root / "untrusted-ca.key",
        "untrusted_cert": root / "untrusted-ca.crt",
        "issued": root / "issued",
    }


def run(arguments: list[str | Path], *, quiet: bool = False) -> None:
    command = [str(argument) for argument in arguments]
    completed = subprocess.run(
        command,
        check=False,
        stdout=subprocess.DEVNULL if quiet else None,
        stderr=subprocess.PIPE,
        text=True,
    )
    if completed.returncode:
        detail = completed.stderr.strip()
        raise RuntimeError(
            f"{Path(command[0]).name} failed"
            + (f": {detail}" if detail else ""))


def require_tools() -> None:
    missing = [
        path for path in (CERTTOOL, UPDATE_CA_CERTIFICATES, OPENSSL)
        if not path.is_file()
    ]
    if missing:
        raise RuntimeError(
            "missing test-PKI tools: "
            + ", ".join(str(path) for path in missing))


def require_root(action: str) -> None:
    if os.geteuid() != 0:
        raise RuntimeError(f"{action} must be run as root")


def ensure_private_directory(directory: Path) -> None:
    directory.mkdir(parents=True, exist_ok=True)
    if directory.is_symlink():
        raise RuntimeError(f"refusing symlinked PKI directory: {directory}")
    directory.chmod(0o700)


def write_template(content: str) -> Path:
    descriptor, raw_path = tempfile.mkstemp(prefix="bongo-test-pki-",
                                            suffix=".tmpl")
    path = Path(raw_path)
    try:
        os.write(descriptor, content.encode("ascii"))
    finally:
        os.close(descriptor)
    path.chmod(0o600)
    return path


def generate_ca(root: Path, kind: str, common_name: str) -> None:
    values = paths(root)
    key = values[f"{kind}_key"]
    certificate = values[f"{kind}_cert"]
    if key.exists() or certificate.exists():
        if key.is_file() and certificate.is_file():
            return
        raise RuntimeError(f"incomplete {kind} test CA below {root}")

    template = write_template(
        f"""cn = {common_name}
organization = Bongo local protocol tests
expiration_days = 3650
ca
cert_signing_key
crl_signing_key
""")
    temporary_key = key.with_suffix(".key.new")
    temporary_certificate = certificate.with_suffix(".crt.new")
    try:
        run([
            CERTTOOL, "--generate-privkey", "--sec-param", "high",
            "--outfile", temporary_key,
        ], quiet=True)
        temporary_key.chmod(0o600)
        run([
            CERTTOOL, "--generate-self-signed",
            "--load-privkey", temporary_key,
            "--template", template,
            "--outfile", temporary_certificate,
        ], quiet=True)
        temporary_certificate.chmod(0o644)
        temporary_key.replace(key)
        temporary_certificate.replace(certificate)
    finally:
        template.unlink(missing_ok=True)
        temporary_key.unlink(missing_ok=True)
        temporary_certificate.unlink(missing_ok=True)


def prepare(root: Path) -> None:
    require_tools()
    if root == DEFAULT_ROOT:
        require_root("preparing the system test PKI")
    ensure_private_directory(root)
    ensure_private_directory(paths(root)["issued"])
    generate_ca(root, "trusted", "Bongo Local Test Root CA")
    generate_ca(root, "untrusted", "Bongo Untrusted Test Root CA")
    print(f"Trusted test CA:   {paths(root)['trusted_cert']}")
    print(f"Untrusted test CA: {paths(root)['untrusted_cert']}")


def normalise_dns_name(value: str) -> str:
    try:
        result = value.rstrip(".").encode("idna").decode("ascii").lower()
    except UnicodeError as error:
        raise RuntimeError(f"invalid DNS name {value!r}: {error}") from error
    labels = result.split(".")
    valid_label = r"[a-z0-9](?:[a-z0-9-]*[a-z0-9])?"
    if (not result or len(result) > 253 or
            any(not label or len(label) > 63 for label in labels) or
            any(not re.fullmatch(valid_label, label) for label in labels)):
        raise RuntimeError(f"invalid DNS name: {value!r}")
    return result


def issue(root: Path, kind: str, name: str, dns_names: list[str],
          ip_addresses: list[str]) -> None:
    require_tools()
    if root == DEFAULT_ROOT:
        require_root("issuing a system test certificate")
    if not SAFE_NAME.fullmatch(name):
        raise RuntimeError(
            "certificate name must contain only lowercase letters, digits, "
            "dot, underscore, and hyphen")
    prepare(root)
    dns_names = list(dict.fromkeys(normalise_dns_name(item)
                                   for item in dns_names))
    if not dns_names:
        raise RuntimeError("at least one --dns-name is required")
    parsed_addresses = [str(ipaddress.ip_address(item))
                        for item in ip_addresses]

    values = paths(root)
    output = values["issued"] / name
    if output.exists():
        raise RuntimeError(
            f"certificate output already exists; remove it first: {output}")
    ensure_private_directory(output)
    ca_key = values[f"{kind}_key"]
    ca_certificate = values[f"{kind}_cert"]
    key = output / "key.pem"
    certificate = output / "cert.pem"
    fullchain = output / "fullchain.pem"
    template_lines = [
        f"cn = {dns_names[0]}",
        "organization = Bongo local protocol tests",
        "expiration_days = 30",
        "signing_key",
        "encryption_key",
        "tls_www_server",
    ]
    template_lines.extend(f"dns_name = {item}" for item in dns_names)
    template_lines.extend(f"ip_address = {item}" for item in parsed_addresses)
    template = write_template("\n".join(template_lines) + "\n")
    try:
        run([
            CERTTOOL, "--generate-privkey", "--sec-param", "high",
            "--outfile", key,
        ], quiet=True)
        key.chmod(0o600)
        run([
            CERTTOOL, "--generate-certificate",
            "--load-privkey", key,
            "--load-ca-certificate", ca_certificate,
            "--load-ca-privkey", ca_key,
            "--template", template,
            "--outfile", certificate,
        ], quiet=True)
        certificate.chmod(0o644)
        fullchain.write_bytes(
            certificate.read_bytes() + ca_certificate.read_bytes())
        fullchain.chmod(0o644)
    except Exception:
        shutil.rmtree(output, ignore_errors=True)
        raise
    finally:
        template.unlink(missing_ok=True)
    print(f"Issued {kind} certificate: {certificate}")
    print(f"Private key: {key}")
    print(f"Full chain: {fullchain}")


def install(root: Path) -> None:
    require_root("installing the trusted test CA")
    prepare(root)
    source = paths(root)["trusted_cert"]
    TRUST_ANCHOR.parent.mkdir(parents=True, exist_ok=True)
    if (TRUST_ANCHOR.exists()
            and TRUST_ANCHOR.read_bytes() != source.read_bytes()):
        raise RuntimeError(
            f"refusing to replace a different trust anchor: {TRUST_ANCHOR}")
    temporary = TRUST_ANCHOR.with_suffix(".crt.new")
    shutil.copyfile(source, temporary)
    temporary.chmod(0o644)
    temporary.replace(TRUST_ANCHOR)
    run([UPDATE_CA_CERTIFICATES])
    run([OPENSSL, "verify", "-CAfile", SYSTEM_BUNDLE, source])
    print(f"Installed trusted test CA: {TRUST_ANCHOR}")
    print("The untrusted test CA was deliberately not installed.")


def remove(root: Path) -> None:
    require_root("removing the trusted test CA")
    source = paths(root)["trusted_cert"]
    if TRUST_ANCHOR.exists():
        if (source.exists()
                and TRUST_ANCHOR.read_bytes() != source.read_bytes()):
            raise RuntimeError(
                f"refusing to remove a different trust anchor: {TRUST_ANCHOR}")
        TRUST_ANCHOR.unlink()
    run([UPDATE_CA_CERTIFICATES])
    print(f"Removed trusted test CA: {TRUST_ANCHOR}")


def status(root: Path) -> None:
    values = paths(root)
    print(f"PKI directory: {root}")
    print(
        "Trusted CA generated: "
        f"{'yes' if values['trusted_cert'].is_file() else 'no'}")
    print(
        "Untrusted CA generated: "
        f"{'yes' if values['untrusted_cert'].is_file() else 'no'}")
    installed = (
        TRUST_ANCHOR.is_file()
        and values["trusted_cert"].is_file()
        and TRUST_ANCHOR.read_bytes() == values["trusted_cert"].read_bytes()
    )
    print(f"Trusted CA installed system-wide: {'yes' if installed else 'no'}")
    print(f"System trust anchor: {TRUST_ANCHOR}")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--root", type=Path, default=DEFAULT_ROOT)
    subparsers = parser.add_subparsers(dest="action", required=True)
    subparsers.add_parser("prepare")
    subparsers.add_parser("install")
    subparsers.add_parser("remove")
    subparsers.add_parser("status")
    issue_parser = subparsers.add_parser("issue")
    issue_parser.add_argument("--ca", choices=("trusted", "untrusted"),
                              default="trusted")
    issue_parser.add_argument("--name", required=True)
    issue_parser.add_argument("--dns-name", action="append", default=[])
    issue_parser.add_argument("--ip-address", action="append", default=[])
    return parser


def main() -> int:
    arguments = build_parser().parse_args()
    try:
        if arguments.action == "prepare":
            prepare(arguments.root)
        elif arguments.action == "install":
            install(arguments.root)
        elif arguments.action == "remove":
            remove(arguments.root)
        elif arguments.action == "status":
            status(arguments.root)
        elif arguments.action == "issue":
            issue(arguments.root, arguments.ca, arguments.name,
                  arguments.dns_name, arguments.ip_address)
    except (OSError, RuntimeError) as error:
        print(f"local-test-pki: {error}", file=sys.stderr)
        return 1
    return 0


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