app/backend/tests/api/routes/test_oauth.py

"""The authorization flow an agent goes through to reach the MCP endpoint."""

import base64
import hashlib
import secrets
import uuid
from urllib.parse import parse_qs, urlparse

import pytest
from fastapi.testclient import TestClient
from sqlmodel import Session, select

from fluksio.core.config import settings
from fluksio.models import OAuthRefreshToken

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",
    }


@pytest.mark.usefixtures("enrolled")
def test_through_the_hub_the_metadata_and_resource_are_the_hubs(
    client: TestClient, superuser_token_headers: dict[str, str]
) -> None:
    """An agent at /i/{id}/mcp on the hub is sent to the hub's addresses."""
    from fluksio.cloud import config as cloud_config
    from tests.utils.portal import INSTANCE_ID, ISSUER

    via = {cloud_config.VIA_HEADER: cloud_config.VIA_PORTAL}
    base = f"{ISSUER}/i/{INSTANCE_ID}"
    document = client.get("/.well-known/oauth-authorization-server", headers=via).json()
    assert document["issuer"] == base
    assert document["token_endpoint"] == f"{base}{PREFIX}/token"
    # No link is attached here to name the portal's page, so the hub's own.
    assert document["authorization_endpoint"] == f"{base}/oauth/authorize"
    # A request made here hears nothing about the hub.
    local = client.get("/.well-known/oauth-authorization-server").json()
    assert local["issuer"] == settings.oauth_issuer

    _verifier, challenge = pkce()
    body = {
        "client_id": register(client)["client_id"],
        "redirect_uri": REDIRECT,
        "code_challenge": challenge,
        "code_challenge_method": "S256",
        "resource": f"{base}/mcp",
    }
    through = client.post(
        f"{PREFIX}/authorize", headers={**superuser_token_headers, **via}, json=body
    )
    assert through.status_code == 200, through.text
    # Only through the hub is that a resource this instance serves.
    here = client.post(
        f"{PREFIX}/authorize", headers=superuser_token_headers, json=body
    )
    assert here.json()["error"] == "invalid_target"


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,
    db: Session,
    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())

    # And the access token already in its hands stops working, rather than
    # outliving the revocation by up to MCP_TOKEN_EXPIRE_MINUTES: the token is
    # stateless, so the client row it names is what withdrawing it removes.
    assert (
        client.post(
            f"{settings.API_V1_STR}/login/test-token",
            headers={"Authorization": f"Bearer {tokens['access_token']}"},
        ).status_code
        == 401
    )

    # 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

    # Went, rather than merely stopped working: nothing here deletes those rows
    # itself, so an orphan is a foreign key the database is not enforcing —
    # which SQLite does not do unless it is asked (`PRAGMA foreign_keys`).
    left = db.exec(
        select(OAuthRefreshToken).where(
            OAuthRefreshToken.client_id == uuid.UUID(client_id)
        )
    ).all()
    assert left == []