Why an AI API key must never go in a web page or app

"""Keep the numbers inside ../widget.html equal to the ones in scenario.py.

    python3 embed_widget.py          rewrite the data block in widget.html
    python3 embed_widget.py --check  exit 1 if widget.html differs (test_widget.py runs the same check)
"""
import json
import re
import sys
from pathlib import Path

from scenario import CAPACITY, DEMO_KEY, PRICE_PER_REQUEST, REFILL_PER_SECOND, STRANGER_REQUESTS, STRANGER_SECONDS

WIDGET = Path(__file__).resolve().parent.parent / "widget.html"
CONFIG_BLOCK = re.compile(r'(<script type="application/json" id="akb-config">)(.*?)(</script>)', re.S)


def config():
    return {"capacity": CAPACITY, "refillPerSecond": REFILL_PER_SECOND, "requests": STRANGER_REQUESTS,
            "seconds": STRANGER_SECONDS, "priceMilliDollars": int(PRICE_PER_REQUEST * 1000), "key": DEMO_KEY}


def embed(html):
    """html with the config block rewritten from scenario.py."""
    block = json.dumps(config(), indent=1)
    return CONFIG_BLOCK.sub(lambda m: m.group(1) + "\n" + block + "\n" + m.group(3), html, count=1)


if __name__ == "__main__":
    current = WIDGET.read_text(encoding="utf-8")
    updated = embed(current)
    if "--check" in sys.argv:
        sys.exit(0 if updated == current else 1)
    WIDGET.write_text(updated, encoding="utf-8")
    print("widget.html data block now matches scenario.py")

"""A small server-side proxy that keeps an AI API key out of the browser. Python standard library only.

The browser sends a chat request to this server. The server checks it, adds the secret key, and
forwards it to the AI provider. The key is read from an environment variable and appears in one place
only: the header of the outgoing request. It is not in any response, error message or log line, and the
server's own access log (see PathFreeHandler) leaves out the request line, so a path that is the key is not logged.

Run it:  AI_API_KEY=... AI_UPSTREAM_URL=... AI_MODELS=demo-small python3 keyproxy.py
"""
import json
import logging
import math
import os
import sys
import time
import urllib.parse
import urllib.request
from dataclasses import dataclass, field
from http import HTTPStatus
from wsgiref.simple_server import WSGIRequestHandler, make_server

log = logging.getLogger("keyproxy")
ROUTE = "/api/chat"
MAX_REPLY_BYTES = 1_000_000


class TokenBucket:
    """Holds up to `capacity` tickets and gains `rate` tickets per second. A request spends one.

    (The word "token" here has nothing to do with the tokens an AI model reads.)
    """

    def __init__(self, capacity, rate, now):
        self.capacity, self.rate = capacity, rate
        self.level, self.last = float(capacity), now

    def take(self, now):
        """Spend one ticket. Return 0.0 if allowed, else the seconds until a ticket is available."""
        self.level = min(self.capacity, self.level + (now - self.last) * self.rate)
        self.last = now
        if self.level >= 1:
            self.level -= 1
            return 0.0
        return (1 - self.level) / self.rate


@dataclass
class Config:
    api_key: str = field(repr=False)  # repr=False: printing the config must not print the key
    upstream_url: str
    allowed_models: frozenset
    max_body_bytes: int = 4096
    max_output_tokens: int = 300
    bucket_capacity: int = 5
    refill_per_second: float = 1.0
    timeout_seconds: float = 10.0


def config_from_env(env=os.environ):
    try:
        models = frozenset(name.strip() for name in env["AI_MODELS"].split(","))
        cfg = Config(env["AI_API_KEY"], env["AI_UPSTREAM_URL"], models)
    except KeyError:
        raise SystemExit("set AI_API_KEY, AI_UPSTREAM_URL and AI_MODELS") from None
    url = urllib.parse.urlsplit(cfg.upstream_url)
    if url.scheme != "https" and url.hostname not in ("127.0.0.1", "localhost"):  # a typo must not send the key in clear text
        raise SystemExit("AI_UPSTREAM_URL must start with https://")
    return cfg


class NoRedirect(urllib.request.HTTPRedirectHandler):
    """Python's urllib would resend the Authorization header to wherever a redirect points."""

    def redirect_request(self, *args, **kwargs):
        return None


