"""The authorization flow an agent goes through to reach the MCP endpoint.""" import base64 import hashlib import secrets from urllib.parse import parse_qs, urlparse import pytest from fastapi.testclient import TestClient from app.core.config import settings PREFIX = f"{settings.API_V1_STR}/oauth" REDIRECT = "http://127.0.0.1:41234/callback" @pytest.fixture(autouse=True) def mcp_on(monkeypatch: pytest.MonkeyPatch) -> None: """The OAuth endpoints only answer when the MCP endpoint is wanted.""" monkeypatch.setattr(settings, "MCP_ENABLED", True) def pkce() -> tuple[str, str]: verifier = secrets.token_urlsafe(48) digest = hashlib.sha256(verifier.encode("ascii")).digest() challenge = base64.urlsafe_b64encode(digest).rstrip(b"=").decode("ascii") return verifier, challenge def register(client: TestClient, **overrides) -> dict: body = {"client_name": "Test agent", "redirect_uris": [REDIRECT], **overrides} return client.post(f"{PREFIX}/register", json=body).json() def approve( client: TestClient, headers: dict[str, str], client_id: str, challenge: str ) -> str: """Walk the consent step and return the code it hands back.""" response = client.post( f"{PREFIX}/authorize", headers=headers, json={ "client_id": client_id, "redirect_uri": REDIRECT, "code_challenge": challenge, "code_challenge_method": "S256", "resource": settings.mcp_resource, }, ) assert response.status_code == 200, response.text url = response.json()["redirect_url"] return parse_qs(urlparse(url).query)["code"][0] def exchange(client: TestClient, client_id: str, code: str, verifier: str): return client.post( f"{PREFIX}/token", data={ "grant_type": "authorization_code", "code": code, "redirect_uri": REDIRECT, "client_id": client_id, "code_verifier": verifier, }, ) def test_the_metadata_says_how_to_get_in(client: TestClient) -> None: document = client.get("/.well-known/oauth-authorization-server").json() assert document["issuer"] == settings.oauth_issuer assert document["code_challenge_methods_supported"] == ["S256"] assert document["token_endpoint_auth_methods_supported"] == ["none"] assert set(document["grant_types_supported"]) == { "authorization_code", "refresh_token", } def test_a_full_authorization_reaches_the_api( client: TestClient, superuser_token_headers: dict[str, str] ) -> None: registered = register(client) verifier, challenge = pkce() code = approve(client, superuser_token_headers, registered["client_id"], challenge) tokens = exchange(client, registered["client_id"], code, verifier).json() assert tokens["token_type"] == "Bearer" assert tokens["scope"] == "mcp" # The token works against the ordinary API, as the user who approved it. whoami = client.post( f"{settings.API_V1_STR}/login/test-token", headers={"Authorization": f"Bearer {tokens['access_token']}"}, ) assert whoami.status_code == 200 def test_a_wrong_verifier_is_refused( client: TestClient, superuser_token_headers: dict[str, str] ) -> None: registered = register(client) _, challenge = pkce() code = approve(client, superuser_token_headers, registered["client_id"], challenge) response = exchange( client, registered["client_id"], code, secrets.token_urlsafe(48) ) assert response.status_code == 400 assert response.json()["error"] == "invalid_grant" def test_replaying_a_code_withdraws_what_it_produced( client: TestClient, superuser_token_headers: dict[str, str] ) -> None: registered = register(client) verifier, challenge = pkce() code = approve(client, superuser_token_headers, registered["client_id"], challenge) first = exchange(client, registered["client_id"], code, verifier).json() replay = exchange(client, registered["client_id"], code, verifier) assert replay.status_code == 400 assert replay.json()["error"] == "invalid_grant" # The refresh token the first exchange handed out is no longer any good: # the replay says someone else may be holding a copy. refreshed = client.post( f"{PREFIX}/token", data={ "grant_type": "refresh_token", "refresh_token": first["refresh_token"], "client_id": registered["client_id"], }, ) assert refreshed.status_code == 400 def test_a_refresh_token_rotates_and_cannot_be_reused( client: TestClient, superuser_token_headers: dict[str, str] ) -> None: registered = register(client) verifier, challenge = pkce() code = approve(client, superuser_token_headers, registered["client_id"], challenge) first = exchange(client, registered["client_id"], code, verifier).json() refresh = lambda token: client.post( # noqa: E731 f"{PREFIX}/token", data={ "grant_type": "refresh_token", "refresh_token": token, "client_id": registered["client_id"], }, ) second = refresh(first["refresh_token"]) assert second.status_code == 200 assert second.json()["refresh_token"] != first["refresh_token"] # Presenting the spent one again is what a stolen token looks like. assert refresh(first["refresh_token"]).status_code == 400 assert refresh(second.json()["refresh_token"]).status_code == 400 def test_an_unregistered_redirect_is_refused( client: TestClient, superuser_token_headers: dict[str, str] ) -> None: registered = register(client) _, challenge = pkce() response = client.post( f"{PREFIX}/authorize", headers=superuser_token_headers, json={ "client_id": registered["client_id"], "redirect_uri": "https://somewhere.else/callback", "code_challenge": challenge, "code_challenge_method": "S256", }, ) assert response.status_code == 400 assert response.json()["error"] == "invalid_request" def test_registration_refuses_plain_http_off_the_loopback( client: TestClient, ) -> None: response = client.post( f"{PREFIX}/register", json={"client_name": "Nope", "redirect_uris": ["http://example.com/cb"]}, ) assert response.status_code == 400 assert response.json()["error"] == "invalid_redirect_uri" def test_approving_needs_a_signed_in_user(client: TestClient) -> None: registered = register(client) _, challenge = pkce() response = client.post( f"{PREFIX}/authorize", json={ "client_id": registered["client_id"], "redirect_uri": REDIRECT, "code_challenge": challenge, "code_challenge_method": "S256", }, ) assert response.status_code == 401 def test_an_unsupported_grant_says_so(client: TestClient) -> None: response = client.post(f"{PREFIX}/token", data={"grant_type": "client_credentials"}) assert response.status_code == 400 assert response.json()["error"] == "unsupported_grant_type" def test_everything_is_refused_while_mcp_is_off( client: TestClient, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr(settings, "MCP_ENABLED", False) assert ( client.post( f"{PREFIX}/register", json={"redirect_uris": [REDIRECT]} ).status_code == 403 ) assert ( client.post( f"{PREFIX}/token", data={"grant_type": "authorization_code"} ).status_code == 403 ) def test_one_agent_can_be_revoked_without_touching_the_others( client: TestClient, superuser_token_headers: dict[str, str], normal_user_token_headers: dict[str, str], ) -> None: registered = register(client, client_name="Doomed agent") client_id = registered["client_id"] verifier, challenge = pkce() code = approve(client, superuser_token_headers, client_id, challenge) tokens = exchange(client, client_id, code, verifier).json() def listed() -> list[dict]: response = client.get(f"{PREFIX}/clients", headers=superuser_token_headers) assert response.status_code == 200 return response.json()["data"] entry = next(c for c in listed() if c["id"] == client_id) assert entry["client_name"] == "Doomed agent" # The column that tells an approved agent from one that only registered. assert entry["active_tokens"] == 1 # Withdrawing access is a superuser's job. assert ( client.get(f"{PREFIX}/clients", headers=normal_user_token_headers).status_code == 403 ) revoked = client.delete( f"{PREFIX}/clients/{client_id}", headers=superuser_token_headers ) assert revoked.status_code == 200, revoked.text assert all(c["id"] != client_id for c in listed()) # Its refresh token went with it, so it cannot mint itself a new one. refreshed = client.post( f"{PREFIX}/token", data={"grant_type": "refresh_token", "refresh_token": tokens["refresh_token"]}, ) assert refreshed.status_code == 400