monky-deployd/tests/fakes.py

"""Fake monky-tenancy agent endpoint + fake OpenBao, as real HTTP servers on 127.0.0.1.

They implement exactly the shapes of record (ADR-0028 + Gate 1 v2): the lease is a
`{login_jwt, ttl_s, mount, role, vault}` deploy grant; OpenBao's `auth/jwt-tenancy/login`
accepts any grant listed in `bao.grants` and pins the token's meta to that grant's env_id."""

from __future__ import annotations

import io
import json
import tarfile
import threading
import time
import uuid
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

from monky_deployd.bundle import bundle_sha

ENV = "env-qa-02"

COMPOSE = """services:
  see:
    image: harbor.tikali.net/monky/see-backend:${SEE_TAG:-develop}
    environment:
      POSTGRES_PASSWORD: ${POSTGRES_PASSWORD}
      GEMINI_API_KEY: ${GEMINI_API_KEY}
      SEE_ADMIN_TOKEN: ${SEE_ADMIN_TOKEN}
    healthcheck:
      test: ["CMD", "true"]
"""


def make_manifest(env=ENV, versions=None, prefix="monky"):
    versions = versions or {}
    entries = [
        {
            "var": "GEMINI_API_KEY",
            "path": f"{prefix}/{env}/see/gemini_api_key",
            "kind": "supplied",
            "version": versions.get("gemini_api_key"),
        },
        {
            "var": "POSTGRES_PASSWORD",
            "path": f"{prefix}/{env}/see/pg_password",
            "kind": "generated",
            "version": versions.get("pg_password", 1),
        },
        {
            "var": "SEE_ADMIN_TOKEN",
            "path": f"{prefix}/{env}/see/admin_token",
            "kind": "generated",
            "version": versions.get("admin_token", 1),
        },
    ]
    return {"vault_mode": "openbao", "entries": entries}


def env_template(manifest):
    lines = ["# .env template — PLACEHOLDERS ONLY", ""]
    for e in manifest["entries"]:
        lines += [f"# [{e['kind']}] openbao: {e['path']}", f"{e['var']}=${{{e['var']}}}"]
    return "\n".join(lines) + "\n"


def make_files(env=ENV, *, compose=COMPOSE, manifest=None, meta=None):
    manifest = manifest or make_manifest(env)
    files = {
        "docker-compose.yml": compose,
        ".env.template": env_template(manifest),
        "secrets.manifest.json": json.dumps(manifest, indent=2, sort_keys=True) + "\n",
    }
    bundle_json = {
        "env_id": env,
        "tier": env.split("-")[1] if env.startswith("env-") else "dev",
        "site": "cbs",
        "bundle": "see",
        "target_class": "docker-host",
        "renderer": "compose",
        "images_policy": "develop",
        "secrets_provider": "openbao",
        "agent": {},
        "files": sorted([*files, "bundle.json"]),
    }
    if meta:
        bundle_json.update(meta)
    files["bundle.json"] = json.dumps(bundle_json, indent=2, sort_keys=True) + "\n"
    return files


def tar_bytes(files: dict[str, str]) -> bytes:
    buf = io.BytesIO()
    with tarfile.open(fileobj=buf, mode="w") as tar:
        for name in sorted(files):
            data = files[name].encode()
            info = tarfile.TarInfo(name=name)
            info.size = len(data)
            info.mode = 0o644
            tar.addfile(info, io.BytesIO(data))
    return buf.getvalue()


class _Server(ThreadingHTTPServer):
    daemon_threads = True
    allow_reuse_address = True


class _Handler(BaseHTTPRequestHandler):
    server_version = "fake/0"
    protocol_version = "HTTP/1.1"

    def log_message(self, *a):  # quiet
        pass

    def _body(self):
        n = int(self.headers.get("Content-Length") or 0)
        raw = self.rfile.read(n) if n else b""
        return json.loads(raw) if raw else {}

    def _send(self, status, body=None, *, raw=None, headers=None, ctype="application/json"):
        data = raw if raw is not None else (json.dumps(body).encode() if body is not None else b"")
        self.send_response(status)
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Length", str(len(data)))
        for k, v in (headers or {}).items():
            self.send_header(k, v)
        self.end_headers()
        self.wfile.write(data)

    def do_GET(self):
        self.server.app.handle(self, "GET", self.path, {})

    def do_POST(self):
        self.server.app.handle(self, "POST", self.path, self._body())


