fix(secrets): recover encrypted tokens on all backends
test / integration-docker (pull_request) Successful in 24s
lint / lint (push) Failing after 1m0s
test / unit (pull_request) Successful in 2m3s
test / integration-firecracker (pull_request) Successful in 4m48s
test / coverage (pull_request) Successful in 20s
test / publish-infra (pull_request) Has been skipped
tracker-policy-pr / check-pr (pull_request) Failing after 11m36s

This commit is contained in:
2026-07-22 17:31:39 +00:00
parent c435e9088e
commit 89058fbaec
25 changed files with 517 additions and 57 deletions
@@ -0,0 +1,160 @@
"""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, call, patch
from bot_bottle.orchestrator.reprovision import reprovision_bottles
from bot_bottle.backend.firecracker import consolidated_launch as fc
from bot_bottle.backend.macos_container import consolidated_launch as mac
from bot_bottle.backend.docker import consolidated_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,
]
self.assertEqual(
1,
reprovision_bottles(
client, {"10.0.0.1": "key-1", "10.0.0.2": "key-2"},
),
)
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:
inspect = _proc(stdout=(
"bot-bottle-infra 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.ConsolidatedLaunchError, "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, {})
if __name__ == "__main__":
unittest.main()
+9
View File
@@ -46,6 +46,15 @@ class TestConsolidatedAgentCompose(unittest.TestCase):
# forwarded secrets are bare names (value inherited from process env).
self.assertIn("CLAUDE_CODE_OAUTH_TOKEN", env)
def test_env_var_secret_stays_a_bare_name(self) -> None:
plan = _plan(with_egress=True, supervise=True, with_git=True)
plan = type(plan)(**{**vars(plan), "env_var_secret": "secret-value"})
env = consolidated_agent_compose(
plan, gateway_ip=_GW, source_ip=_IP, network=_NET,
)["services"]["agent"]["environment"]
self.assertIn("ENV_VAR_SECRET", env)
self.assertNotIn("ENV_VAR_SECRET=secret-value", env)
if __name__ == "__main__":
unittest.main()
+5 -1
View File
@@ -91,6 +91,7 @@ class TestTeardownWarning(unittest.TestCase):
bottle_id="b1", identity_token="t", source_ip="172.20.0.4",
network="bot-bottle-gateway", gateway_ip="172.20.0.2",
orchestrator_url="http://orch:8099",
env_var_secret="encryption-key",
)
images = BottleImages(agent="bot-bottle-claude:latest", sidecar="bot-bottle-sidecars:latest")
@@ -107,7 +108,7 @@ class TestTeardownWarning(unittest.TestCase):
mock.patch.object(
launch_mod, "write_compose_file", return_value=Path("/tmp/compose.yml"),
), \
mock.patch.object(launch_mod, "compose_up"), \
mock.patch.object(launch_mod, "compose_up") as compose_up, \
mock.patch.object(launch_mod, "compose_dump_logs"), \
mock.patch.object(
launch_mod, "compose_down",
@@ -122,6 +123,9 @@ class TestTeardownWarning(unittest.TestCase):
self.assertIn("bot-bottle: warning:", output)
self.assertIn("bot-bottle-test-teardown-abc", output)
self.assertIn("compose-down", output)
self.assertEqual(
"encryption-key", compose_up.call_args.kwargs["env"]["ENV_VAR_SECRET"],
)
if __name__ == "__main__":
+1
View File
@@ -66,6 +66,7 @@ class TestNetpoolRenderers(unittest.TestCase):
self.assertIn("chown node:node /home/node", util._GUEST_INIT)
self.assertIn("chmod 755 /home/node", util._GUEST_INIT)
self.assertIn("mount -t tmpfs -o mode=0755 tmpfs /run", util._GUEST_INIT)
def test_nixos_module_is_non_invasive(self):
# The NixOS module must NOT flip the host firewall backend or
+8
View File
@@ -62,6 +62,14 @@ class TestProcessScan(unittest.TestCase):
with patch.object(fc_cleanup.subprocess, "run", return_value=_proc(returncode=1)):
self.assertEqual((set(), []), fc_cleanup._scan_processes(Path("/x")))
def test_live_run_dirs_returns_paths_in_stable_order(self):
with patch.object(fc_cleanup, "_run_root", return_value=Path("/run")), \
patch.object(fc_cleanup, "_scan_processes",
return_value=({"/run/b", "/run/a"}, [])):
self.assertEqual(
(Path("/run/a"), Path("/run/b")), fc_cleanup.live_run_dirs(),
)
def test_orphan_run_dirs_excludes_live_and_missing_root(self):
with tempfile.TemporaryDirectory() as tmp:
run_root = Path(tmp)
@@ -49,6 +49,7 @@ def _plan(
agent_git_gate_url: str = "",
agent_supervise_url: str = "",
image_policy: str = "fresh",
env_var_secret: str = "",
) -> MacosContainerBottlePlan:
routes_path = stage_dir / "routes.yaml"
routes_path.write_text("routes: []\n", encoding="utf-8")
@@ -80,6 +81,7 @@ def _plan(
),
agent_git_gate_url=agent_git_gate_url,
agent_supervise_url=agent_supervise_url,
env_var_secret=env_var_secret,
))
@@ -185,6 +187,12 @@ class TestAgentRunArgv(unittest.TestCase):
"bot-bottle-mac-gateway", self.argv[self.argv.index("--network") + 1],
)
def test_env_var_secret_is_in_configured_container_environment(self) -> None:
argv = _agent_run_argv(
_plan(Path(self._tmp.name), env_var_secret="key-material"), _endpoint(),
)
self.assertIn("ENV_VAR_SECRET=key-material", argv)
def test_never_pins_an_ip(self) -> None:
"""Apple Container 1.0.0 has no --ip: the address is DHCP-assigned and
read back after start."""
+15
View File
@@ -28,6 +28,21 @@ class TestMacosContainerAvailability(unittest.TestCase):
class TestMacosContainerCommands(unittest.TestCase):
def test_read_container_env(self):
completed = util.subprocess.CompletedProcess(
args=[], returncode=0, stdout="secret\n", stderr="",
)
with patch.object(util, "_run_container_op", return_value=completed) as run:
self.assertEqual("secret", util.read_container_env("bottle", "KEY"))
run.assert_called_once_with(["container", "exec", "bottle", "printenv", "KEY"])
def test_read_container_env_returns_empty_on_failure(self):
completed = util.subprocess.CompletedProcess(
args=[], returncode=1, stdout="", stderr="missing",
)
with patch.object(util, "_run_container_op", return_value=completed):
self.assertEqual("", util.read_container_env("bottle", "KEY"))
def test_dns_server_prefers_direct_host_ipv4_resolver(self):
scutil = util.subprocess.CompletedProcess(
args=[],
+21
View File
@@ -76,6 +76,27 @@ class TestTeardown(unittest.TestCase):
self.assertEqual("DELETE", m.call_args.args[0].get_method())
class TestReprovisionGateway(unittest.TestCase):
def setUp(self) -> None:
self.c = OrchestratorClient("http://orch:8080")
def test_success_posts_key(self) -> None:
with patch(_URLOPEN, return_value=_resp(200, {"reprovisioned": True})) as opened:
self.assertTrue(self.c.reprovision_gateway("b1", "key"))
request = opened.call_args.args[0]
self.assertEqual("POST", request.get_method())
self.assertEqual({"env_var_secret": "key"}, json.loads(request.data))
def test_missing_stored_secret_is_false(self) -> None:
with patch(_URLOPEN, side_effect=_http_error(404)):
self.assertFalse(self.c.reprovision_gateway("b1", "key"))
def test_other_status_raises(self) -> None:
with patch(_URLOPEN, side_effect=_http_error(400)):
with self.assertRaises(OrchestratorClientError):
self.c.reprovision_gateway("b1", "key")
class TestHealthAndPolicy(unittest.TestCase):
def setUp(self) -> None:
self.c = OrchestratorClient("http://orch:8080")
@@ -6,6 +6,7 @@ server tests), plus one real-socket round-trip to prove the handler wiring.
from __future__ import annotations
import base64
import json
import secrets
import sqlite3
@@ -63,6 +64,42 @@ class TestDispatch(unittest.TestCase):
self.assertTrue(payload["bottle_id"])
self.assertTrue(payload["identity_token"])
def test_register_and_reprovision_encrypted_tokens(self) -> None:
key = base64.urlsafe_b64encode(b"unit-test-key").rstrip(b"=").decode()
status, payload = dispatch(
self.orch, "POST", "/bottles", _body({
"source_ip": "10.243.0.11",
"tokens": {"EGRESS_TOKEN_0": "upstream-secret"},
"env_var_secret": key,
}),
)
self.assertEqual(201, status)
bottle_id = payload["bottle_id"]
self.orch._tokens.clear()
status, response = dispatch(
self.orch, "POST", f"/bottles/{bottle_id}/reprovision_gateway",
_body({"env_var_secret": key}),
)
self.assertEqual((200, {"reprovisioned": True}), (status, response))
self.assertEqual(
{"EGRESS_TOKEN_0": "upstream-secret"}, self.orch.tokens_for(bottle_id),
)
def test_reprovision_validates_request_and_missing_rows(self) -> None:
status, _ = dispatch(
self.orch, "POST", "/bottles/b1/reprovision_gateway", b"not-json",
)
self.assertEqual(400, status)
status, _ = dispatch(
self.orch, "POST", "/bottles/b1/reprovision_gateway", _body({}),
)
self.assertEqual(400, status)
status, _ = dispatch(
self.orch, "POST", "/bottles/b1/reprovision_gateway",
_body({"env_var_secret": "key"}),
)
self.assertEqual(404, status)
def test_register_requires_source_ip(self) -> None:
status, _ = dispatch(self.orch, "POST", "/bottles", _body({}))
self.assertEqual(400, status)
+15 -9
View File
@@ -68,15 +68,21 @@ class TestDecryptErrors(unittest.TestCase):
def test_wrong_key_raises_value_error(self) -> None:
ct = encrypt_value(self.secret, "secret-token")
other_key = new_env_var_secret()
# Wrong key produces garbage bytes; decrypt_value raises ValueError
# when the result is non-UTF-8 (which is very likely for 12-char data).
# We allow it to succeed only if garbage happens to be valid UTF-8, but
# the plaintext must not match.
try:
result = decrypt_value(other_key, ct)
self.assertNotEqual("secret-token", result)
except ValueError:
pass
with self.assertRaisesRegex(ValueError, "authentication failed"):
decrypt_value(other_key, ct)
def test_modified_ciphertext_raises_value_error(self) -> None:
ct = encrypt_value(self.secret, "secret-token")
replacement = "A" if ct[-10] != "A" else "B"
tampered = ct[:-10] + replacement + ct[-9:]
with self.assertRaisesRegex(ValueError, "authentication failed"):
decrypt_value(self.secret, tampered)
def test_modified_tag_raises_value_error(self) -> None:
ct = encrypt_value(self.secret, "secret-token")
replacement = "A" if ct[-1] != "A" else "B"
with self.assertRaisesRegex(ValueError, "authentication failed"):
decrypt_value(self.secret, ct[:-1] + replacement)
def test_truncated_blob_raises_value_error(self) -> None:
with self.assertRaises(ValueError):
+20
View File
@@ -15,6 +15,7 @@ from bot_bottle.orchestrator.broker import LaunchBroker, LaunchRequest, StubBrok
from bot_bottle.orchestrator.registry import RegistryStore
from bot_bottle.orchestrator.service import Orchestrator
from bot_bottle.orchestrator.gateway import Gateway
from bot_bottle.orchestrator.secret_store import new_env_var_secret
from bot_bottle.store_manager import StoreManager
from bot_bottle.supervise import (
Proposal,
@@ -117,6 +118,25 @@ class TestOrchestrator(unittest.TestCase):
rec = self.orch.launch_bottle("10.243.0.6")
self.assertEqual({}, self.orch.tokens_for(rec.bottle_id))
def test_encrypted_tokens_can_be_reprovisioned_after_memory_loss(self) -> None:
key = new_env_var_secret()
rec = self.orch.launch_bottle(
"10.243.0.12", tokens={"EGRESS_TOKEN_0": "secret"},
env_var_secret=key,
)
self.assertNotEqual({}, self.store.get_agent_secrets(rec.bottle_id))
self.orch._tokens.clear()
self.assertTrue(self.orch.reprovision_from_secret(rec.bottle_id, key))
self.assertEqual({"EGRESS_TOKEN_0": "secret"}, self.orch.tokens_for(rec.bottle_id))
def test_reprovision_rejects_missing_rows_and_wrong_key(self) -> None:
self.assertFalse(self.orch.reprovision_from_secret("missing", new_env_var_secret()))
rec = self.orch.launch_bottle(
"10.243.0.13", tokens={"K": "value"},
env_var_secret=new_env_var_secret(),
)
self.assertFalse(self.orch.reprovision_from_secret(rec.bottle_id, new_env_var_secret()))
def test_set_policy_live_reload(self) -> None:
rec = self.orch.launch_bottle("10.243.0.3")
self.assertTrue(self.orch.set_policy(rec.bottle_id, '{"x":1}'))