Topic 5: Module 7 lab

"""Module 7 lab: start the protected notes server and check every authorization rule end to end.

Runs on port 8079 against a throwaway copy of the notes. The signing key is
generated fresh for this run and only ever lives in environment variables.
"""
import collections
import os
import re
import secrets
import shutil
import subprocess
import sys
import tempfile
import time
from pathlib import Path

PORT = 8079
URL = f"http://127.0.0.1:{PORT}/mcp"
os.environ["NOTES_AUTH_KEY"] = secrets.token_urlsafe(32)
os.environ["NOTES_AUTH_ISSUER"] = "https://auth.sleeplab.example/realms/lab"
os.environ["NOTES_RESOURCE_URL"] = URL

import anyio  # noqa: E402
import httpx2  # noqa: E402

from examples.m07_scoped_calls import connect  # noqa: E402
from examples.m07_step_up import parse_challenge  # noqa: E402
from examples.m07_token_matrix import probe  # noqa: E402
from examples.m07_tokens import make_tokens  # noqa: E402

checks: list[tuple[str, bool]] = []


def check(name: str, ok: bool) -> None:
    checks.append((name, ok))


def start_server(workdir: Path) -> tuple[subprocess.Popen, Path]:
    notes = workdir / "notes"
    shutil.copytree("notes", notes)
    log_path = workdir / "server.log"
    env = dict(os.environ, NOTES_DIR=str(notes), NOTES_TRANSPORT="streamable-http", NOTES_PORT=str(PORT))
    proc = subprocess.Popen(
        [sys.executable, "-m", "notes_assistant.server"],
        env=env, stdout=subprocess.DEVNULL, stderr=log_path.open("w"),
    )
    for _ in range(50):
        try:
            httpx2.get(f"http://127.0.0.1:{PORT}/.well-known/oauth-protected-resource/mcp", timeout=0.5)
            return proc, log_path
        except httpx2.TransportError:
            time.sleep(0.1)
    raise RuntimeError("server did not start; read server.log")


async def mcp_checks(tokens: dict[str, str]) -> None:
    async with connect(tokens["read"]) as client:
        found = await client.call_tool("search_notes", {"query": "sleep"})
        check("read token can search", not found.is_error and len(found.structured_content["hits"]) > 0)
        refused = await client.call_tool("create_note", {"title": "Lab read attempt", "body": "no"})
        check("read token cannot create", refused.is_error and "notes:write" in refused.content[0].text)
    async with connect(tokens["read+write"]) as client:
        made = await client.call_tool("create_note", {"title": "Lab write attempt", "body": "yes", "tags": ["lab"]})
        check("write token can create", not made.is_error)


def main() -> None:
    workdir = Path(tempfile.mkdtemp(prefix="m07-lab-"))
    proc, log_path = start_server(workdir)
    started = time.perf_counter()
    try:
        # 1. Discovery: the 401 points at the metadata, the metadata points at the issuer.
        door = httpx2.post(URL, json={})
        challenge = parse_challenge(door.headers["www-authenticate"])
        metadata = httpx2.get(challenge["resource_metadata"]).json()
        check("no token gets 401 + resource_metadata", door.status_code == 401 and "resource_metadata" in challenge)
        check("metadata names this resource", metadata["resource"] == URL)
        check("metadata names the issuer", metadata["authorization_servers"] == [os.environ["NOTES_AUTH_ISSUER"]])

        # 2. Token matrix at the HTTP layer.
        tokens = make_tokens()
        expected = {"read": 200, "read+write": 200, "write-only": 403, "wrong-audience": 401,
                    "expired": 401, "wrong-key": 401, "wrong-issuer": 401}
        for name, status in expected.items():
            got, _ = probe(tokens[name])
            check(f"{name} token gets HTTP {status}", got == status)

        # 3. Scope rules inside the tools, through a real MCP client.
        anyio.run(mcp_checks, tokens)

        # 4. The server log names each rejection reason and never contains a token.
        time.sleep(0.2)
        log = log_path.read_text()
        reasons = collections.Counter(re.findall(r"rejected token: (\w+)", log))
        check("log names InvalidAudienceError", reasons["InvalidAudienceError"] > 0)
        check("log contains no token", not any(t in log for t in tokens.values()))
        check("note written to the scratch copy only", (workdir / "notes" / "lab-write-attempt.md").exists()
              and not Path("notes/lab-write-attempt.md").exists())
    finally:
        proc.terminate()
        proc.wait(timeout=10)

    for name, ok in checks:
        print(f"{'PASS' if ok else 'FAIL'}  {name}")
    print("rejections in server log:", dict(sorted(reasons.items())))
    print(f"{sum(ok for _, ok in checks)}/{len(checks)} checks passed in {time.perf_counter() - started:.1f} s")


if __name__ == "__main__":
    main()