agent-credential-broker/tests/fake_bao.py

"""Shared fake OpenBao server used by the offline test suites.

ThreadingHTTPServer implementing the subset of the Vault/OpenBao HTTP
API that the ``openbao`` backend exercises:

    POST /v1/auth/approle/login            {role_id, secret_id}
    GET  /v1/auth/token/lookup-self
    POST /v1/auth/token/renew-self
    POST /v1/auth/token/revoke-self
    GET  /v1/<mount>/metadata/<p>/?list=true   (or LIST verb)
    GET  /v1/<mount>/metadata/<p>/<key>
    GET  /v1/<mount>/data/<p>/<key>

Every request ``(method, path)`` is recorded in ``server.calls`` so
tests can assert that unlock/list_metadata NEVER touch ``/data/*`` --
the key property of the design (the ref index is metadata-only).

Secrets live in ``server.secrets`` keyed by full path relative to the
mount, e.g. ``"agent/gitea.admin"`` -> ``{"data": {...},
"custom_metadata": {...}}``. Paths listed in ``server.forbidden``
answer 403 on ``data/*`` reads (token rejection mid-session fixture).
"""

from __future__ import annotations

import json
import threading
import urllib.parse
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

ROLE_ID = "fixture-role-id-DO-NOT-USE"
SECRET_ID = "fixture-secret-id-DO-NOT-USE"
TOKEN_VALUE = "fixture-bao-token-DO-NOT-USE"

PW = "example-password-DO-NOT-USE"
TOKEN = "example-token-DO-NOT-USE"


def fixture_secrets() -> dict:
    """KV v2 secrets under mount `secret`, prefix `agent`."""
    return {
        "agent/gitea.admin": {
            "data": {"username": "gitea-admin", "password": PW,
                     "uris": ["https://git2.paginica.com"],
                     "notes": "bot account"},
            "custom_metadata": {
                "agent.ref": "gitea.admin",
                "agent.adapters": "browser,http,git",
                "agent.allowed_urls": "https://git2.paginica.com",
                "agent.auth_type": "bearer",
                "agent.confirmation": "AUTO",
                "label": "gitea admin",
                "uris": "https://git2.paginica.com"},
        },
        "agent/test.service": {
            "data": {"password": TOKEN,
                     "uris": "https://service.example,https://alt.example"},
            "custom_metadata": {
                "agent.ref": "test.service",
                "agent.adapters": "http,cli-profile",
                "agent.allowed_urls": "https://service.example",
                "agent.auth_type": "bearer",
                "agent.confirmation": "AUTO",
                "agent.cli_profile": "example-echo",
                "label": "api token"},
        },
        # nested path: id is the full path relative to the prefix
        "agent/infra/router": {
            "data": {"username": "root", "password": PW},
            "custom_metadata": {
                "agent.ref": "infra.router",
                "agent.adapters": "ssh",
                "agent.allowed_hosts": "192.0.2.1",
                "agent.confirmation": "AUTO"},
        },
        # no agent.ref -> invisible to the ref index (same as a BW item
        # without the field)
        "agent/not-for-agents": {
            "data": {"username": "x", "password": PW},
            "custom_metadata": {"label": "private"}},
    }


class FakeBaoServer(ThreadingHTTPServer):
    daemon_threads = True

    def __init__(self, addr, handler, *, secrets=None):
        super().__init__(addr, handler)
        self.secrets = dict(secrets if secrets is not None else fixture_secrets())
        self.tokens: set[str] = set()
        self.calls: list[tuple[str, str]] = []
        self.role_id = ROLE_ID
        self.secret_id = SECRET_ID
        # knobs tests can tweak before a request
        self.lookup_ttl = 3600          # ttl reported by lookup-self
        self.lookup_renewable = True
        self.renew_lease = 3600         # lease reported by renew-self
        self.forbidden: set[str] = set()   # data paths answering 403
        self._seq = 0


def start_server(secrets=None) -> tuple[FakeBaoServer, threading.Thread, int]:
    srv = FakeBaoServer(("127.0.0.1", 0), _Handler, secrets=secrets)
    t = threading.Thread(target=srv.serve_forever,
                         kwargs={"poll_interval": 0.05}, daemon=True)
    t.start()
    return srv, t, srv.server_address[1]


