bot-bottle/bot_bottle/backend/firecracker/infra_launch.py

"""Consolidated bottle launch sequence for the Firecracker backend
(PRD 0070, Stage B).

The orchestrator control plane and the shared gateway run in **two** persistent
per-host **infra VMs** (`infra_vm.py`), split per PRD 0070 — not Docker
containers. Agent VMs reach the gateway's egress / supervise / git-http ports
via a PREROUTING DNAT on their own host-side TAP IP that redirects to the
**gateway** VM (see `scripts/firecracker-netpool.sh`) — never the orchestrator,
so a breached agent has no L3 route to the control plane. The host CLI reaches
the control plane over HTTP at the orchestrator VM's guest IP.

Attribution is by the agent VM's guest IP, unspoofable by construction: the
/31 point-to-point TAP + the `bot_bottle_fc` nft table ensure only the
expected VM can source-IP that address.

Sequence:
  1. ensure the infra VMs (orchestrator control plane + gateway data plane)
     are up (a singleton pair — a prior launcher may already have booted them);
  2. register the bottle on the orchestrator by its guest IP (attribution key)
     → bottle id + identity token;
  3. provision its git-gate repos/creds into the gateway VM (over SSH);
  4. fetch the shared gateway CA for the provisioner to install in the rootfs.

The TAP slot allocation, rootfs build, and VM boot are the caller's job.
"""

from __future__ import annotations

import json
import subprocess
from dataclasses import dataclass
from pathlib import Path

from ...egress import EgressPlan
from ...git_gate import GitGatePlan
from ...log import info
from ...orchestrator.client import OrchestratorClient, OrchestratorClientError
from ...orchestrator.lifecycle import (
    OrchestratorStartError,  # re-exported so callers can catch it
)
from ...orchestrator.reprovision import reprovision_bottles
from ...orchestrator.store.secret_store import ENV_VAR_SECRET_NAME
from ..base import InfraLaunchError
from ..provision_bottle import deprovision_bottle, provision_bottle
from ..util import AGENT_CA_PATH
from . import cleanup, util
from .gateway import FirecrackerGateway
from .infra import FirecrackerInfraService

_ENV_VAR_SECRET_PATH = "/run/bot-bottle/env-var-secret"


@dataclass(frozen=True)
class LaunchContext:
    """What the Firecracker launch needs from the consolidated sequence."""

    bottle_id: str
    identity_token: str
    source_ip: str          # the VM's guest IP — the attribution key
    gateway_ca_pem: str     # the shared gateway CA the provisioner installs
    orchestrator_url: str
    env_var_secret: str = ""  # encryption key injected into the agent's env


def _guest_ip_from_config(config_path: Path) -> str:
    """Read the kernel's configured guest IP from a Firecracker config."""
    try:
        config = json.loads(config_path.read_text())
        args = config["boot-source"]["boot_args"]
        ip_arg = next(part for part in args.split() if part.startswith("ip="))
        return ip_arg.removeprefix("ip=").split(":", 1)[0]
    except (OSError, ValueError, KeyError, TypeError, StopIteration):
        return ""


def persist_env_var_secret(private_key: Path, guest_ip: str, secret: str) -> None:
    """Mirror the exec-time key into guest tmpfs for restart recovery."""
    proc = subprocess.run(
        util.ssh_base_argv(private_key, guest_ip)
        + [f"umask 077; mkdir -p /run/bot-bottle; cat > {_ENV_VAR_SECRET_PATH}"],
        input=secret, capture_output=True, text=True, check=False,
    )
    if proc.returncode != 0:
        raise InfraLaunchError(
            f"failed to persist {ENV_VAR_SECRET_NAME} in agent VM: "
            f"{proc.stderr.strip() or '<no stderr>'}"
        )