def make_app(cfg, clock=time.monotonic):
    """Return a WSGI app. `clock` is a parameter so tests can move time without waiting."""
    buckets = {}  # one TokenBucket per client address
    opener = urllib.request.build_opener(NoRedirect)

    def app(environ, start_response):
        client = environ.get("REMOTE_ADDR", "unknown")  # never X-Forwarded-For: a client can write that

        def reply(status, body, extra=()):
            log.info("client=%s status=%d", client, status)  # no path, no headers, no body
            start_response(f"{status} {HTTPStatus(status).phrase}", [
                ("Content-Type", "application/json"), ("Content-Length", str(len(body))), *extra])
            return [body]

        def fail(status, message, extra=()):
            return reply(status, json.dumps({"error": message}).encode(), extra)  # fixed text, never input

        if environ.get("PATH_INFO") != ROUTE:
            return fail(404, "not found")
        if environ.get("REQUEST_METHOD") != "POST":
            return fail(405, "use POST", [("Allow", "POST")])

        if client not in buckets:
            buckets[client] = TokenBucket(cfg.bucket_capacity, cfg.refill_per_second, clock())
        wait = buckets[client].take(clock())
        if wait:
            return fail(429, "too many requests", [("Retry-After", str(math.ceil(wait)))])

        length = environ.get("CONTENT_LENGTH") or ""
        if not (length.isascii() and length.isdigit()):
            return fail(411, "length required")
        length = length.lstrip("0") or "0"  # leading zeros do not change the number
        if len(length) > 9 or int(length) > cfg.max_body_bytes:  # refuse before reading a byte of the body
            return fail(413, "request too large")
        raw = environ["wsgi.input"].read(int(length))

        try:
            data = json.loads(raw)
            model, messages = data["model"], data["messages"]
        except (ValueError, KeyError, TypeError, RecursionError):
            return fail(400, "send JSON with model and messages")
        if not isinstance(model, str) or not isinstance(messages, list) or not messages:
            return fail(400, "send JSON with model and messages")
        if model not in cfg.allowed_models:
            return fail(400, "model not allowed")

        # The server builds the upstream request: the client cannot add options such as a longer reply.
        upstream_body = json.dumps(
            {"model": model, "messages": messages, "max_tokens": cfg.max_output_tokens}).encode()
        request = urllib.request.Request(cfg.upstream_url, upstream_body, method="POST", headers={
            "Authorization": "Bearer " + cfg.api_key, "Content-Type": "application/json"})
        try:
            with opener.open(request, timeout=cfg.timeout_seconds) as upstream:
                answer = upstream.read(MAX_REPLY_BYTES)
        except Exception as exc:  # an error body can repeat what we sent, so log only a code or a class name
            log.warning("client=%s upstream=%s", client, getattr(exc, "code", type(exc).__name__))
            return fail(502, "upstream failed")
        return reply(200, answer.replace(cfg.api_key.encode(), b"[removed]"))

    return app


class PathFreeHandler(WSGIRequestHandler):
    """The standard handler prints each request line, path included, to stderr. This one prints the address and status."""

    def log_request(self, code="-", size="-"):
        sys.stderr.write(f"client={self.client_address[0]} status={code}\n")

    def log_message(self, format, *args):  # protocol errors quote the request line, so print none of it
        sys.stderr.write(f"client={self.client_address[0]} protocol error\n")


def main():
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")
    with make_server("127.0.0.1", 8080, make_app(config_from_env()), handler_class=PathFreeHandler) as server:
        server.serve_forever()


if __name__ == "__main__":
    main()

"""The pretend scenario in widget.html: a stranger sends many requests. How many does each setup let through?

Every number is an example you can edit. The price per request is the lesson's own arithmetic: 500 input
tokens at $2 per million plus 400 output tokens at $10 per million (the example in "What is a token?").
"""
from decimal import Decimal

from keyproxy import TokenBucket

CAPACITY = 5  # tickets in each bucket
REFILL_PER_SECOND = 1.0
STRANGER_REQUESTS = 1000
STRANGER_SECONDS = 10
PRICE_PER_REQUEST = (500 * Decimal(2) + 400 * Decimal(10)) / 1_000_000  # dollars, pretend


DEMO_KEY = "EXAMPLE-demo-key-do-not-use"  # the made-up key the widget shows; it opens nothing


def stranger_outcomes(limited, addresses=1):
    """One True or False per request the stranger sends: was it accepted?

    The requests are evenly spaced over STRANGER_SECONDS and take the stranger's addresses in turn.
    limited=False is a key in the page, or a server with no limits: every request is accepted.
    """
    buckets = [TokenBucket(CAPACITY, REFILL_PER_SECOND, 0.0) for _ in range(addresses)]
    return [not limited or buckets[i % addresses].take(i * STRANGER_SECONDS / STRANGER_REQUESTS) == 0.0
            for i in range(STRANGER_REQUESTS)]


def stranger_run(limited, addresses=1):
    """How many of the stranger's requests are accepted."""
    return sum(stranger_outcomes(limited, addresses))


