key_agent.py

#!/usr/bin/env python3
"""Diskless API-key injecting agent.

Sits between nginx and the upstream backend. nginx proxies to this agent,
which attaches the `Authorization` header from a key held ONLY in this
process's memory, then forwards the request upstream.

Why this removes the key from disk entirely:

  * The key is resolved at startup (Vault / AWS / env) into a Python
    variable. It is never written to a file, never placed in systemd's
    EnvironmentFile, and never rendered into any config.
  * nginx.conf contains no secret and no reference to the key. If the agent
    is compromised in memory, a restart re-fetches; there is no at-rest copy.
  * Optionally, the agent can re-fetch on a timer and on auth failures, so
    the credential can stay short-lived.

Trade-off vs the EnvironmentFile design: this adds a local hop and a second
process to supervise. It buys you "no at-rest key".
"""
from __future__ import annotations

import json
import os
import subprocess
import sys
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from urllib import error, request

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
LISTEN_HOST = os.environ.get("AGENT_HOST", "127.0.0.1")
LISTEN_PORT = int(os.environ.get("AGENT_PORT", "5001"))
UPSTREAM_URL = os.environ.get("UPSTREAM_URL", "http://127.0.0.1:5000").rstrip("/")

# Auth header style: "bearer" or "header" (x-api-key).
AUTH_STYLE = os.environ.get("AUTH_STYLE", "bearer").lower()
AUTH_HEADER = os.environ.get("AUTH_HEADER", "Authorization")
AUTH_SCHEME = os.environ.get("AUTH_SCHEME", "Bearer")

# Rotation: re-fetch every N seconds (0 disables).
REFRESH_SECONDS = int(os.environ.get("KEY_REFRESH_SECONDS", "0"))

DEFAULT_API_KEY = "__REPLACE_ME__"

# ---------------------------------------------------------------------------
# Key resolution. Returns an in-memory string. NEVER writes to disk.
# ---------------------------------------------------------------------------
def _run(cmd: list[str]) -> str | None:
    """Run a command, return stdout stripped, or None on failure."""
    try:
        out = subprocess.run(cmd, capture_output=True, text=True, timeout=30, check=False)
        if out.returncode == 0 and out.stdout.strip():
            return out.stdout.strip()
    except (subprocess.SubprocessError, OSError):
        pass
    return None


def fetch_from_vault() -> str | None:
    path = os.environ.get("VAULT_SECRET_PATH")
    if not path:
        return None
    addr = os.environ.get("VAULT_ADDR")
    if addr:
        os.environ["VAULT_ADDR"] = addr

    # Token may come from the environment (short-lived, injected by an
    # orchestrator) or from a token file read into memory at startup.
    token = os.environ.get("VAULT_TOKEN")
    if not token:
        token_file = os.environ.get("VAULT_TOKEN_FILE", "/etc/vault-token")
        if os.path.isfile(token_file):
            try:
                with open(token_file, encoding="utf-8") as fh:
                    token = fh.read().strip()
            except OSError:
                token = None
    if not token:
        # AppRole: log in and keep the returned token in memory only.
        role_id = os.environ.get("VAULT_ROLE_ID")
        secret_id_file = os.environ.get("VAULT_SECRET_ID_FILE", "/etc/vault-approle-secret-id")
        if role_id and os.path.isfile(secret_id_file):
            try:
                with open(secret_id_file, encoding="utf-8") as fh:
                    secret_id = fh.read().strip()
            except OSError:
                return None
            raw = _run([
                "vault", "write", "-format=json", "auth/approle/login",
                f"role_id={role_id}", f"secret_id={secret_id}",
            ])
            if raw:
                try:
                    token = json.loads(raw).get("auth", {}).get("client_token")
                except json.JSONDecodeError:
                    token = None
    if not token:
        return None

    os.environ["VAULT_TOKEN"] = token
    field = os.environ.get("VAULT_SECRET_FIELD", "api_key")
    return _run(["vault", "kv", "get", f"-field={field}", path])


def fetch_from_aws() -> str | None:
    secret_id = os.environ.get("AWS_SECRET_ID")
    if not secret_id:
        return None
    raw = _run([
        "aws", "secretsmanager", "get-secret-value",
        "--secret-id", secret_id,
        "--query", "SecretString", "--output", "text",
    ])
    if not raw:
        return None
    field = os.environ.get("AWS_SECRET_FIELD", "api_key")
    try:
        data = json.loads(raw)
        if isinstance(data, dict):
            return data.get(field) or data.get("api_key") or data.get("API_KEY")
        return raw
    except json.JSONDecodeError:
        # SecretString was a bare value, not JSON.
        return raw


