teleo-infrastructure/scripts/bootstrap_clone_kb_gate.py

#!/usr/bin/env python3
"""Install guarded review/apply prerequisites and ephemeral credentials in a disposable clone."""

from __future__ import annotations

import argparse
import json
import os
import secrets
import subprocess
import sys
from pathlib import Path
from typing import Any

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))

import apply_proposal as ap  # noqa: E402
import run_leo_clone_bound_handler_checkpoint as bound  # noqa: E402


def _psql(
    container: str,
    database: str,
    sql: str,
    *,
    secret_values: tuple[str, ...] = (),
) -> None:
    command = [
        bound.SYSTEM_DOCKER,
        "exec",
        "-i",
        container,
        "psql",
        "-U",
        "postgres",
        "-d",
        database,
        "-v",
        "ON_ERROR_STOP=1",
        "-At",
        "-q",
    ]
    result = subprocess.run(
        command,
        input=sql,
        text=True,
        capture_output=True,
        env={"PATH": bound.SYSTEM_EXEC_PATH},
        check=False,
    )
    if result.returncode != 0:
        raise bound.CheckpointError(
            f"clone gate SQL failed ({result.returncode}): "
            + bound.redact_text(result.stderr or result.stdout, secret_values=secret_values)
        )


def _verify_login(container: str, database: str, role: str, password: str) -> None:
    command = [
        bound.SYSTEM_DOCKER,
        "exec",
        "-e",
        "PGPASSWORD",
        "-i",
        container,
        "psql",
        "-U",
        role,
        "-h",
        "127.0.0.1",
        "-d",
        database,
        "-v",
        "ON_ERROR_STOP=1",
        "-At",
        "-q",
        "-c",
        "select current_user;",
    ]
    result = subprocess.run(
        command,
        text=True,
        capture_output=True,
        env={"PATH": bound.SYSTEM_EXEC_PATH, "PGPASSWORD": password},
        check=False,
    )
    if result.returncode != 0 or result.stdout.strip() != role:
        raise bound.CheckpointError(f"ephemeral {role} clone login verification failed")


def _write_secret(path: Path, key: str, value: str) -> None:
    descriptor = os.open(
        path,
        os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0),
        0o600,
    )
    with os.fdopen(descriptor, "w", encoding="utf-8") as handle:
        handle.write(f"{key}={value}\n")
        handle.flush()
        os.fsync(handle.fileno())


def bootstrap(container: str, database: str, output_dir: Path) -> dict[str, Any]:
    target = bound.container_identity(container)
    production = bound.container_identity(bound.PRODUCTION_CONTAINER)
    isolation = bound.validate_disposable_target_identity(target, production)
    target_id = str(target["id"])
    output_dir = output_dir.resolve(strict=True)
    stat = output_dir.stat()
    if not output_dir.is_dir() or stat.st_uid != os.geteuid() or stat.st_mode & 0o077:
        raise bound.CheckpointError("--output-dir must be an existing private directory owned by the current user")
    review_file = output_dir / "kb-review.env"
    apply_file = output_dir / "kb-apply.env"
    if review_file.exists() or review_file.is_symlink() or apply_file.exists() or apply_file.is_symlink():
        raise bound.CheckpointError("clone credential files already exist; use a fresh private run directory")
    review_password = secrets.token_urlsafe(36)
    apply_password = secrets.token_urlsafe(36)
    prereqs = (HERE / "kb_apply_prereqs.sql").read_text(encoding="utf-8")
    _psql(
        target_id,
        database,
        """do $bootstrap$
begin
  if not exists (select 1 from pg_roles where rolname = 'kb_apply') then
    create role kb_apply login;
  end if;
end
$bootstrap$;
""",
    )
    _psql(target_id, database, prereqs)
    _psql(
        target_id,
        database,
        "alter role kb_review password "
        + ap.sql_literal(review_password)
        + ";\nalter role kb_apply password "
        + ap.sql_literal(apply_password)
        + ";\n",
        secret_values=(review_password, apply_password),
    )
    _verify_login(target_id, database, "kb_review", review_password)
    _verify_login(target_id, database, "kb_apply", apply_password)
    manifest = bound._psql_json(
        target_id,
        database,
        """\\set ON_ERROR_STOP on
begin transaction read only;
select jsonb_build_object(
  'approval_table', to_regclass('kb_stage.kb_proposal_approvals')::text,
  'review_principals_table', to_regclass('kb_stage.kb_review_principals')::text,
  'review_role_exists', exists(select 1 from pg_roles where rolname = 'kb_review'),
  'apply_role_exists', exists(select 1 from pg_roles where rolname = 'kb_apply')
)::text;
rollback;
""",
    )
    if manifest != {
        "approval_table": "kb_stage.kb_proposal_approvals",
        "review_principals_table": "kb_stage.kb_review_principals",
        "review_role_exists": True,
        "apply_role_exists": True,
    }:
        raise bound.CheckpointError("clone gate manifest is incomplete: " + json.dumps(manifest, sort_keys=True))
    _write_secret(review_file, "KB_REVIEW_PASSWORD", review_password)
    try:
        _write_secret(apply_file, "KB_APPLY_DB_PASSWORD", apply_password)
    except BaseException:
        review_file.unlink(missing_ok=True)
        raise
    review_password = ""
    apply_password = ""
    return {
        "status": "ready",
        "bound_container_id": target_id,
        "database": database,
        "isolation_checks": isolation,
        "gate_manifest": manifest,
        "review_secret_file": str(review_file),
        "apply_secret_file": str(apply_file),
        "secret_values_or_hashes_recorded": False,
        "login_roles_verified": ["kb_review", "kb_apply"],
    }


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--container", required=True)
    parser.add_argument("--db", required=True)
    parser.add_argument("--output-dir", required=True, type=Path)
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    try:
        if args.container == bound.PRODUCTION_CONTAINER:
            raise bound.CheckpointError("--container must be a disposable non-production target")
        result = bootstrap(args.container, args.db, args.output_dir)
    except (OSError, bound.CheckpointError) as exc:
        print(json.dumps({"status": "rejected", "error": str(exc)}, sort_keys=True))
        return 2
    print(json.dumps(result, sort_keys=True))
    return 0


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