def bill(accepted):
    return accepted * PRICE_PER_REQUEST

"""Print every number the lesson computes. Run it from this folder: python3 show_numbers.py"""
from decimal import Decimal

from keyproxy import TokenBucket
from scenario import (CAPACITY, PRICE_PER_REQUEST, REFILL_PER_SECOND, STRANGER_REQUESTS, STRANGER_SECONDS,
                      bill, stranger_run)

print(f"== Pretend price of one chat request: ${PRICE_PER_REQUEST} (500 input tokens at $2/M + 400 output tokens at $10/M)")

print("\n== A bucket of 3 tickets that gains 1 ticket per second, requests at the times shown")
bucket = TokenBucket(3, 1.0, 0.0)
print("  time (s) | tickets before | outcome")
for t in (0.0, 0.1, 0.2, 0.3, 0.8, 1.3, 1.4):
    wait = bucket.take(t)
    before = bucket.level + (0 if wait else 1)
    outcome = "accepted" if not wait else f"refused, try again in {wait:.1f} s"
    print(f"  {t:8.1f} | {before:14.1f} | {outcome}")

print(f"\n== The stranger: {STRANGER_REQUESTS:,} requests in {STRANGER_SECONDS} s "
      f"({STRANGER_REQUESTS // STRANGER_SECONDS}/s), bucket of {CAPACITY} gaining {REFILL_PER_SECOND:g}/s per address")
rows = [("Key in the page", False, 1), ("Server, no limits", False, 1), ("Server, rate limit, 1 address", True, 1)]
rows += [(f"Server, rate limit, {n} addresses", True, n) for n in (20, 100)]
print("  setup | accepted | pretend bill")
for name, limited, addresses in rows:
    accepted = stranger_run(limited, addresses)
    print(f"  {name} | {accepted:,} | ${bill(accepted):.3f}")

last = (STRANGER_REQUESTS - 1) * STRANGER_SECONDS / STRANGER_REQUESTS
print("\n== Why one address gets 14, not 15")
print(f"  {CAPACITY} tickets at the start + {last:g} s x {REFILL_PER_SECOND:g}/s before the last request "
      f"= {CAPACITY + last * REFILL_PER_SECOND:g} tickets, and only whole requests count")

"""Guards the lesson's claim that deleting a committed secret does not remove it from git history."""
import os
import subprocess

KEY = "EXAMPLE-demo-key-0000-not-a-real-key"


def git(repo, *args):
    env = {**os.environ, "GIT_CONFIG_GLOBAL": os.devnull, "GIT_CONFIG_SYSTEM": os.devnull}
    run = subprocess.run(["git", "-C", str(repo), "-c", "user.name=demo", "-c", "user.email=demo@example.invalid",
                          "-c", "commit.gpgsign=false", *args], capture_output=True, text=True, env=env)
    assert run.returncode == 0, run.stderr
    return run.stdout


def test_a_deleted_secret_is_still_in_history_and_gitignore_stops_the_first_commit(tmp_path):
    git(tmp_path, "init", "-q")
    (tmp_path / ".gitignore").write_text(".env\n")
    (tmp_path / ".env").write_text(f"AI_API_KEY={KEY}\n")
    git(tmp_path, "add", "--all")
    git(tmp_path, "commit", "-q", "-m", "first")
    assert ".env" not in git(tmp_path, "ls-files")  # the ignore rule kept it out
    git(tmp_path, "add", "--force", ".env")  # one --force, the mistake
    git(tmp_path, "commit", "-q", "-m", "oops")
    git(tmp_path, "rm", "-q", "--cached", ".env")
    (tmp_path / ".env").unlink()
    git(tmp_path, "commit", "-q", "-m", "remove the file")
    assert not (tmp_path / ".env").exists() and ".env" not in git(tmp_path, "ls-files")
    assert KEY in git(tmp_path, "log", "--all", "-p")  # gone from the folder, still in the history

"""Tests for keyproxy.py. Nothing here touches the internet: the "AI provider" is a fake server on 127.0.0.1."""
import io
import json
import logging
import re
import socket
import textwrap
import threading
import urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from wsgiref.simple_server import WSGIRequestHandler, make_server

import pytest

import keyproxy
from keyproxy import Config, PathFreeHandler, TokenBucket, config_from_env, make_app

KEY = "EXAMPLE-demo-key-0000-not-a-real-key"
GOOD = json.dumps({"model": "demo-small", "messages": [{"role": "user", "content": "hi"}]}).encode()


class QuietServer(ThreadingHTTPServer):
    def handle_error(self, request, client_address):  # the slow-provider test hangs up on purpose
        pass