class FakeBao:
    """auth/jwt-tenancy/login, token lookup-self/renew-self/revoke-self, KV-v2 reads."""

    def __init__(self, env=ENV):
        self.env = env
        self.grants: dict[str, str] = {}  # jwt -> env_id
        self.tokens: dict[str, dict] = {}
        self.kv: dict[str, list[dict]] = {}  # "<env>/see/<name>" -> versions [ {value} ]
        self.calls: list[tuple] = []
        self.login_fail = False
        self.token_ttl = 86400
        self.max_ttl = 30 * 86400
        self.mount = "jwt-tenancy"
        self.role = "see-env"
        self.kv_mount = "monky"
        self.seed_secrets(env)

    def seed_secrets(self, env):
        self.kv[f"{env}/see/gemini_api_key"] = [
            {"value": "AIzaSy-FAKE-GEMINI-KEY-0001"},
            {"value": "AIzaSy-FAKE-GEMINI-KEY-0002"},
        ]
        self.kv[f"{env}/see/pg_password"] = [{"value": 'pg-s3cret-"quoted"$dollar'}]
        self.kv[f"{env}/see/admin_token"] = [{"value": "adm1n-t0ken-abcdef0123456789"}]

    def grant(self, env=None, jti=None) -> str:
        jwt = "eyJhbGciOiJFUzI1NiJ9." + uuid.uuid4().hex + uuid.uuid4().hex + "." + uuid.uuid4().hex + uuid.uuid4().hex
        self.grants[jwt] = json.dumps({"env_id": env or self.env, "jti": jti or uuid.uuid4().hex})
        return jwt

    def mint(self, env=None, *, age_s=0, ttl=None) -> str:
        tok = "hvs." + uuid.uuid4().hex + uuid.uuid4().hex[:8]
        self.tokens[tok] = {
            "accessor": "acc-" + uuid.uuid4().hex[:12],
            "env_id": env or self.env,
            "grant_jti": uuid.uuid4().hex,
            "creation_time": int(time.time()) - age_s,
            "ttl": ttl if ttl is not None else self.token_ttl,
            "renewable": True,
            "revoked": False,
        }
        return tok

    def _auth(self, h):
        tok = h.headers.get("X-Vault-Token")
        t = self.tokens.get(tok)
        if not t or t["revoked"]:
            return None, tok
        return t, tok

    def handle(self, h, method, path, body):
        self.calls.append((method, path.split("?")[0]))
        if method == "POST" and path == f"/v1/auth/{self.mount}/login":
            if self.login_fail:
                return h._send(400, {"errors": ["error validating token: expired"]})
            if body.get("role") != self.role:
                return h._send(400, {"errors": [f"role {body.get('role')!r} could not be found"]})
            # the jwt mount keeps no jti/replay state: a known grant logs in until it expires;
            # a SUPERSEDED grant is refused by tenancy (FakeTenancy.superseded_jtis), not here
            meta = self.grants.get(body.get("jwt", ""))
            if meta is None:
                return h._send(400, {"errors": ["error validating token: unknown grant"]})
            m = json.loads(meta)
            tok = self.mint(m["env_id"])
            self.tokens[tok]["grant_jti"] = m["jti"]
            t = self.tokens[tok]
            return h._send(
                200,
                {
                    "auth": {
                        "client_token": tok,
                        "accessor": t["accessor"],
                        "lease_duration": t["ttl"],
                        "renewable": True,
                        "token_policies": ["see-env"],
                        "metadata": {"env_id": m["env_id"], "grant_jti": m["jti"], "role": self.role},
                    }
                },
            )
        t, tok = self._auth(h)
        if t is None:
            return h._send(403, {"errors": ["permission denied"]})
        if method == "GET" and path == "/v1/auth/token/lookup-self":
            return h._send(
                200,
                {
                    "data": {
                        "accessor": t["accessor"],
                        "creation_time": t["creation_time"],
                        "ttl": t["ttl"],
                        "renewable": t["renewable"],
                        "explicit_max_ttl": self.max_ttl,
                        "meta": {"env_id": t["env_id"], "grant_jti": t["grant_jti"]},
                    }
                },
            )
        if method == "POST" and path == "/v1/auth/token/renew-self":
            age = int(time.time()) - t["creation_time"]
            t["ttl"] = max(0, min(self.token_ttl, self.max_ttl - age))
            return h._send(200, {"auth": {"client_token": tok, "lease_duration": t["ttl"], "renewable": True}})
        if method == "POST" and path == "/v1/auth/token/revoke-self":
            t["revoked"] = True
            return h._send(204)
        if method == "GET" and path.startswith(f"/v1/{self.kv_mount}/data/"):
            p, _, q = path.partition("?")
            key = p[len(f"/v1/{self.kv_mount}/data/") :]
            env = key.split("/")[0]
            if env != t["env_id"]:
                return h._send(403, {"errors": ["1 error occurred:\n\t* permission denied\n\n"]})
            versions = self.kv.get(key)
            if not versions:
                return h._send(404, {"errors": []})
            want = None
            for part in q.split("&"):
                if part.startswith("version="):
                    want = int(part.split("=", 1)[1])
            idx = (want or len(versions)) - 1
            if idx < 0 or idx >= len(versions):
                return h._send(404, {"errors": []})
            return h._send(200, {"data": {"data": versions[idx], "metadata": {"version": idx + 1}}})
        return h._send(404, {"errors": [f"no handler for {method} {path}"]})