def attach_ca_to_agent_vm(run_dir: Path, ca_pem: str) -> None:
    """(Re)attach one running agent VM (identified by its `run_dir`) to the
    current gateway: replace its trusted gateway CA with `ca_pem` and rebuild
    its trust store over SSH (PRD 0083). Unconditional install — there is one
    gateway, so no fingerprint match is needed. Bare-pipe input keeps the cert
    off argv.

    Fail hard: raises `InfraLaunchError` (naming the bottle) if the run dir is
    malformed or the push fails — the base reconcile aggregates and surfaces it
    rather than leaving the agent silently unable to reach the gateway."""
    guest_ip = _guest_ip_from_config(run_dir / "config.json")
    private_key = run_dir / "bottle_id_ed25519"
    if not guest_ip or not private_key.is_file():
        raise InfraLaunchError(
            f"{run_dir.name}: cannot resolve guest IP / SSH key to attach it "
            f"to the gateway"
        )
    install = (
        "umask 022; mkdir -p /usr/local/share/ca-certificates && "
        f"cat > {AGENT_CA_PATH} && chmod 644 {AGENT_CA_PATH} && "
        "update-ca-certificates"
    )
    try:
        proc = subprocess.run(
            util.ssh_base_argv(private_key, guest_ip) + [install],
            input=ca_pem, capture_output=True, text=True, check=False,
        )
    except OSError as exc:
        raise InfraLaunchError(f"{run_dir.name} ({guest_ip}): CA push failed: {exc}") from exc
    if proc.returncode != 0:
        raise InfraLaunchError(
            f"{run_dir.name} ({guest_ip}): CA push failed: "
            f"{proc.stderr.strip() or '<no stderr>'}"
        )


def _reprovision_running_bottles(client: OrchestratorClient) -> None:
    """Read keys from live agent VMs and restore the restarted gateway."""
    try:
        secrets_by_ip: dict[str, str] = {}
        for run_dir in cleanup.live_run_dirs():
            guest_ip = _guest_ip_from_config(run_dir / "config.json")
            private_key = run_dir / "bottle_id_ed25519"
            if not guest_ip or not private_key.is_file():
                continue
            proc = subprocess.run(
                util.ssh_base_argv(private_key, guest_ip)
                + [f"cat {_ENV_VAR_SECRET_PATH}"],
                capture_output=True, text=True, check=False,
            )
            if proc.returncode == 0 and proc.stdout.strip():
                secrets_by_ip[guest_ip] = proc.stdout.strip()
        count = reprovision_bottles(client, secrets_by_ip)
        if count:
            info(f"reprovisioned egress tokens for {count} Firecracker bottle(s)")
    except (OSError, OrchestratorClientError) as exc:
        info(f"egress token reprovision skipped: {exc}")


def launch_consolidated(
    egress_plan: EgressPlan,
    git_gate_plan: GitGatePlan,
    *,
    guest_ip: str,
    image_ref: str = "",
    tokens: dict[str, str] | None = None,
) -> LaunchContext:
    """Ensure the infra VM is up, register the bottle by its guest IP, and
    provision its git-gate state into the gateway VM. Returns the context the
    agent-VM launch needs. Raises on failure — the caller tears down."""
    service = FirecrackerInfraService()
    url = service.ensure_running()
    client = OrchestratorClient(url)
    _reprovision_running_bottles(client)

    # Read the gateway's provisioning transport + CA off the Gateway service.
    gateway = service.gateway()
    transport = gateway.provisioning_transport()
    reg = provision_bottle(
        client, guest_ip, egress_plan, git_gate_plan, transport,
        image_ref=image_ref, tokens=tokens,
    )
    # The shared gateway CA every agent on this host trusts for TLS
    # interception — fetched from the gateway VM over SSH.
    return LaunchContext(
        bottle_id=reg.bottle_id,
        identity_token=reg.identity_token,
        source_ip=guest_ip,
        gateway_ca_pem=gateway.ca_cert_pem(),
        orchestrator_url=url,
        env_var_secret=reg.env_var_secret,
    )


def deprovision_consolidated(
    bottle_id: str, *, orchestrator_url: str, timeout: float | None = None,
) -> None:
    """Deregister the bottle and remove its git-gate state from the gateway
    VM. Both steps are idempotent so this is safe from a cleanup trap. Does
    NOT stop the infra VMs — they're persistent per-host singletons shared by
    every bottle."""
    deprovision_bottle(bottle_id, FirecrackerGateway().provisioning_transport(),
                   orchestrator_url=orchestrator_url, timeout=timeout)


__all__ = [
    "LaunchContext",
    "launch_consolidated",
    "attach_ca_to_agent_vm",
    "deprovision_consolidated",
    "InfraLaunchError",
    "OrchestratorStartError",
]