class FakeProvider:
    """A local stand-in for the AI provider. `mode` picks how it misbehaves; `seen` records what arrived."""

    def __init__(self, mode="ok", redirect_to=None):
        self.mode, self.redirect_to, self.seen = mode, redirect_to, []
        outer = self

        class Handler(BaseHTTPRequestHandler):
            def do_POST(self):
                self.answer(self.rfile.read(int(self.headers["Content-Length"])))

            def do_GET(self):  # urllib turns a redirected POST into a GET
                self.answer(b"null")

            def answer(self, body):
                outer.seen.append({"auth": self.headers.get("Authorization"), "body": json.loads(body)})
                if outer.mode == "slow":
                    threading.Event().wait(0.5)
                echo = self.headers.get("Authorization", "")
                status, payload, headers, message = 200, {"answer": "pretend answer"}, {}, None
                if outer.mode == "echo-key":  # a careless provider that repeats the header it was given
                    payload = {"answer": "pretend answer", "debug": echo}
                if outer.mode == "error-with-key":
                    status, payload = 401, {"error": f"bad credential {echo}"}
                if outer.mode == "key-in-reason":  # the key rides in the status line, which urllib keeps as HTTPError.reason
                    status, message = 401, f"bad credential {echo}"
                if outer.mode == "redirect":
                    status, headers = 302, {"Location": outer.redirect_to}
                data = json.dumps(payload).encode()
                self.send_response(status, message)
                for name, value in headers.items():
                    self.send_header(name, value)
                self.send_header("Content-Length", str(len(data)))
                self.end_headers()
                self.wfile.write(data)

            def log_message(self, *args):
                pass

        self.server = QuietServer(("127.0.0.1", 0), Handler)
        self.url = f"http://127.0.0.1:{self.server.server_port}/v1/chat"
        threading.Thread(target=self.server.serve_forever, args=(0.01,), daemon=True).start()

    def close(self):
        self.server.shutdown()
        self.server.server_close()


@pytest.fixture
def provider():
    made = []

    def make(**kwargs):
        made.append(FakeProvider(**kwargs))
        return made[-1]

    yield make
    for p in made:
        p.close()


class Clock:
    def __init__(self):
        self.now = 1000.0

    def __call__(self):
        return self.now


def cfg_for(provider, **kwargs):
    return Config(KEY, provider.url, frozenset({"demo-small"}), **kwargs)


class NoRead(io.BytesIO):
    def read(self, *args):
        raise AssertionError("the body was read")


def call(app, body=GOOD, method="POST", path="/api/chat", client="203.0.113.7", length="auto", **extra):
    """Run a WSGI app once. Returns (status code, headers dict, body bytes)."""
    environ = {"REQUEST_METHOD": method, "PATH_INFO": path, "REMOTE_ADDR": client,
               "wsgi.input": body if isinstance(body, io.BytesIO) else io.BytesIO(body), **extra}
    if length == "auto":
        length = str(len(body.getvalue() if isinstance(body, io.BytesIO) else body))
    if length is not None:
        environ["CONTENT_LENGTH"] = length
    out = {}

    def start_response(status, headers):
        out["status"], out["headers"] = int(status.split()[0]), dict(headers)

    parts = app(environ, start_response)
    return out["status"], out["headers"], b"".join(parts)


# ---- the token bucket

def test_bucket_allows_a_burst_then_says_how_long_to_wait():
    bucket = TokenBucket(capacity=3, rate=1.0, now=0.0)
    assert [bucket.take(0.0) for _ in range(3)] == [0.0, 0.0, 0.0]
    assert bucket.take(0.0) == 1.0  # empty: one ticket arrives in 1 second
    assert bucket.take(0.5) == pytest.approx(0.5)  # a refused request spends nothing; half a ticket has arrived
    assert bucket.take(1.0) == 0.0


def test_bucket_never_holds_more_than_its_capacity():
    bucket = TokenBucket(capacity=3, rate=1.0, now=0.0)
    for _ in range(3):
        bucket.take(0.0)
    assert [bucket.take(1000.0) for _ in range(4)] == [0.0, 0.0, 0.0, 1.0]  # waiting 1,000 s buys only 3


# ---- what a good request does

def test_good_request_is_forwarded_with_the_key_and_the_server_sets_the_limits(provider):
    p = provider()
    body = json.dumps({"model": "demo-small", "messages": [{"role": "user", "content": "hi"}],
                       "max_tokens": 99999, "stream": True, "tools": [{"name": "x"}]}).encode()
    status, headers, out = call(make_app(cfg_for(p)), body)
    assert (status, json.loads(out)) == (200, {"answer": "pretend answer"})
    assert headers["Content-Type"] == "application/json"
    assert p.seen == [{"auth": f"Bearer {KEY}", "body": {
        "model": "demo-small", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 300}}]


