bot-bottle/tests/unit/test_backend_secret_reprovision.py

"""Backend-agnostic and backend-specific encrypted-secret recovery tests."""

from __future__ import annotations

import json
import subprocess
import tempfile
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import Mock, patch

from bot_bottle.orchestrator.reprovision import reprovision_bottles
from bot_bottle.backend.base import BottleBackend, InfraLaunchError
from bot_bottle.backend.gateway_attach import GatewayAttachResources
from bot_bottle.backend.firecracker import infra_launch as fc
from bot_bottle.backend.macos_container import infra_launch as mac
from bot_bottle.backend.docker import infra_launch as docker
from bot_bottle.orchestrator.client import OrchestratorClientError


def _proc(returncode: int = 0, stdout: str = "", stderr: str = ""):
    return subprocess.CompletedProcess([], returncode, stdout=stdout, stderr=stderr)


class TestSharedReprovision(unittest.TestCase):
    def test_joins_registry_records_by_source_ip(self) -> None:
        client = Mock()
        client.list_bottles.return_value = [
            {"bottle_id": "b1", "source_ip": "10.0.0.1"},
            {"bottle_id": "b2", "source_ip": "10.0.0.2"},
            {"bottle_id": 3, "source_ip": "10.0.0.3"},
        ]
        client.reprovision_gateway.side_effect = [True, False]
        count = reprovision_bottles(
            client, {"10.0.0.1": " key-1\n", "10.0.0.2": "key-2"},
        )
        self.assertEqual(1, count)
        self.assertEqual(
            [("b1", "key-1"), ("b2", "key-2")],
            [call.args for call in client.reprovision_gateway.call_args_list],
        )

    def test_one_failure_does_not_block_other_bottles(self) -> None:
        client = Mock()
        client.list_bottles.return_value = [
            {"bottle_id": "b1", "source_ip": "10.0.0.1"},
            {"bottle_id": "b2", "source_ip": "10.0.0.2"},
        ]
        client.reprovision_gateway.side_effect = [
            OrchestratorClientError("bad key"), True,
        ]
        with patch("bot_bottle.orchestrator.reprovision.debug") as debug:
            count = reprovision_bottles(
                client, {"10.0.0.1": "key-1", "10.0.0.2": "key-2"},
            )
        self.assertEqual(1, count)
        self.assertEqual("b1", debug.call_args.kwargs["context"]["bottle_id"])
        self.assertNotIn("bad key", repr(debug.call_args))


class TestMacosReprovision(unittest.TestCase):
    def test_reads_configured_container_env_and_reprovisions(self) -> None:
        endpoint = mac.GatewayEndpoint("http://orch", "10.0.0.9", "PEM", "net")
        agent = SimpleNamespace(slug="demo")
        client = Mock()
        with patch.object(mac, "enumerate_active", return_value=[agent]), \
             patch.object(mac.container_mod, "inspect_container_network_ip",
                          return_value="10.0.0.1"), \
             patch.object(mac.container_mod, "read_container_env", return_value="key"), \
             patch.object(mac, "OrchestratorClient", return_value=client), \
             patch.object(mac, "reprovision_bottles", return_value=1) as restore, \
             patch.object(mac, "info"):
            mac._reprovision_running_bottles(endpoint)
        restore.assert_called_once_with(client, {"10.0.0.1": "key"})

    def test_enumeration_failure_is_best_effort(self) -> None:
        endpoint = mac.GatewayEndpoint("http://orch", "10.0.0.9", "PEM", "net")
        with patch.object(mac, "enumerate_active",
                          side_effect=mac.EnumerationError("failed")), \
             patch.object(mac, "info") as info:
            mac._reprovision_running_bottles(endpoint)
        self.assertIn("skipped", info.call_args.args[0])


