atlas-iac/services/hermes/scripts/sync_runtime_credentials.py

#!/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())