def test_printing_the_config_does_not_print_the_key(provider):
    cfg = cfg_for(provider())
    assert KEY not in repr(cfg) and KEY not in str(cfg)


def test_config_comes_from_environment_variables_and_names_only_the_variables_when_one_is_missing():
    env = {"AI_API_KEY": KEY, "AI_UPSTREAM_URL": "http://127.0.0.1:1/x", "AI_MODELS": "demo-small, demo-large"}
    assert config_from_env(env).allowed_models == {"demo-small", "demo-large"}  # a space after the comma is not part of the name
    with pytest.raises(SystemExit) as exc:
        config_from_env({"AI_UPSTREAM_URL": "x", "AI_MODELS": "m"})
    assert "AI_API_KEY" in str(exc.value) and KEY not in str(exc.value)


@pytest.mark.parametrize("url", ["http://api.example.com/v1/chat", "api.example.com/v1/chat",
                                 "http://localhost.evil.example/v1/chat", "http://127.0.0.1.evil.example/v1/chat"])
def test_an_upstream_address_that_is_not_https_is_refused_unless_it_is_this_computer(url):
    """Guards a typo (http for https) that would send the key in clear text."""
    env = {"AI_API_KEY": KEY, "AI_UPSTREAM_URL": url, "AI_MODELS": "demo-small"}
    with pytest.raises(SystemExit) as exc:
        config_from_env(env)
    assert "https" in str(exc.value) and KEY not in str(exc.value)
    env["AI_UPSTREAM_URL"] = "https://api.example.com/v1/chat"
    assert config_from_env(env).upstream_url == env["AI_UPSTREAM_URL"]


# ---- rate limit

def test_sixth_request_in_a_burst_is_refused_and_never_reaches_the_provider(provider):
    p, clock = provider(), Clock()
    app = make_app(cfg_for(p, bucket_capacity=5, refill_per_second=1.0), clock)
    codes = [call(app)[0] for _ in range(6)]
    assert codes == [200] * 5 + [429]
    status, headers, _ = call(app)
    assert (status, headers["Retry-After"]) == (429, "1")
    assert len(p.seen) == 5
    clock.now += 1.0
    assert call(app)[0] == 200


def test_retry_after_rounds_a_fraction_of_a_second_up(provider):
    """Guards ceil() turning into int(): a wait of 0.4 s must say 1, because 0 tells the client to retry at once."""
    clock = Clock()
    app = make_app(cfg_for(provider(), bucket_capacity=1, refill_per_second=1.0), clock)
    assert call(app)[0] == 200
    clock.now += 0.6
    status, headers, _ = call(app)
    assert (status, headers["Retry-After"]) == (429, "1")


def test_each_client_address_has_its_own_bucket(provider):
    app = make_app(cfg_for(provider(), bucket_capacity=1), Clock())
    assert call(app, client="203.0.113.1")[0] == 200
    assert call(app, client="203.0.113.1")[0] == 429
    assert call(app, client="203.0.113.2")[0] == 200


def test_a_forwarded_for_header_cannot_buy_a_fresh_bucket(provider):
    app = make_app(cfg_for(provider(), bucket_capacity=1), Clock())
    assert call(app, HTTP_X_FORWARDED_FOR="198.51.100.1")[0] == 200
    assert call(app, HTTP_X_FORWARDED_FOR="198.51.100.2")[0] == 429


# ---- request size

def test_body_over_the_cap_is_refused_before_it_is_read(provider):
    p = provider()
    app = make_app(cfg_for(p, max_body_bytes=100))
    status, _, _ = call(app, NoRead(b"x" * 101), length="101")
    assert status == 413 and p.seen == []
    assert call(app, NoRead(b"x"), length="9999999999")[0] == 413
    assert call(app, NoRead(b"x"), length="9" * 5000)[0] == 413  # int() refuses more than 4,300 digits, which was a 500


@pytest.mark.parametrize("length", [None, "", "abc", "-5", "²", "1.5"])
def test_missing_or_odd_content_length_is_refused(provider, length):
    assert call(make_app(cfg_for(provider())), length=length)[0] == 411


def test_body_exactly_at_the_cap_passes(provider):
    app = make_app(cfg_for(provider(), max_body_bytes=len(GOOD)))
    assert call(app)[0] == 200


@pytest.mark.parametrize("body, length, status", [(GOOD, str(len(GOOD)).zfill(12), 200), (b"", "0" * 12, 400)],
                         ids=["padded-real-length", "padded-zero"])
def test_a_content_length_with_leading_zeros_is_the_number_it_spells(provider, body, length, status):
    """Guards the nine-digit bound counting padding (a valid length got 413) and int("") on an all-zero length."""
    assert call(make_app(cfg_for(provider())), body, length=length)[0] == status


