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

#!/usr/bin/env python3
"""Reconcile coding-client OAuth documents between runtime storage and Vault."""

from __future__ import annotations

import json
import os
import time
import uuid
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 _credential_value(value: object, required: tuple[str, ...]) -> str | None:
    """Return a complete Vault credential document or reject invalid content."""
    if not isinstance(value, str) or not value.strip():
        return None
    try:
        document = json.loads(value)
    except 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.strip()


def _write_credential(path: Path, value: str) -> None:
    """Atomically restore one private runtime credential from Vault."""
    temporary = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp")
    flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0)
    descriptor = os.open(temporary, flags, 0o600)
    try:
        remaining = memoryview((value + "\n").encode("utf-8"))
        while remaining:
            remaining = remaining[os.write(descriptor, remaining) :]
        os.fsync(descriptor)
        os.fchmod(descriptor, 0o600)
    finally:
        os.close(descriptor)
    try:
        os.replace(temporary, path)
    finally:
        temporary.unlink(missing_ok=True)


def sync_once(token: str) -> tuple[list[str], list[str]]:
    """Persist valid rotations and restore incomplete runtime credentials."""
    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)
    persisted: list[str] = []
    restored: list[str] = []
    for field, (path, required) in CREDENTIALS.items():
        runtime_value = _credential_document(path, required)
        vault_value = _credential_value(data.get(field), required)
        if runtime_value is not None:
            if runtime_value != vault_value:
                updated[field] = runtime_value
                persisted.append(field)
        elif vault_value is not None:
            _write_credential(path, vault_value)
            restored.append(field)
    if persisted:
        _request(
            "POST",
            SECRET_ENDPOINT,
            {"options": {"cas": version}, "data": updated},
            token=token,
        )
    return persisted, restored


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()
            persisted, restored = sync_once(token)
            if persisted:
                print(
                    "Persisted refreshed runtime credentials: "
                    + ", ".join(persisted),
                    flush=True,
                )
            if restored:
                print(
                    "Restored incomplete runtime credentials from Vault: "
                    + ", ".join(restored),
                    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())