class TestDockerReprovision(unittest.TestCase):
    def test_maps_network_containers_to_keys(self) -> None:
        # The gateway container (excluded — it's on the data network but isn't
        # an agent) + one agent + a malformed line.
        inspect = _proc(stdout=(
            "bot-bottle-orch-gateway 172.18.0.2/16\n"
            "bot-bottle-a 172.18.0.3/16\n"
            "malformed\n"
        ))
        key = _proc(stdout="secret\n")
        client = Mock()
        with patch.object(docker, "OrchestratorClient", return_value=client), \
             patch.object(docker, "run_docker", side_effect=[inspect, key]), \
             patch.object(docker, "reprovision_bottles", return_value=1) as restore, \
             patch.object(docker.log, "info"):
            docker._reprovision_running_bottles("http://orch")
        restore.assert_called_once_with(client, {"172.18.0.3": "secret"})

    def test_missing_docker_is_best_effort(self) -> None:
        with patch.object(docker, "run_docker", side_effect=FileNotFoundError("docker")), \
             patch.object(docker.log, "info") as info:
            docker._reprovision_running_bottles("http://orch")
        self.assertIn("skipped", info.call_args.args[0])


class TestFirecrackerReprovision(unittest.TestCase):
    def _run_dir(self, root: Path, ip: str = "10.243.0.3") -> Path:
        run_dir = root / "demo"
        run_dir.mkdir()
        (run_dir / "bottle_id_ed25519").write_text("key")
        (run_dir / "config.json").write_text(json.dumps({
            "boot-source": {"boot_args": f"root=/dev/vda ip={ip}::gw:mask::eth0:off"}
        }))
        return run_dir

    def test_extracts_guest_ip_from_config(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            run_dir = self._run_dir(Path(tmp))
            self.assertEqual("10.243.0.3", fc._guest_ip_from_config(run_dir / "config.json"))
            self.assertEqual("", fc._guest_ip_from_config(run_dir / "missing.json"))

    def test_persists_key_over_stdin_not_argv(self) -> None:
        with patch.object(fc.util, "ssh_base_argv", return_value=["ssh", "guest"]), \
             patch.object(fc.subprocess, "run", return_value=_proc()) as run:
            fc.persist_env_var_secret(Path("/key"), "10.0.0.1", "super-secret")
        self.assertEqual("super-secret", run.call_args.kwargs["input"])
        self.assertNotIn("super-secret", " ".join(run.call_args.args[0]))

    def test_persist_failure_is_fatal_to_launch(self) -> None:
        with patch.object(fc.util, "ssh_base_argv", return_value=["ssh", "guest"]), \
             patch.object(fc.subprocess, "run", return_value=_proc(1, stderr="denied")):
            with self.assertRaisesRegex(fc.InfraLaunchError, "denied"):
                fc.persist_env_var_secret(Path("/key"), "10.0.0.1", "secret")

    def test_reads_live_vm_key_and_reprovisions(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            run_dir = self._run_dir(Path(tmp))
            client = Mock()
            with patch.object(fc.cleanup, "live_run_dirs", return_value=(run_dir,)), \
                 patch.object(fc.util, "ssh_base_argv", return_value=["ssh", "guest"]), \
                 patch.object(fc.subprocess, "run", return_value=_proc(stdout="secret\n")), \
                 patch.object(fc, "reprovision_bottles", return_value=1) as restore, \
                 patch.object(fc, "info"):
                fc._reprovision_running_bottles(client)
        restore.assert_called_once_with(client, {"10.243.0.3": "secret"})

    def test_unreadable_vm_is_skipped(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            run_dir = self._run_dir(Path(tmp))
            client = Mock()
            with patch.object(fc.cleanup, "live_run_dirs", return_value=(run_dir,)), \
                 patch.object(fc.util, "ssh_base_argv", return_value=["ssh", "guest"]), \
                 patch.object(fc.subprocess, "run", return_value=_proc(1)), \
                 patch.object(fc, "reprovision_bottles", return_value=0) as restore:
                fc._reprovision_running_bottles(client)
        restore.assert_called_once_with(client, {})


class TestAttachFlow(unittest.TestCase):
    """PRD 0083 / #519 review: the base backend owns the reconcile flow and
    fails hard — gather resources, attempt every running bottle, then raise an
    aggregate if any failed (never silently skip)."""

    def _fake(self, **primitives: object) -> SimpleNamespace:
        base: dict[str, object] = {
            "_gateway_attach_resources": Mock(
                return_value=GatewayAttachResources(ca_pem="PEM")),
            "_running_bottles": Mock(return_value=["a", "b", "c"]),
            "_attach_bottle_to_gateway": Mock(),
        }
        base.update(primitives)
        return SimpleNamespace(**base)

    def test_attaches_every_running_bottle_once(self) -> None:
        fake = self._fake()
        BottleBackend.attach_bottled_agents_to_gateway(fake)  # type: ignore[arg-type]
        self.assertEqual(3, fake._attach_bottle_to_gateway.call_count)
        # Each bottle got the single gathered resource set.
        res = fake._gateway_attach_resources.return_value
        for call in fake._attach_bottle_to_gateway.call_args_list:
            self.assertIs(res, call.args[1])

    def test_attempts_all_then_raises_aggregate_on_failure(self) -> None:
        def attach(bottle: str, _resources: object) -> None:
            if bottle in ("b", "c"):
                raise InfraLaunchError(f"{bottle} unreachable")
        fake = self._fake(_attach_bottle_to_gateway=Mock(side_effect=attach))
        with self.assertRaises(InfraLaunchError) as ctx:
            BottleBackend.attach_bottled_agents_to_gateway(fake)  # type: ignore[arg-type]
        msg = str(ctx.exception)
        # Every bottle was attempted (not fail-fast), and both failures surface.
        self.assertEqual(3, fake._attach_bottle_to_gateway.call_count)
        self.assertIn("b unreachable", msg)
        self.assertIn("c unreachable", msg)
        self.assertIn("2", msg)

    def test_resource_gathering_failure_fails_hard_before_enumerating(self) -> None:
        fake = self._fake(
            _gateway_attach_resources=Mock(side_effect=RuntimeError("no CA yet")))
        with self.assertRaises(RuntimeError):
            BottleBackend.attach_bottled_agents_to_gateway(fake)  # type: ignore[arg-type]
        fake._running_bottles.assert_not_called()


class TestFirecrackerAttachPrimitive(unittest.TestCase):
    """PRD 0083: attach_ca_to_agent_vm pushes the current gateway CA into one
    running agent VM over SSH, failing hard."""

    def _run_dir(self, root: Path, ip: str = "10.243.0.3") -> Path:
        run_dir = root / "demo"
        run_dir.mkdir()
        (run_dir / "bottle_id_ed25519").write_text("key")
        (run_dir / "config.json").write_text(json.dumps({
            "boot-source": {"boot_args": f"root=/dev/vda ip={ip}::gw:mask::eth0:off"}
        }))
        return run_dir

    def test_pushes_ca_over_stdin_and_rebuilds_trust_store(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            run_dir = self._run_dir(Path(tmp))
            with patch.object(fc.util, "ssh_base_argv", return_value=["ssh", "guest"]), \
                 patch.object(fc.subprocess, "run", return_value=_proc()) as run:
                fc.attach_ca_to_agent_vm(run_dir, "PEM-DATA")
        # The cert goes over stdin (off argv), and the trust store is rebuilt.
        self.assertEqual("PEM-DATA", run.call_args.kwargs["input"])
        script = run.call_args.args[0][-1]
        self.assertIn(fc.AGENT_CA_PATH, script)
        self.assertIn("update-ca-certificates", script)

    def test_malformed_run_dir_fails_hard(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            empty = Path(tmp) / "no-config"
            empty.mkdir()
            with self.assertRaises(fc.InfraLaunchError):
                fc.attach_ca_to_agent_vm(empty, "PEM")

    def test_push_failure_fails_hard(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            run_dir = self._run_dir(Path(tmp))
            with patch.object(fc.util, "ssh_base_argv", return_value=["ssh", "guest"]), \
                 patch.object(fc.subprocess, "run", return_value=_proc(1, stderr="down")):
                with self.assertRaisesRegex(fc.InfraLaunchError, "down"):
                    fc.attach_ca_to_agent_vm(run_dir, "PEM")


class TestDockerAttachPrimitives(unittest.TestCase):
    def test_running_agent_containers_excludes_gateway(self) -> None:
        # network inspect lists the gateway (excluded) + one agent + a blank line.
        inspect = _proc(stdout="bot-bottle-orch-gateway\nbot-bottle-a\n\n")
        with patch.object(docker, "run_docker", return_value=inspect):
            names = docker.running_agent_containers(
                network="net", gateway_name="bot-bottle-orch-gateway")
        self.assertEqual(["bot-bottle-a"], names)

    def test_running_agent_containers_fails_hard_on_docker_error(self) -> None:
        with patch.object(docker, "run_docker", side_effect=FileNotFoundError("docker")):
            with self.assertRaises(FileNotFoundError):
                docker.running_agent_containers()

    def test_push_ca_cp_then_rebuilds_trust_store(self) -> None:
        with patch.object(
            docker, "run_docker", side_effect=[_proc(), _proc(), _proc()],
        ) as run:
            docker.push_ca_to_container("bot-bottle-a", "PEM")
        argvs = [c.args[0] for c in run.call_args_list]
        cp = next(a for a in argvs if a[:2] == ["docker", "cp"])
        self.assertEqual(f"bot-bottle-a:{docker.AGENT_CA_PATH}", cp[-1])
        self.assertTrue(any("update-ca-certificates" in a[-1] for a in argvs))

    def test_push_ca_fails_hard(self) -> None:
        with patch.object(docker, "run_docker", return_value=_proc(1, stderr="gone")):
            with self.assertRaisesRegex(docker.InfraLaunchError, "gone"):
                docker.push_ca_to_container("bot-bottle-a", "PEM")

    def test_gateway_attach_resources_reads_the_threaded_infra_gateway(self) -> None:
        # On the bring-up reconcile path the backend is threaded with the infra
        # whose gateway just cold-booted, so it reads THAT gateway's CA — not a
        # fresh default-named DockerInfraService that wouldn't exist for a
        # non-default instance (#519 review).
        from bot_bottle.backend.docker.backend import DockerBottleBackend
        infra = Mock()
        infra.gateway.return_value.ca_cert_pem.return_value = "ITEST-CA"
        res = DockerBottleBackend(infra=infra)._gateway_attach_resources()
        self.assertEqual("ITEST-CA", res.ca_pem)
        infra.gateway.assert_called_once()

    def test_gateway_attach_resources_defaults_to_the_host_singleton(self) -> None:
        # Ordinary construction (no infra) falls back to the per-host default.
        from bot_bottle.backend.docker.backend import DockerBottleBackend
        with patch("bot_bottle.backend.docker.infra.DockerInfraService") as InfraCls:
            InfraCls.return_value.gateway.return_value.ca_cert_pem.return_value = "DEFAULT-CA"
            res = DockerBottleBackend()._gateway_attach_resources()
        InfraCls.assert_called_once_with()
        self.assertEqual("DEFAULT-CA", res.ca_pem)


class TestMacosAttachPrimitives(unittest.TestCase):
    def test_running_agent_containers_from_enumeration(self) -> None:
        with patch.object(mac, "enumerate_active",
                          return_value=[SimpleNamespace(slug="demo")]):
            names = mac.running_agent_containers()
        self.assertEqual([f"{mac.CONTAINER_NAME_PREFIX}demo"], names)

    def test_running_agent_containers_fails_hard(self) -> None:
        with patch.object(mac, "enumerate_active",
                          side_effect=mac.EnumerationError("failed")):
            with self.assertRaises(mac.EnumerationError):
                mac.running_agent_containers()

    def test_push_ca_cp_then_rebuilds_trust_store(self) -> None:
        with patch.object(mac.container_mod, "run_container_argv",
                          return_value=_proc()) as run:
            mac.push_ca_to_container(f"{mac.CONTAINER_NAME_PREFIX}demo", "PEM")
        argvs = [c.args[0] for c in run.call_args_list]
        self.assertTrue(any(mac.AGENT_CA_PATH in " ".join(a) for a in argvs))
        self.assertTrue(any("update-ca-certificates" in a[-1] for a in argvs))

    def test_push_ca_fails_hard(self) -> None:
        with patch.object(mac.container_mod, "run_container_argv",
                          return_value=_proc(1, stderr="no")):
            with self.assertRaisesRegex(mac.InfraLaunchError, "no"):
                mac.push_ca_to_container("x", "PEM")


if __name__ == "__main__":
    unittest.main()