# ---- model allowlist and shape

def test_a_model_that_is_not_on_the_list_is_refused_and_not_forwarded(provider):
    p = provider()
    body = json.dumps({"model": "demo-huge-and-costly", "messages": [{"role": "user", "content": "hi"}]}).encode()
    status, _, out = call(make_app(cfg_for(p)), body)
    assert (status, json.loads(out)) == (400, {"error": "model not allowed"}) and p.seen == []


@pytest.mark.parametrize("body", [
    b"", b"not json", b"[]", b"7", b'{"model": "demo-small"}', b'{"messages": [{"role": "user"}]}',
    b'{"model": ["demo-small"], "messages": [{"role": "user"}]}',
    b'{"model": "demo-small", "messages": []}', b'{"model": "demo-small", "messages": "hi"}', b"\xff\xfe"])
def test_malformed_json_is_refused_without_a_crash(provider, body):
    p = provider()
    assert call(make_app(cfg_for(p)), body)[0] == 400 and p.seen == []


def test_a_parser_that_runs_out_of_stack_gives_400_not_a_crash(monkeypatch):
    """Guards Python 3.11 and earlier, where json.loads raises RecursionError on about a thousand open brackets."""
    def too_deep(*args):
        raise RecursionError

    monkeypatch.setattr(json, "loads", too_deep)
    cfg = Config(KEY, "http://127.0.0.1:9/v1/chat", frozenset({"demo-small"}))
    assert call(make_app(cfg), b"[" * 4000)[0] == 400


def test_only_post_on_the_one_route_is_served(provider):
    app = make_app(cfg_for(provider()))
    assert call(app, path="/")[0] == 404 and call(app, path="/api/chat/")[0] == 404
    status, headers, _ = call(app, method="GET")
    assert (status, headers["Allow"]) == (405, "POST")


# ---- upstream trouble

def test_provider_error_becomes_a_plain_502(provider):
    p = provider(mode="error-with-key")
    status, _, out = call(make_app(cfg_for(p)))
    assert (status, json.loads(out)) == (502, {"error": "upstream failed"})


def test_provider_that_never_answers_gives_502_after_the_timeout(provider):
    p = provider(mode="slow")
    assert call(make_app(cfg_for(p, timeout_seconds=0.05)))[0] == 502


def test_nothing_listening_gives_502():
    cfg = Config(KEY, "http://127.0.0.1:9/v1/chat", frozenset({"demo-small"}))
    assert call(make_app(cfg))[0] == 502


def test_a_redirect_is_not_followed_so_the_key_is_not_sent_elsewhere(provider):
    elsewhere = provider()
    redirecting = provider(mode="redirect", redirect_to=elsewhere.url)
    status, _, _ = call(make_app(cfg_for(redirecting)))
    assert status == 502
    assert redirecting.seen and elsewhere.seen == []


# ---- the key never comes back out

def test_a_provider_that_echoes_the_key_is_scrubbed_and_the_test_is_not_vacuous(provider):
    p = provider(mode="echo-key")
    status, _, out = call(make_app(cfg_for(p)))
    assert p.seen[0]["auth"] == f"Bearer {KEY}"  # the provider really did receive it, and repeated it
    assert status == 200 and KEY.encode() not in out and b"[removed]" in out


@pytest.mark.parametrize("mode", ["ok", "echo-key", "error-with-key", "key-in-reason", "slow", "redirect"])
def test_the_key_is_in_no_response_body_header_or_log_line(provider, caplog, mode):
    caplog.set_level(logging.DEBUG)
    elsewhere = provider()
    p = provider(mode=mode, redirect_to=elsewhere.url)
    clock = Clock()
    app = make_app(cfg_for(p, bucket_capacity=1, max_body_bytes=200, timeout_seconds=0.05), clock)
    seen = []
    seen.append(call(app))  # a good request, whatever the provider then does
    seen.append(call(app, client="203.0.113.9"))
    seen.append(call(app))  # runs out of tickets: 429
    seen.append(call(app, b"x" * 300, client="203.0.113.10"))  # 413
    seen.append(call(app, b"not json", client="203.0.113.11"))  # 400
    seen.append(call(app, method="GET", client="203.0.113.12"))  # 405
    seen.append(call(app, path=f"/{KEY}", client="203.0.113.13"))  # a path that contains the key: 404
    seen.append(call(app, length=None, client="203.0.113.14"))  # 411
    assert {s[0] for s in seen} >= {400, 404, 405, 411, 413, 429}
    for status, headers, body in seen:
        assert KEY.encode() not in body, status
        assert all(KEY not in f"{name}: {value}" for name, value in headers.items()), status
    assert caplog.records
    for record in caplog.records:
        text = " ".join([record.getMessage(), str(record.args), str(record.exc_text), record.pathname])
        assert KEY not in text