def resolve_key() -> str:
    """Resolve the key into memory. Precedence: env -> Vault -> AWS -> placeholder."""
    key = os.environ.get("UPSTREAM_API_KEY")
    if key:
        print("[agent] key from environment", flush=True)
        return key

    key = fetch_from_vault()
    if key:
        print("[agent] key fetched from Vault (in memory only)", flush=True)
        return key

    key = fetch_from_aws()
    if key:
        print("[agent] key fetched from AWS Secrets Manager (in memory only)", flush=True)
        return key

    print("[agent][warn] no key resolved; using placeholder", flush=True)
    return DEFAULT_API_KEY


class KeyStore:
    """Thread-safe in-memory holder; supports periodic refresh."""

    def __init__(self, initial: str) -> None:
        self._key = initial
        self._lock = threading.Lock()

    def get(self) -> str:
        with self._lock:
            return self._key

    def refresh(self) -> None:
        new = resolve_key()
        with self._lock:
            if new != self._key:
                print("[agent] key rotated in memory", flush=True)
            self._key = new


KEYSTORE = KeyStore(DEFAULT_API_KEY)


def auth_value(key: str) -> tuple[str, str]:
    """Return (header_name, header_value) for the configured style."""
    if AUTH_STYLE == "header":
        return AUTH_HEADER, key
    return AUTH_HEADER, f"{AUTH_SCHEME} {key}".strip()


# ---------------------------------------------------------------------------
# Forwarding handler
# ---------------------------------------------------------------------------
HOP_BY_HOP = {
    "connection", "keep-alive", "proxy-authenticate", "proxy-authorization",
    "te", "trailers", "transfer-encoding", "upgrade", "host", "content-length",
}


class AgentHandler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"
    server_version = "KeyAgent/1.0"

    def _proxy(self) -> None:
        length = int(self.headers.get("Content-Length") or 0)
        body = self.rfile.read(length) if length else None

        headers = {
            k: v for k, v in self.headers.items()
            if k.lower() not in HOP_BY_HOP
        }

        # The whole point: overwrite any client-supplied credential.
        name, value = auth_value(KEYSTORE.get())
        headers[name] = value

        target = f"{UPSTREAM_URL}{self.path}"
        req = request.Request(target, data=body, headers=headers, method=self.command)

        try:
            with request.urlopen(req, timeout=60) as resp:
                payload = resp.read()
                self.send_response(resp.status)
                for k, v in resp.headers.items():
                    if k.lower() not in HOP_BY_HOP:
                        self.send_header(k, v)
                self.send_header("Content-Length", str(len(payload)))
                self.end_headers()
                self.wfile.write(payload)
        except error.HTTPError as exc:
            payload = exc.read()
            self.send_response(exc.code)
            self.send_header("Content-Type", exc.headers.get("Content-Type", "text/plain"))
            self.send_header("Content-Length", str(len(payload)))
            self.end_headers()
            self.wfile.write(payload)
        except (error.URLError, OSError) as exc:
            msg = f"upstream unreachable: {exc}".encode()
            self.send_response(502)
            self.send_header("Content-Type", "text/plain")
            self.send_header("Content-Length", str(len(msg)))
            self.end_headers()
            self.wfile.write(msg)

    do_GET = do_POST = do_PUT = do_PATCH = do_DELETE = _proxy

    def log_message(self, fmt, *args):
        sys.stderr.write("[agent] %s\n" % (fmt % args))


def main() -> None:
    KEYSTORE.refresh()

    if REFRESH_SECONDS > 0:
        def loop():
            while True:
                time.sleep(REFRESH_SECONDS)
                try:
                    KEYSTORE.refresh()
                except Exception as exc:  # noqa: BLE001 - keep the agent alive
                    print(f"[agent][warn] refresh failed: {exc}", flush=True)

        threading.Thread(target=loop, daemon=True).start()
        print(f"[agent] rotation every {REFRESH_SECONDS}s", flush=True)

    server = ThreadingHTTPServer((LISTEN_HOST, LISTEN_PORT), AgentHandler)
    print(f"[agent] listening on {LISTEN_HOST}:{LISTEN_PORT} -> {UPSTREAM_URL}", flush=True)
    server.serve_forever()


if __name__ == "__main__":
    main()