stroblme
6c710e7f8d
Superuser-only management for agents that registered themselves: list them with whether anyone approved them, and withdraw one without rotating the signing key and cutting off every other agent.
268 lines
8.9 KiB
Python
268 lines
8.9 KiB
Python
"""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
|