# ---- through a real server socket

def stderr_after_raw_request(handler_class, request_line, capsys, cfg):
    """Serve one raw request through a real wsgiref server and return what the server printed to stderr."""
    server = make_server("127.0.0.1", 0, make_app(cfg), handler_class=handler_class)
    threading.Thread(target=server.serve_forever, args=(0.01,), daemon=True).start()
    try:
        with socket.create_connection(("127.0.0.1", server.server_port), timeout=5) as conn:
            conn.sendall(f"{request_line}\r\nContent-Length: 0\r\n\r\n".encode())
            while conn.recv(65536):  # the server closes after it has answered and logged
                pass
    finally:
        server.shutdown()
        server.server_close()
    return capsys.readouterr().err


@pytest.mark.parametrize("request_line", [f"POST /{KEY} HTTP/1.1", f"GET /x {KEY}"], ids=["key-as-path", "key-in-a-bad-request-line"])
def test_the_servers_own_access_log_leaves_out_the_request_line(capsys, request_line):
    """Guards PathFreeHandler: the standard handler prints the whole request line, so a key used as the path was logged."""
    cfg = Config(KEY, "http://127.0.0.1:9/v1/chat", frozenset({"demo-small"}))
    assert KEY in stderr_after_raw_request(WSGIRequestHandler, request_line, capsys, cfg)  # the stock handler really prints it
    printed = stderr_after_raw_request(PathFreeHandler, request_line, capsys, cfg)
    assert KEY not in printed and re.search(r"client=127\.0\.0\.1 (status=\d+|protocol error)", printed)


def test_main_serves_with_the_path_free_handler(monkeypatch):
    """Guards the wiring: a main() that leaves handler_class out serves with the standard handler and logs the key."""
    class Stop(Exception):
        pass

    seen = {}

    def fake_make_server(host, port, app, handler_class=None):
        seen["handler_class"] = handler_class
        raise Stop  # no socket is opened

    monkeypatch.setattr(logging, "basicConfig", lambda **kwargs: None)
    monkeypatch.setattr("keyproxy.make_server", fake_make_server)
    monkeypatch.setenv("AI_API_KEY", KEY)
    monkeypatch.setenv("AI_UPSTREAM_URL", "https://api.example.com/v1/chat")
    monkeypatch.setenv("AI_MODELS", "demo-small")
    with pytest.raises(Stop):
        keyproxy.main()
    assert seen["handler_class"] is PathFreeHandler


def test_end_to_end_through_a_real_wsgi_server(provider):
    class Quiet(WSGIRequestHandler):
        def log_message(self, *args):
            pass

    p = provider()
    server = make_server("127.0.0.1", 0, make_app(cfg_for(p)), handler_class=Quiet)
    threading.Thread(target=server.serve_forever, args=(0.01,), daemon=True).start()
    try:
        request = urllib.request.Request(f"http://127.0.0.1:{server.server_port}/api/chat", GOOD,
                                         headers={"Content-Type": "application/json"})
        with urllib.request.urlopen(request, timeout=5) as response:
            body = response.read()
            headers = str(response.headers)
    finally:
        server.shutdown()
        server.server_close()
    assert json.loads(body) == {"answer": "pretend answer"} and KEY not in headers
    assert p.seen[0]["auth"] == f"Bearer {KEY}"


def test_every_python_block_in_the_lesson_is_a_verbatim_excerpt_of_the_source():
    """The lesson quotes this file. If the code changes and the lesson does not, the reader sees code that was never tested."""
    lesson = (Path(__file__).parent.parent / "lesson.md").read_text()
    source = Path(__file__).with_name("keyproxy.py").read_text().splitlines()
    blocks = re.findall(r"^```python\n(.*?)^```", lesson, flags=re.S | re.M)
    assert blocks, "the lesson has no python blocks to check"
    for block in blocks:
        lines = block.rstrip("\n").splitlines()
        windows = (textwrap.dedent("\n".join(source[i:i + len(lines)])) for i in range(len(source) - len(lines) + 1))
        assert textwrap.dedent("\n".join(lines)) in windows, f"not in keyproxy.py:\n{block}"

"""Guards for ../widget.html: its data block, what it may show, and that its bucket counts like keyproxy.py."""
import json
import re
import subprocess

import embed_widget
from embed_widget import WIDGET, embed
from scenario import DEMO_KEY, bill, stranger_outcomes

HTML = WIDGET.read_text(encoding="utf-8")
SNIPPETS = json.loads(re.search(r'<script type="application/json" id="akb-text">(.*?)</script>', HTML, re.S).group(1))

