Compare commits
1 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 01cee056be |
@@ -7,10 +7,13 @@ imports it rather than re-implementing it.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
|
||||
from ..egress import EgressPlan
|
||||
from ..git_gate import GitGatePlan
|
||||
from ..orchestrator.client import OrchestratorClient
|
||||
from ..orchestrator.client import OrchestratorClient, RegisteredBottle
|
||||
from ..orchestrator.registration import registration_inputs
|
||||
from ..orchestrator.secret_store import new_env_var_secret
|
||||
from .docker.gateway_provision import GatewayTransport, deprovision_git_gate, provision_git_gate
|
||||
|
||||
|
||||
@@ -23,21 +26,26 @@ def provision_bottle(
|
||||
*,
|
||||
image_ref: str = "",
|
||||
tokens: dict[str, str] | None = None,
|
||||
):
|
||||
) -> RegisteredBottle:
|
||||
"""Register the bottle and provision its git-gate state. Rolls back the
|
||||
registration if provisioning fails so no orphan is left. Returns the
|
||||
`RegisteredBottle` from the orchestrator."""
|
||||
registration if provisioning fails so no orphan is left.
|
||||
|
||||
Generates a fresh ENV_VAR_SECRET, passes it to the orchestrator so it can
|
||||
encrypt the token values at rest, and stamps the secret onto the returned
|
||||
``RegisteredBottle`` so callers can inject it into the agent container's
|
||||
environment."""
|
||||
inputs = registration_inputs(egress_plan)
|
||||
env_var_secret = new_env_var_secret()
|
||||
reg = client.register_bottle(
|
||||
source_ip, image_ref=image_ref, policy=inputs.policy,
|
||||
metadata=inputs.metadata, tokens=tokens,
|
||||
metadata=inputs.metadata, tokens=tokens, env_var_secret=env_var_secret,
|
||||
)
|
||||
try:
|
||||
provision_git_gate(transport, reg.bottle_id, git_gate_plan)
|
||||
except Exception:
|
||||
client.teardown_bottle(reg.bottle_id)
|
||||
raise
|
||||
return reg
|
||||
return dataclasses.replace(reg, env_var_secret=env_var_secret)
|
||||
|
||||
|
||||
def teardown_consolidated(
|
||||
|
||||
@@ -39,6 +39,10 @@ class DockerBottlePlan(BottlePlan):
|
||||
# (egress proxy credentials, git-gate/supervise headers); set by launch
|
||||
# from the orchestrator registration. Empty pre-registration.
|
||||
identity_token: str = ""
|
||||
# Encryption key for the agent's stored egress secrets; injected into the
|
||||
# agent container as ENV_VAR_SECRET via the compose subprocess env (bare
|
||||
# name — value never written to the compose file). Empty pre-registration.
|
||||
env_var_secret: str = ""
|
||||
|
||||
@property
|
||||
def container_name(self) -> str:
|
||||
|
||||
@@ -17,6 +17,7 @@ from __future__ import annotations
|
||||
from typing import Any
|
||||
|
||||
from ...egress import egress_agent_env_entries
|
||||
from ...orchestrator.secret_store import ENV_VAR_SECRET_NAME
|
||||
from ..util import AGENT_CA_BUNDLE, AGENT_CA_PATH
|
||||
from .bottle_plan import DockerBottlePlan
|
||||
from .egress import EGRESS_PORT
|
||||
@@ -58,6 +59,10 @@ def consolidated_agent_compose(
|
||||
# the secret value never lands on argv or in the compose file.
|
||||
for name in sorted(plan.forwarded_env.keys()):
|
||||
env.append(name)
|
||||
# ENV_VAR_SECRET: bare name so the value comes from the compose subprocess
|
||||
# env (set in launch.py) and is never written to the compose file on disk.
|
||||
if getattr(plan, "env_var_secret", ""):
|
||||
env.append(ENV_VAR_SECRET_NAME)
|
||||
env.extend(egress_agent_env_entries(plan.egress_plan))
|
||||
|
||||
service: dict[str, Any] = {
|
||||
|
||||
@@ -15,12 +15,14 @@ from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from ... import log
|
||||
from ...docker_cmd import run_docker
|
||||
from ...egress import EgressPlan
|
||||
from ...git_gate import GitGatePlan
|
||||
from ...orchestrator.client import OrchestratorClient
|
||||
from ...orchestrator.gateway import GATEWAY_NETWORK
|
||||
from ...orchestrator.lifecycle import INFRA_NAME, OrchestratorService
|
||||
from ...orchestrator.secret_store import ENV_VAR_SECRET_NAME
|
||||
from ..consolidated_util import provision_bottle
|
||||
from ..consolidated_util import teardown_consolidated as _teardown_util
|
||||
from .gateway_provision import DockerGatewayTransport
|
||||
@@ -41,6 +43,7 @@ class LaunchContext:
|
||||
network: str # the shared gateway network to attach to
|
||||
gateway_ip: str # the gateway's address — the agent's proxy target
|
||||
orchestrator_url: str
|
||||
env_var_secret: str = "" # encryption key injected into the agent's env
|
||||
|
||||
|
||||
def _network_cidr(network: str) -> str:
|
||||
@@ -85,6 +88,66 @@ def _network_container_ips(network: str) -> list[str]:
|
||||
return ips
|
||||
|
||||
|
||||
def _reprovision_running_bottles(
|
||||
orchestrator_url: str,
|
||||
network: str = GATEWAY_NETWORK,
|
||||
infra_name: str = INFRA_NAME,
|
||||
) -> None:
|
||||
"""Re-inject egress tokens for any registered bottles that lost their
|
||||
in-memory tokens (e.g., after an infra container restart).
|
||||
|
||||
For each registered bottle whose source IP maps to a live container on the
|
||||
gateway network, reads ENV_VAR_SECRET via ``docker exec … printenv`` and
|
||||
calls ``POST /bottles/<id>/reprovision_gateway``. Idempotent — a no-op
|
||||
when the orchestrator already has all tokens loaded. Best-effort: a single
|
||||
container exec failure never blocks a new bottle launch."""
|
||||
client = OrchestratorClient(orchestrator_url)
|
||||
bottles = client.list_bottles()
|
||||
if not bottles:
|
||||
return
|
||||
|
||||
# Build {source_ip: container_name} from live containers on the gateway
|
||||
# network, excluding the infra container itself.
|
||||
proc = run_docker([
|
||||
"docker", "network", "inspect",
|
||||
"--format", "{{range .Containers}}{{.Name}} {{.IPv4Address}}\n{{end}}",
|
||||
network,
|
||||
])
|
||||
ip_to_container: dict[str, str] = {}
|
||||
for line in proc.stdout.splitlines():
|
||||
parts = line.strip().split()
|
||||
if len(parts) >= 2 and parts[0] != infra_name:
|
||||
ip = parts[1].split("/", 1)[0]
|
||||
if ip:
|
||||
ip_to_container[ip] = parts[0]
|
||||
|
||||
reprovisioned = 0
|
||||
for bottle in bottles:
|
||||
bottle_id = bottle.get("bottle_id")
|
||||
source_ip = bottle.get("source_ip")
|
||||
if not isinstance(bottle_id, str) or not isinstance(source_ip, str):
|
||||
continue
|
||||
container_name = ip_to_container.get(source_ip)
|
||||
if not container_name:
|
||||
continue
|
||||
proc = run_docker(
|
||||
["docker", "exec", container_name, "printenv", ENV_VAR_SECRET_NAME]
|
||||
)
|
||||
if proc.returncode != 0 or not proc.stdout.strip():
|
||||
continue
|
||||
try:
|
||||
if client.reprovision_gateway(bottle_id, proc.stdout.strip()):
|
||||
reprovisioned += 1
|
||||
except Exception: # noqa: BLE001 — best-effort, never block a launch
|
||||
pass
|
||||
|
||||
if reprovisioned:
|
||||
log.info(
|
||||
"reprovisioned egress tokens",
|
||||
context={"count": reprovisioned},
|
||||
)
|
||||
|
||||
|
||||
def launch_consolidated(
|
||||
egress_plan: EgressPlan,
|
||||
git_gate_plan: GitGatePlan,
|
||||
@@ -96,9 +159,14 @@ def launch_consolidated(
|
||||
network: str = GATEWAY_NETWORK,
|
||||
) -> LaunchContext:
|
||||
"""Ensure the infra container is up, allocate + register the bottle, and
|
||||
provision its git-gate state. Returns the agent's attach context."""
|
||||
provision its git-gate state. Returns the agent's attach context.
|
||||
|
||||
Also reprovisiones egress tokens for any already-running bottles that lost
|
||||
their in-memory credentials (e.g. after an infra container restart), so
|
||||
they regain egress access before the new bottle is registered."""
|
||||
service = service or OrchestratorService()
|
||||
url = service.ensure_running()
|
||||
_reprovision_running_bottles(url, network=network, infra_name=infra_name)
|
||||
client = OrchestratorClient(url)
|
||||
|
||||
cidr = _network_cidr(network)
|
||||
@@ -117,6 +185,7 @@ def launch_consolidated(
|
||||
network=network,
|
||||
gateway_ip=gateway_ip,
|
||||
orchestrator_url=url,
|
||||
env_var_secret=reg.env_var_secret,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -186,6 +186,7 @@ def launch(
|
||||
agent_git_gate_url=git_gate_url,
|
||||
agent_supervise_url=supervise_url,
|
||||
identity_token=ctx.identity_token,
|
||||
env_var_secret=ctx.env_var_secret,
|
||||
)
|
||||
|
||||
# Step 5: render + up the agent-only compose, pinned on the shared
|
||||
@@ -198,7 +199,12 @@ def launch(
|
||||
project = compose_project_name(plan.slug)
|
||||
# Forwarded vars (OAuth token, host interpolations) flow through the
|
||||
# subprocess env as bare names so values never land in the file.
|
||||
# ENV_VAR_SECRET follows the same pattern: bare name in the compose
|
||||
# spec, value only in the subprocess env so it is never written to disk.
|
||||
compose_env: dict[str, str] = {**os.environ, **plan.forwarded_env}
|
||||
if plan.env_var_secret:
|
||||
from ...orchestrator.secret_store import ENV_VAR_SECRET_NAME
|
||||
compose_env[ENV_VAR_SECRET_NAME] = plan.env_var_secret
|
||||
info(
|
||||
f"docker compose up -d (project {project}, agent on shared "
|
||||
f"gateway {ctx.gateway_ip}, ip {ctx.source_ip})"
|
||||
|
||||
@@ -50,6 +50,7 @@ class LaunchContext:
|
||||
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 launch_consolidated(
|
||||
@@ -80,6 +81,7 @@ def launch_consolidated(
|
||||
source_ip=guest_ip,
|
||||
gateway_ca_pem=infra.gateway_ca_pem(),
|
||||
orchestrator_url=url,
|
||||
env_var_secret=reg.env_var_secret,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -72,6 +72,7 @@ class LaunchContext:
|
||||
gateway_ip: str
|
||||
network: str
|
||||
orchestrator_url: str
|
||||
env_var_secret: str = "" # encryption key injected into the agent's env
|
||||
|
||||
|
||||
def ensure_gateway(
|
||||
@@ -152,6 +153,7 @@ def register_agent(
|
||||
gateway_ip=endpoint.gateway_ip,
|
||||
network=endpoint.network,
|
||||
orchestrator_url=endpoint.orchestrator_url,
|
||||
env_var_secret=reg.env_var_secret,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -41,10 +41,13 @@ class OrchestratorClientError(RuntimeError):
|
||||
@dataclass(frozen=True)
|
||||
class RegisteredBottle:
|
||||
"""What `POST /bottles` returns: the minted bottle id and the per-bottle
|
||||
identity token the agent presents for app-layer attribution."""
|
||||
identity token the agent presents for app-layer attribution. `env_var_secret`
|
||||
is set by the caller (not from the server response) and carries the
|
||||
encryption key so it can be injected into the agent container's env."""
|
||||
|
||||
bottle_id: str
|
||||
identity_token: str
|
||||
env_var_secret: str = ""
|
||||
|
||||
|
||||
class OrchestratorClient:
|
||||
@@ -120,17 +123,21 @@ class OrchestratorClient:
|
||||
metadata: str = "",
|
||||
policy: str = "",
|
||||
tokens: dict[str, str] | None = None,
|
||||
env_var_secret: str = "",
|
||||
) -> RegisteredBottle:
|
||||
"""Register a bottle and broker its launch (`POST /bottles`). `tokens`
|
||||
are the per-bottle egress auth values (env_name -> value) the
|
||||
orchestrator holds in memory for the gateway to inject. Returns the
|
||||
minted id + identity token."""
|
||||
orchestrator holds in memory for the gateway to inject. When
|
||||
*env_var_secret* is provided, the orchestrator also encrypts the token
|
||||
values and stores them in ``bottled_agent_secrets`` for restart
|
||||
recovery. Returns the minted id + identity token."""
|
||||
payload = self._ok("POST", "/bottles", {
|
||||
"source_ip": source_ip,
|
||||
"image_ref": image_ref,
|
||||
"metadata": metadata,
|
||||
"policy": policy,
|
||||
"tokens": tokens or {},
|
||||
"env_var_secret": env_var_secret,
|
||||
})
|
||||
bottle_id = payload.get("bottle_id")
|
||||
token = payload.get("identity_token")
|
||||
@@ -138,6 +145,24 @@ class OrchestratorClient:
|
||||
raise OrchestratorClientError("register: response missing bottle_id/identity_token")
|
||||
return RegisteredBottle(bottle_id=bottle_id, identity_token=token)
|
||||
|
||||
def reprovision_gateway(self, bottle_id: str, env_var_secret: str) -> bool:
|
||||
"""Re-inject a bottle's egress tokens from its ENV_VAR_SECRET
|
||||
(`POST /bottles/<id>/reprovision_gateway`). Returns True when the
|
||||
orchestrator successfully decrypted and restored the tokens, False
|
||||
when it had no stored secrets for this bottle (404)."""
|
||||
status, _ = self._request(
|
||||
"POST",
|
||||
f"/bottles/{bottle_id}/reprovision_gateway",
|
||||
{"env_var_secret": env_var_secret},
|
||||
)
|
||||
if status == 404:
|
||||
return False
|
||||
if not 200 <= status < 300:
|
||||
raise OrchestratorClientError(
|
||||
f"reprovision_gateway {bottle_id}: HTTP {status}"
|
||||
)
|
||||
return True
|
||||
|
||||
def teardown_bottle(self, bottle_id: str) -> bool:
|
||||
"""Tear a bottle down (`DELETE /bottles/<id>`). False if the
|
||||
orchestrator didn't know it (404) — idempotent for cleanup paths."""
|
||||
|
||||
@@ -9,9 +9,13 @@ vsock / unix-socket portability caveats):
|
||||
GET /bottles -> 200 {"bottles": [ <redacted record>, ...]}
|
||||
POST /bottles -> 201 {"bottle_id","identity_token"} (launch)
|
||||
body: {"source_ip", ["image_ref"],
|
||||
["metadata"], ["policy"]}
|
||||
["metadata"], ["policy"],
|
||||
["tokens"], ["env_var_secret"]}
|
||||
PUT /bottles/<bottle_id>/policy -> 200 {"updated": true} | 404 (live reload)
|
||||
body: {"policy"}
|
||||
POST /bottles/<bottle_id>/reprovision_gateway
|
||||
-> 200 {"reprovisioned": true} | 404
|
||||
body: {"env_var_secret"}
|
||||
DELETE /bottles/<bottle_id> -> 200 {"torn_down": true} | 404 (teardown)
|
||||
POST /reconcile -> 200 {"reaped": [bottle_id, ...]}
|
||||
body: {"live_source_ips": [...],
|
||||
@@ -116,12 +120,14 @@ def dispatch( # pylint: disable=too-many-return-statements,too-many-branches
|
||||
tokens = {
|
||||
k: v for k, v in raw_tokens.items() if isinstance(k, str) and isinstance(v, str)
|
||||
} if isinstance(raw_tokens, dict) else {}
|
||||
env_var_secret = data.get("env_var_secret", "")
|
||||
rec = orch.launch_bottle(
|
||||
source_ip,
|
||||
image_ref=image_ref if isinstance(image_ref, str) else "",
|
||||
metadata=metadata if isinstance(metadata, str) else "",
|
||||
policy=policy if isinstance(policy, str) else "",
|
||||
tokens=tokens,
|
||||
env_var_secret=env_var_secret if isinstance(env_var_secret, str) else "",
|
||||
)
|
||||
return 201, {"bottle_id": rec.bottle_id, "identity_token": rec.identity_token}
|
||||
|
||||
@@ -138,6 +144,23 @@ def dispatch( # pylint: disable=too-many-return-statements,too-many-branches
|
||||
return 200, {"updated": True}
|
||||
return 404, {"error": "no such bottle"}
|
||||
|
||||
if (
|
||||
method == "POST"
|
||||
and route.startswith("/bottles/")
|
||||
and route.endswith("/reprovision_gateway")
|
||||
):
|
||||
bottle_id = route[len("/bottles/") : -len("/reprovision_gateway")]
|
||||
try:
|
||||
data = _parse_json_object(body)
|
||||
except ValueError as e:
|
||||
return 400, {"error": f"invalid JSON: {e}"}
|
||||
env_var_secret = data.get("env_var_secret")
|
||||
if not isinstance(env_var_secret, str) or not env_var_secret:
|
||||
return 400, {"error": "env_var_secret (string) is required"}
|
||||
if orch.reprovision_from_secret(bottle_id, env_var_secret):
|
||||
return 200, {"reprovisioned": True}
|
||||
return 404, {"error": "no stored secrets for this bottle"}
|
||||
|
||||
if method == "DELETE" and route.startswith("/bottles/"):
|
||||
bottle_id = route[len("/bottles/"):]
|
||||
if orch.teardown_bottle(bottle_id):
|
||||
|
||||
@@ -113,6 +113,22 @@ _MIGRATIONS = TableMigrations(
|
||||
# egress allowlist / routes / git config selected by source IP. The
|
||||
# multi-tenant gateway resolves it per request via `attribute`.
|
||||
"ALTER TABLE orchestrator_bottles ADD COLUMN policy TEXT NOT NULL DEFAULT ''",
|
||||
# v4 — per-bottle encrypted egress secrets (PRD prd-new-secret-provider).
|
||||
# One row per env-var: key (env-var name) is plaintext for auditing;
|
||||
# value is the encrypted token string. The encryption key (ENV_VAR_SECRET)
|
||||
# lives only in the agent's environment — a row alone cannot recover the
|
||||
# credential.
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS bottled_agent_secrets (
|
||||
bottled_agent_id TEXT NOT NULL,
|
||||
key TEXT NOT NULL,
|
||||
value TEXT NOT NULL,
|
||||
type TEXT NOT NULL DEFAULT 'injected_env_var'
|
||||
)
|
||||
""",
|
||||
# v5 — index for fast per-bottle lookups and bulk DELETE on teardown.
|
||||
"CREATE INDEX IF NOT EXISTS idx_bottled_agent_secrets_id "
|
||||
"ON bottled_agent_secrets (bottled_agent_id, type)",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -326,6 +342,57 @@ class RegistryStore(DbStore):
|
||||
return None
|
||||
return rec
|
||||
|
||||
# --- encrypted egress secret store ------------------------------------
|
||||
|
||||
def store_agent_secrets(
|
||||
self,
|
||||
bottle_id: str,
|
||||
encrypted_values: dict[str, str],
|
||||
secret_type: str = "injected_env_var",
|
||||
) -> None:
|
||||
"""Replace all stored secrets for *bottle_id* with *encrypted_values*
|
||||
(env-var name → encrypted ciphertext). Deletes then re-inserts so a
|
||||
re-registration is always consistent with the current token set."""
|
||||
with self._connection() as conn:
|
||||
conn.execute(
|
||||
"DELETE FROM bottled_agent_secrets "
|
||||
"WHERE bottled_agent_id = ? AND type = ?",
|
||||
(bottle_id, secret_type),
|
||||
)
|
||||
conn.executemany(
|
||||
"INSERT INTO bottled_agent_secrets "
|
||||
"(bottled_agent_id, key, value, type) VALUES (?, ?, ?, ?)",
|
||||
[(bottle_id, k, v, secret_type) for k, v in encrypted_values.items()],
|
||||
)
|
||||
self._chmod()
|
||||
|
||||
def get_agent_secrets(
|
||||
self,
|
||||
bottle_id: str,
|
||||
secret_type: str = "injected_env_var",
|
||||
) -> dict[str, str]:
|
||||
"""Return {env_var_name: encrypted_value} for *bottle_id*, or {} if none."""
|
||||
with self._connection() as conn:
|
||||
rows = conn.execute(
|
||||
"SELECT key, value FROM bottled_agent_secrets "
|
||||
"WHERE bottled_agent_id = ? AND type = ?",
|
||||
(bottle_id, secret_type),
|
||||
).fetchall()
|
||||
return {row[0]: row[1] for row in rows}
|
||||
|
||||
def delete_agent_secrets(
|
||||
self,
|
||||
bottle_id: str,
|
||||
secret_type: str = "injected_env_var",
|
||||
) -> None:
|
||||
"""Remove all stored secrets for *bottle_id* (e.g. on teardown)."""
|
||||
with self._connection() as conn:
|
||||
conn.execute(
|
||||
"DELETE FROM bottled_agent_secrets "
|
||||
"WHERE bottled_agent_id = ? AND type = ?",
|
||||
(bottle_id, secret_type),
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BottleRecord",
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
"""Symmetric encryption for per-bottle egress secrets (PRD prd-new-secret-provider).
|
||||
|
||||
Each agent receives a random ENV_VAR_SECRET at startup — passed as an env var,
|
||||
never logged or persisted. The host uses this key to encrypt each egress auth
|
||||
token value before writing it to the bottled_agent_secrets table; the DB rows
|
||||
(ciphertext, plaintext env-var name) without the key are insufficient to
|
||||
recover the credentials.
|
||||
|
||||
On orchestrator restart the in-memory token map is lost. The host-side
|
||||
reattachment path reads ENV_VAR_SECRET from the running agent container via
|
||||
``docker exec … printenv ENV_VAR_SECRET`` and posts it to
|
||||
``POST /bottles/<id>/reprovision_gateway``; the orchestrator decrypts the
|
||||
stored rows and re-populates ``_tokens``.
|
||||
|
||||
Encryption scheme: HMAC-SHA256 used as a PRF in CTR mode (stdlib-only,
|
||||
no external deps). Each value is encrypted independently. The output blob is
|
||||
``nonce (16 bytes) || ciphertext`` encoded as URL-safe base64 (no padding).
|
||||
|
||||
keystream_block_i = HMAC-SHA256(key, nonce || i.to_bytes(4, "big"))
|
||||
ciphertext_i = plaintext_i XOR keystream_block_i[:len(plaintext_i)]
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import hmac
|
||||
import secrets
|
||||
|
||||
_KEY_BYTES = 32 # 256-bit key from ENV_VAR_SECRET
|
||||
_NONCE_BYTES = 16 # 128-bit random nonce per encrypt call
|
||||
_BLOCK = 32 # HMAC-SHA256 output width == one keystream block
|
||||
|
||||
# Env-var name the agent container receives at startup.
|
||||
ENV_VAR_SECRET_NAME = "ENV_VAR_SECRET"
|
||||
|
||||
|
||||
def new_env_var_secret() -> str:
|
||||
"""Generate a fresh ENV_VAR_SECRET: 32 random bytes as URL-safe base64."""
|
||||
return base64.urlsafe_b64encode(secrets.token_bytes(_KEY_BYTES)).rstrip(b"=").decode()
|
||||
|
||||
|
||||
def _b64dec(s: str) -> bytes:
|
||||
return base64.urlsafe_b64decode(s + "=" * (-len(s) % 4))
|
||||
|
||||
|
||||
def _keystream(key: bytes, nonce: bytes, block_index: int) -> bytes:
|
||||
return hmac.new(
|
||||
key, nonce + block_index.to_bytes(4, "big"), hashlib.sha256
|
||||
).digest()
|
||||
|
||||
|
||||
def encrypt_value(secret_b64: str, plaintext: str) -> str:
|
||||
"""Encrypt a single string value with *secret_b64* (the ENV_VAR_SECRET).
|
||||
|
||||
Returns a URL-safe base64 blob ``nonce || ciphertext`` suitable for
|
||||
the ``bottled_agent_secrets.value`` column."""
|
||||
key = _b64dec(secret_b64)
|
||||
pt = plaintext.encode()
|
||||
nonce = secrets.token_bytes(_NONCE_BYTES)
|
||||
ct = bytearray()
|
||||
for i in range(0, len(pt), _BLOCK):
|
||||
chunk = pt[i : i + _BLOCK]
|
||||
ks = _keystream(key, nonce, i)[: len(chunk)]
|
||||
ct.extend(p ^ k for p, k in zip(chunk, ks))
|
||||
return base64.urlsafe_b64encode(nonce + bytes(ct)).rstrip(b"=").decode()
|
||||
|
||||
|
||||
def decrypt_value(secret_b64: str, blob_b64: str) -> str:
|
||||
"""Decrypt a blob produced by :func:`encrypt_value`.
|
||||
|
||||
Returns the original plaintext string. Raises ``ValueError`` for malformed
|
||||
input or a key mismatch (wrong key produces garbage, not an error, unless
|
||||
the plaintext is non-UTF-8 — treat all such failures as wrong key)."""
|
||||
key = _b64dec(secret_b64)
|
||||
try:
|
||||
blob = _b64dec(blob_b64)
|
||||
except Exception as exc:
|
||||
raise ValueError(f"invalid ciphertext blob: {exc}") from exc
|
||||
if len(blob) < _NONCE_BYTES:
|
||||
raise ValueError("ciphertext blob too short")
|
||||
nonce, ciphertext = blob[:_NONCE_BYTES], blob[_NONCE_BYTES:]
|
||||
pt = bytearray()
|
||||
for i in range(0, len(ciphertext), _BLOCK):
|
||||
chunk = ciphertext[i : i + _BLOCK]
|
||||
ks = _keystream(key, nonce, i)[: len(chunk)]
|
||||
pt.extend(c ^ k for c, k in zip(chunk, ks))
|
||||
try:
|
||||
return bytes(pt).decode()
|
||||
except UnicodeDecodeError as exc:
|
||||
raise ValueError(f"decryption produced non-UTF-8 output (wrong key?): {exc}") from exc
|
||||
|
||||
|
||||
__all__ = ["ENV_VAR_SECRET_NAME", "new_env_var_secret", "encrypt_value", "decrypt_value"]
|
||||
@@ -87,13 +87,22 @@ class Orchestrator:
|
||||
metadata: str = "",
|
||||
policy: str = "",
|
||||
tokens: dict[str, str] | None = None,
|
||||
env_var_secret: str = "",
|
||||
) -> BottleRecord:
|
||||
"""Register a bottle (with its gateway policy + in-memory egress auth
|
||||
tokens) and broker its launch. Rolls the registry entry back if the
|
||||
launch doesn't take, so a failure leaves no orphan."""
|
||||
launch doesn't take, so a failure leaves no orphan.
|
||||
|
||||
When *env_var_secret* is provided alongside *tokens*, the token values
|
||||
are also encrypted and written to ``bottled_agent_secrets`` so they can
|
||||
survive an orchestrator restart (see ``reprovision_from_secret``)."""
|
||||
rec = self.registry.register(source_ip, metadata=metadata, policy=policy)
|
||||
if tokens:
|
||||
self._tokens[rec.bottle_id] = dict(tokens)
|
||||
if env_var_secret:
|
||||
from .secret_store import encrypt_value
|
||||
encrypted = {k: encrypt_value(env_var_secret, v) for k, v in tokens.items()}
|
||||
self.registry.store_agent_secrets(rec.bottle_id, encrypted)
|
||||
req = LaunchRequest(
|
||||
op="launch",
|
||||
bottle_id=rec.bottle_id,
|
||||
@@ -284,6 +293,26 @@ class Orchestrator:
|
||||
))
|
||||
return True, ""
|
||||
|
||||
# --- secret reprovision -----------------------------------------------
|
||||
|
||||
def reprovision_from_secret(self, bottle_id: str, env_var_secret: str) -> bool:
|
||||
"""Re-inject a bottle's egress tokens from its ENV_VAR_SECRET.
|
||||
|
||||
Reads the encrypted rows from ``bottled_agent_secrets``, decrypts each
|
||||
value with *env_var_secret*, and restores ``_tokens[bottle_id]``.
|
||||
Returns True on success, False when no stored secrets exist for this
|
||||
bottle or decryption fails (wrong key / corrupt data)."""
|
||||
from .secret_store import decrypt_value
|
||||
encrypted = self.registry.get_agent_secrets(bottle_id)
|
||||
if not encrypted:
|
||||
return False
|
||||
try:
|
||||
self._tokens[bottle_id] = {k: decrypt_value(env_var_secret, v)
|
||||
for k, v in encrypted.items()}
|
||||
except ValueError:
|
||||
return False
|
||||
return True
|
||||
|
||||
# --- consolidated gateway ----------------------------------------------
|
||||
|
||||
def ensure_gateway(self) -> None:
|
||||
|
||||
Reference in New Issue
Block a user