142 lines
4.5 KiB
Python
142 lines
4.5 KiB
Python
#!/usr/bin/env python3
|
|
"""Persist refreshed coding-client OAuth documents back into Vault."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from urllib.error import HTTPError, URLError
|
|
from urllib.request import Request, urlopen
|
|
|
|
|
|
VAULT_ADDR = os.environ.get(
|
|
"VAULT_ADDR", "http://vault.vault.svc.cluster.local:8200"
|
|
).rstrip("/")
|
|
VAULT_ROLE = os.environ.get(
|
|
"HERMES_CREDENTIAL_SYNC_VAULT_ROLE", "hermes-credential-sync"
|
|
)
|
|
SECRET_ENDPOINT = "/v1/kv/data/atlas/hermes/agent-tokens"
|
|
JWT_PATH = Path("/var/run/secrets/kubernetes.io/serviceaccount/token")
|
|
CREDENTIALS = {
|
|
"claude_credentials_json": (
|
|
Path("/runtime-access/claude/.credentials.json"),
|
|
("claudeAiOauth", "refreshToken"),
|
|
),
|
|
"codex_auth_json": (
|
|
Path("/runtime-access/codex/auth.json"),
|
|
("tokens", "refresh_token"),
|
|
),
|
|
}
|
|
|
|
|
|
def _request(
|
|
method: str,
|
|
path: str,
|
|
payload: dict[str, Any] | None = None,
|
|
*,
|
|
token: str = "",
|
|
) -> dict[str, Any]:
|
|
"""Issue one bounded Vault JSON request."""
|
|
headers = {"Content-Type": "application/json"}
|
|
if token:
|
|
headers["X-Vault-Token"] = token
|
|
request = Request(
|
|
VAULT_ADDR + path,
|
|
data=(json.dumps(payload).encode("utf-8") if payload is not None else None),
|
|
headers=headers,
|
|
method=method,
|
|
)
|
|
with urlopen(request, timeout=10) as response:
|
|
body = response.read(4 << 20)
|
|
value = json.loads(body)
|
|
if not isinstance(value, dict):
|
|
raise RuntimeError("Vault returned a non-object response")
|
|
return value
|
|
|
|
|
|
def _login() -> str:
|
|
"""Exchange the pod identity for the narrow credential-sync policy."""
|
|
response = _request(
|
|
"POST",
|
|
"/v1/auth/kubernetes/login",
|
|
{"role": VAULT_ROLE, "jwt": JWT_PATH.read_text(encoding="utf-8").strip()},
|
|
)
|
|
token = response.get("auth", {}).get("client_token")
|
|
if not isinstance(token, str) or not token:
|
|
raise RuntimeError("Vault login returned no client token")
|
|
return token
|
|
|
|
|
|
def _credential_document(path: Path, required: tuple[str, ...]) -> str | None:
|
|
"""Return one complete credential JSON document or ignore a partial write."""
|
|
try:
|
|
value = path.read_text(encoding="utf-8").strip()
|
|
document = json.loads(value)
|
|
except (OSError, UnicodeError, json.JSONDecodeError):
|
|
return None
|
|
current = document
|
|
for key in required:
|
|
if not isinstance(current, dict) or key not in current:
|
|
return None
|
|
current = current[key]
|
|
if not isinstance(current, str) or not current:
|
|
return None
|
|
return value
|
|
|
|
|
|
def sync_once(token: str) -> list[str]:
|
|
"""CAS-update only credential fields whose runtime documents changed."""
|
|
response = _request("GET", SECRET_ENDPOINT, token=token)
|
|
envelope = response.get("data")
|
|
if not isinstance(envelope, dict):
|
|
raise RuntimeError("Vault KV response has no data envelope")
|
|
data = envelope.get("data")
|
|
metadata = envelope.get("metadata")
|
|
if not isinstance(data, dict) or not isinstance(metadata, dict):
|
|
raise RuntimeError("Vault KV response has an invalid shape")
|
|
version = metadata.get("version")
|
|
if not isinstance(version, int) or version < 1:
|
|
raise RuntimeError("Vault KV response has no version")
|
|
updated = dict(data)
|
|
changed: list[str] = []
|
|
for field, (path, required) in CREDENTIALS.items():
|
|
value = _credential_document(path, required)
|
|
if value is not None and value != data.get(field):
|
|
updated[field] = value
|
|
changed.append(field)
|
|
if changed:
|
|
_request(
|
|
"POST",
|
|
SECRET_ENDPOINT,
|
|
{"options": {"cas": version}, "data": updated},
|
|
token=token,
|
|
)
|
|
return changed
|
|
|
|
|
|
def main() -> int:
|
|
"""Continuously persist refresh-token rotation without exposing values."""
|
|
interval = max(30, int(os.environ.get("HERMES_CREDENTIAL_SYNC_INTERVAL", "60")))
|
|
while True:
|
|
try:
|
|
token = _login()
|
|
changed = sync_once(token)
|
|
if changed:
|
|
print(
|
|
"Persisted refreshed runtime credentials: " + ", ".join(changed),
|
|
flush=True,
|
|
)
|
|
except (OSError, RuntimeError, HTTPError, URLError, ValueError) as error:
|
|
print(
|
|
f"Runtime credential sync deferred: {type(error).__name__}",
|
|
flush=True,
|
|
)
|
|
time.sleep(interval)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|