# Reads the bucket and bill code out of the widget between its marker comments, runs it, and compares with Python.
PARITY = r"""
import fs from 'node:fs';
const [fixture, widget] = process.argv.slice(2);
const fx = JSON.parse(fs.readFileSync(fixture, 'utf8'));
const html = fs.readFileSync(widget, 'utf8');
const cfg = JSON.parse(html.match(/<script type="application\/json" id="akb-config">([\s\S]*?)<\/script>/)[1]);
const js = html.slice(html.indexOf('/* ---------- the bucket and the stranger'), html.indexOf('/* ---------- state'));
const api = new Function('cfg', "'use strict';" + js + '; return {outcomes, dollars};')(cfg);
const bad = [];
fx.runs.forEach((r) => {
  if (JSON.stringify(api.outcomes(r.limited, r.addresses)) !== JSON.stringify(r.outcomes)) bad.push(['outcomes', r.limited, r.addresses]);
});
fx.money.forEach(([n, s]) => { if (api.dollars(n) !== s) bad.push(['dollars', n, api.dollars(n), s]); });
console.log(JSON.stringify({ runs: fx.runs.length, money: fx.money.length, bad: bad.slice(0, 5) }));
process.exit(bad.length ? 1 : 0);
"""


def money(accepted):
    text = f"{bill(accepted):.3f}"
    return "$" + (text[:-1] if text.endswith("0") else text)


def test_widget_data_block_matches_scenario():
    """Guards a scenario.py edit that was not copied into the widget (run embed_widget.py)."""
    assert embed(HTML) == HTML


def test_every_element_the_script_looks_up_exists_inside_the_root():
    """Guards a widget that never starts: a lookup of an id that is not in the page throws in the browser."""
    root = HTML[HTML.index('<div class="akb-root">'):HTML.index("<script")]
    lookups = set(re.findall(r"\$\('([\w-]+)'\)", HTML[HTML.index("<script>"):]))
    assert lookups and all(f'id="{i}"' in root for i in lookups), sorted(i for i in lookups if f'id="{i}"' not in root)


def test_widget_is_small_and_asks_for_nothing_online():
    """The shown code samples mention fetch() and a URL, so look at the page without its two data blocks."""
    assert len(HTML.encode("utf-8")) < 40 * 1024
    page = re.sub(r'<script type="application/json" id="akb-(config|text)">.*?</script>', "", HTML, flags=re.S)
    for banned in ("http://", "https://", "fetch(", "XMLHttpRequest", "WebSocket", "sendBeacon", "import("):
        assert banned not in page, banned
    assert " src=" not in page and "@import" not in page
    assert set(re.findall(r"https?://[a-z0-9./-]+", HTML)) == {"https://api.example.com/v1/chat"}


def test_widget_has_no_form_element():
    """Guards a dead Send button: the lesson page loads widget.html in a sandboxed frame, where a form cannot submit."""
    assert "<form" not in HTML and 'type="submit"' not in HTML


def test_the_key_is_shown_only_where_it_is_really_in_the_page():
    """Guards the lesson's central picture: with the key on the server, no panel text may contain it."""
    assert "EXAMPLE" in DEMO_KEY and "demo" in DEMO_KEY
    assert {name for name, text in SNIPPETS.items() if "{KEY}" in text} == {"sourcePage", "networkPage"}
    assert not [text for text in SNIPPETS.values() if DEMO_KEY in text]  # the key itself comes only from the config block


def test_widget_bucket_and_bill_match_the_python_code(tmp_path):
    """Guards the widget's JavaScript drifting from keyproxy.py: its counts must be the lesson's."""
    runs = [{"limited": limited, "addresses": n, "outcomes": stranger_outcomes(limited, n)}
            for limited, n in [(False, 1), (True, 1), (True, 2), (True, 3), (True, 7), (True, 20), (True, 100)]]
    fixture = {"runs": runs, "money": [[n, money(n)] for n in range(1001)]}
    (tmp_path / "fixture.json").write_text(json.dumps(fixture), encoding="utf-8")
    (tmp_path / "parity.mjs").write_text(PARITY, encoding="utf-8")
    run = subprocess.run(["node", str(tmp_path / "parity.mjs"), str(tmp_path / "fixture.json"), str(WIDGET)],
                         capture_output=True, text=True, timeout=60)
    assert run.returncode == 0, run.stdout + run.stderr
    assert json.loads(run.stdout) == {"runs": 7, "money": 1001, "bad": []}
    assert sum(runs[1]["outcomes"]) == 14 and sum(runs[0]["outcomes"]) == 1000  # the numbers the lesson quotes