def _children(secrets: dict, rel: str) -> list[str]:
    """Direct children of metadata dir ``rel`` (dirs end with '/')."""
    keys = set()
    for path in secrets:
        if rel:
            if not path.startswith(rel + "/"):
                continue
            rest = path[len(rel) + 1:]
        else:
            rest = path
        if not rest:
            continue
        seg, _, tail = rest.partition("/")
        keys.add(seg + "/" if tail else seg)
    return sorted(keys)


class _Handler(BaseHTTPRequestHandler):
    server: FakeBaoServer

    def log_message(self, *a):
        pass

    # -- helpers -------------------------------------------------------

    def _record(self, path):
        self.server.calls.append((self.command, path))

    def _json(self, status, obj):
        body = json.dumps(obj).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _empty(self, status):
        self.send_response(status)
        self.send_header("Content-Length", "0")
        self.end_headers()

    def _body(self) -> dict:
        try:
            n = int(self.headers.get("Content-Length") or 0)
        except ValueError:
            n = 0
        if n <= 0:
            return {}
        try:
            return json.loads(self.rfile.read(n))
        except Exception:
            return {}

    def _authed(self) -> bool:
        tok = self.headers.get("X-Vault-Token")
        if tok and tok in self.server.tokens:
            return True
        self._json(403, {"errors": ["permission denied"]})
        return False

    # -- routing -------------------------------------------------------

    def do_GET(self):
        self._dispatch("GET")

    def do_POST(self):
        self._dispatch("POST")

    def do_LIST(self):
        self._dispatch("LIST")

    def do_PUT(self):
        self._dispatch("PUT")

    def _dispatch(self, method):
        parsed = urllib.parse.urlparse(self.path)
        path = parsed.path
        query = urllib.parse.parse_qs(parsed.query)
        self._record(path)
        parts = [p for p in path.split("/") if p]
        srv = self.server

        # /v1/auth/approle/login
        if parts == ["v1", "auth", "approle", "login"]:
            if method != "POST":
                return self._json(405, {"errors": ["unsupported operation"]})
            creds = self._body()
            if (creds.get("role_id") == srv.role_id
                    and creds.get("secret_id") == srv.secret_id):
                srv._seq += 1
                tok = f"{TOKEN_VALUE}-{srv._seq}"
                srv.tokens.add(tok)
                return self._json(200, {"auth": {
                    "client_token": tok, "renewable": True,
                    "lease_duration": 3600, "policies": ["agent-broker"]}})
            return self._json(400, {"errors": ["invalid secret_id supplied"]})

        # /v1/auth/token/{lookup-self,renew-self,revoke-self}
        if parts[:3] == ["v1", "auth", "token"]:
            if not self._authed():
                return
            tok = self.headers.get("X-Vault-Token")
            op = parts[3] if len(parts) > 3 else ""
            if op == "lookup-self" and method == "GET":
                return self._json(200, {"data": {
                    "id": tok, "ttl": srv.lookup_ttl,
                    "renewable": srv.lookup_renewable,
                    "policies": ["agent-broker"]}})
            if op == "renew-self" and method == "POST":
                return self._json(200, {"auth": {
                    "client_token": tok, "renewable": True,
                    "lease_duration": srv.renew_lease}})
            if op == "revoke-self" and method == "POST":
                srv.tokens.discard(tok)
                return self._empty(204)
            return self._json(404, {"errors": ["unsupported path"]})

        # /v1/<mount>/{metadata,data}/<rel...>
        if len(parts) >= 4 and parts[0] == "v1" and parts[2] in ("metadata", "data"):
            if not self._authed():
                return
            kind = parts[2]
            rel = "/".join(parts[3:])
            if method not in ("GET", "LIST"):
                return self._json(405, {"errors": ["unsupported operation"]})
            if method == "LIST" or "list" in query:
                keys = _children(srv.secrets, rel)
                if not keys and rel not in srv.secrets:
                    return self._json(404, {"errors": []})
                return self._json(200, {"data": {"keys": keys}})
            entry = srv.secrets.get(rel)
            if entry is None:
                return self._json(404, {"errors": []})
            cm = entry.get("custom_metadata") or {}
            if kind == "metadata":
                return self._json(200, {"data": {
                    "custom_metadata": cm, "version": 1,
                    "created_time": "2026-01-01T00:00:00Z"}})
            if rel in srv.forbidden:
                return self._json(403, {"errors": ["permission denied"]})
            return self._json(200, {"data": {
                "data": entry.get("data") or {},
                "metadata": {"custom_metadata": cm, "version": 1,
                             "created_time": "2026-01-01T00:00:00Z"}}})

        return self._json(404, {"errors": ["unsupported path"]})