class FakeTenancy:
    """The agent entrypoint. `bao` validates bearers (a token for another env -> AGENT_ENV_MISMATCH)."""

    def __init__(self, bao: FakeBao, env=ENV):
        self.bao = bao
        self.env = env
        self.files = make_files(env)
        self.desired_sha = bundle_sha(self.files)
        self.action = "apply"
        self.purge_volumes = False
        self.checkins: list[dict] = []
        self.reports: list[dict] = []
        self.leases: list[dict] = []
        self.lease_shape = "jwt"  # or "approle" to simulate the un-migrated tenancy
        self.lease_limit = 5
        self.superseded_jtis: set[str] = set()
        self.down = False  # network down -> connection refused simulated by test stopping server
        self.kv_mount = "monky"

    def set_files(self, files):
        self.files = files
        self.desired_sha = bundle_sha(files)

    def _principal(self, h):
        auth = h.headers.get("Authorization", "")
        tok = auth.split(" ", 1)[1] if auth.lower().startswith("bearer ") else None
        t = self.bao.tokens.get(tok or "")
        if not t or t["revoked"]:
            return None, ("AGENT_UNAUTHENTICATED", "token unknown, expired or revoked")
        if t["grant_jti"] in self.superseded_jtis:
            return None, ("AGENT_UNAUTHENTICATED", "grant superseded")
        return t, None

    def _err(self, h, status, code, detail, headers=None):
        return h._send(status, {"code": code, "detail": detail}, headers=headers)

    def handle(self, h, method, path, body):
        t, err = self._principal(h)
        if err:
            return self._err(h, 401, *err)
        env_in = body.get("env_id") if body else None
        if path.startswith("/v1/agent/bundle/"):
            env_in = path.split("/")[4]
        if env_in and env_in != t["env_id"]:
            return self._err(h, 403, "AGENT_ENV_MISMATCH", f"token is pinned to {t['env_id']}, not {env_in}")
        if method == "POST" and path == "/v1/agent/checkin":
            self.checkins.append(body)
            action = self.action
            if action == "apply" and body.get("applied_sha") == self.desired_sha:
                action = "none"
            return h._send(
                200,
                {
                    "env_id": self.env,
                    "desired_sha": None if self.action == "down" else self.desired_sha,
                    "action": action,
                    "purge_volumes": self.purge_volumes if action == "down" else False,
                    "bundle_url": f"/v1/agent/bundle/{self.env}/{self.desired_sha}",
                    "checkin_interval_s": 60,
                    "vault": {
                        "addr": "https://bao.cbs.tikali.net:8200",
                        "mount": self.kv_mount,
                        "prefix": f"{self.env}/see",
                    },
                },
            )
        if method == "GET" and path.startswith("/v1/agent/bundle/"):
            sha = path.rsplit("/", 1)[1]
            if sha != self.desired_sha:
                return self._err(h, 404, "BUNDLE_NOT_FOUND", f"no published bundle {sha}")
            return h._send(
                200,
                raw=tar_bytes(self.files),
                ctype="application/x-tar",
                headers={"Cache-Control": "no-store", "X-Bundle-Sha": sha},
            )
        if method == "POST" and path == "/v1/agent/lease":
            if len(self.leases) >= self.lease_limit:
                return self._err(h, 429, "LEASE_RATE_LIMITED", "5 leases per hour", headers={"Retry-After": "600"})
            self.leases.append(body)
            if self.lease_shape == "approle":
                return h._send(
                    200,
                    {
                        "env_id": self.env,
                        "wrapping_token": "hvs.wrap-legacy",
                        "wrap_ttl_s": 3600,
                        "role_id": "role-id-see-env",
                        "vault": {},
                    },
                )
            jwt = self.bao.grant(self.env)
            return h._send(
                200,
                {
                    "env_id": self.env,
                    "login_jwt": jwt,
                    "ttl_s": 3600,
                    "mount": self.bao.mount,
                    "role": self.bao.role,
                    # tenancy AgentLeaseOut: `addr` is top-level; KV mount/prefix come from checkin
                    "addr": "https://bao.cbs.tikali.net:8200",
                },
                headers={"Cache-Control": "no-store"},
            )
        if method == "POST" and path == "/v1/agent/report":
            self.reports.append(body)
            if body.get("result") not in ("applied", "failed", "down"):
                return self._err(h, 400, "BAD_REQUEST", "unknown result")
            return h._send(200, {"env_id": self.env, "accepted": True, "result": body["result"]})
        return self._err(h, 404, "NOT_FOUND", path)


def serve(app):
    srv = _Server(("127.0.0.1", 0), _Handler)
    srv.app = app
    th = threading.Thread(target=srv.serve_forever, daemon=True)
    th.start()
    return srv, srv.server_address[1]