"""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 0081 / #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 0081: 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") 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()