AegisGitea-MCP/tests/test_oauth_helpers.py

"""Shared builders for the per-user OAuth tests: settings, a fake clock and the browser flow."""

from __future__ import annotations

import base64
import hashlib
import json
import re
import secrets
import time
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from typing import Any
from urllib.parse import parse_qs, urlsplit

import httpx
import httpx2
from mcp import Client
from mcp.client.streamable_http import streamable_http_client
from pytest_httpx import HTTPXMock
from starlette.applications import Starlette
from starlette.testclient import TestClient

from aegis_gitea_mcp.audit import AuditLogger
from aegis_gitea_mcp.config import Settings
from aegis_gitea_mcp.oauth.components import build_oauth_components
from aegis_gitea_mcp.server import build_server
from aegis_gitea_mcp.transports import HttpAuthExtras, build_http_app
from tests.test_executor_helpers import API, BEARER_TOKEN, GITEA_TOKEN

BASE = "https://mcp.example.com"
RESOURCE = f"{BASE}/mcp"
GITEA = "https://gitea.example.com"
TOKEN_URL = f"{GITEA}/login/oauth/access_token"
USER_URL = f"{API}/user"
APP_CLIENT_ID = "gitea-app-client-id"
APP_CLIENT_SECRET = "gitea-app-client-secret-0123456789abcdef"
ENCRYPTION_KEY = "token-encryption-key-0123456789abcdefghijklmnop"
USER_TOKEN = "gitea-user-access-token-0123456789abcdef"
USER_REFRESH = "gitea-user-refresh-token-0123456789abcdef"
GITEA_CODE = "gitea-authorization-code-0123456789"
REDIRECT = "http://localhost:3000/callback"
CLIENT_STATE = "client-state-value-123"

__all__ = [
    "API",
    "BEARER_TOKEN",
    "GITEA_TOKEN",
]


class FakeClock:
    """A controllable wall clock starting at the real time."""

    def __init__(self) -> None:
        self.now = time.time()

    def __call__(self) -> float:
        return self.now

    def advance(self, seconds: float) -> None:
        self.now += seconds


def oauth_settings(**overrides: Any) -> Settings:
    """A complete per-user OAuth configuration with a static bearer and an admin token."""
    values: dict[str, Any] = {
        "gitea_url": GITEA,
        "gitea_token": GITEA_TOKEN,
        "mcp_bearer_token": BEARER_TOKEN,
        "mcp_port": 8080,
        "oauth_enabled": True,
        "gitea_oauth_client_id": APP_CLIENT_ID,
        "gitea_oauth_client_secret": APP_CLIENT_SECRET,
        "token_encryption_key": ENCRYPTION_KEY,
        "public_base_url": BASE,
    }
    values.update(overrides)
    return Settings(**values)


@dataclass
class OAuthApp:
    """The HTTP app with its OAuth pieces."""

    app: Starlette
    extras: HttpAuthExtras
    settings: Settings
    clock: FakeClock

    @property
    def provider(self) -> Any:
        return self.extras.auth_server_provider

    @property
    def store(self) -> Any:
        return self.provider._store


def build_oauth_app(
    settings: Settings | None = None,
    clock: FakeClock | None = None,
    audit: AuditLogger | None = None,
) -> OAuthApp:
    """Build the full HTTP app the way ``run_http`` does, with a fake clock."""
    settings = settings or oauth_settings()
    clock = clock or FakeClock()
    extras = build_oauth_components(settings, clock=clock, audit=audit)
    server = build_server(settings, resolver=extras.resolver, executor=extras.executor)
    app = build_http_app(
        settings,
        server=server,
        token_verifier=extras.token_verifier,
        auth_server_provider=extras.auth_server_provider,
        extra_routes=extras.extra_routes,
        client_registration_options=extras.client_registration_options,
        revocation_options=extras.revocation_options,
        rate_limited_paths=extras.rate_limited_paths,
    )
    return OAuthApp(app=app, extras=extras, settings=settings, clock=clock)


def browser(app: Starlette) -> TestClient:
    """A browser-like client (cookies, no redirect following). No lifespan is started."""
    return TestClient(app, base_url=BASE, follow_redirects=False)


def pkce() -> tuple[str, str]:
    verifier = secrets.token_urlsafe(48)
    digest = hashlib.sha256(verifier.encode()).digest()
    return verifier, base64.urlsafe_b64encode(digest).decode().rstrip("=")


def query_of(location: str) -> dict[str, str]:
    return {key: values[0] for key, values in parse_qs(urlsplit(location).query).items()}


def register(client: TestClient, redirect_uris: list[str] | None = None, **extra: Any) -> Any:
    body = {
        "redirect_uris": redirect_uris or [REDIRECT],
        "client_name": "Test client",
        **extra,
    }
    return client.post("/register", json=body)


def mock_gitea_login(
    httpx_mock: HTTPXMock,
    *,
    login: str = "alice",
    is_admin: bool = False,
    access: str = USER_TOKEN,
    refresh: str | None = USER_REFRESH,
    expires_in: int | None = 3600,
) -> None:
    body: dict[str, Any] = {"access_token": access, "token_type": "bearer"}
    if refresh is not None:
        body["refresh_token"] = refresh
    if expires_in is not None:
        body["expires_in"] = expires_in
    httpx_mock.add_response(method="POST", url=TOKEN_URL, json=body)
    httpx_mock.add_response(
        method="GET", url=USER_URL, json={"login": login, "is_admin": is_admin, "id": 7}
    )


@dataclass
class SignIn:
    """Everything a completed browser flow produced."""

    client_id: str
    client_secret: str
    verifier: str
    code: str
    tokens: dict[str, Any] = field(default_factory=dict)
    gitea_url: str = ""


CSRF_FIELD = re.compile(r'name="csrf" value="([^"]+)"')


def start_flow(client: TestClient, registration: dict[str, Any], challenge: str) -> str:
    """``/authorize`` -> consent page -> approve; return Gitea's authorization URL."""
    authorize = client.get(
        "/authorize",
        params={
            "client_id": registration["client_id"],
            "redirect_uri": REDIRECT,
            "response_type": "code",
            "code_challenge": challenge,
            "code_challenge_method": "S256",
            "state": CLIENT_STATE,
            "resource": RESOURCE,
        },
    )
    assert authorize.status_code == 302, authorize.text
    consent_url = authorize.headers["location"]
    assert consent_url.startswith(f"{BASE}/oauth/consent?request=")
    page = client.get(consent_url)
    assert page.status_code == 200
    match = CSRF_FIELD.search(page.text)
    assert match is not None
    approved = client.post(
        "/oauth/consent",
        data={
            "request": query_of(consent_url)["request"],
            "csrf": match.group(1),
            "decision": "approve",
        },
        headers={"Origin": BASE},
    )
    assert approved.status_code == 303, approved.text
    return str(approved.headers["location"])


def sign_in(
    client: TestClient, httpx_mock: HTTPXMock, *, login: str = "alice", **gitea: Any
) -> SignIn:
    """Run the whole flow (DCR, consent, Gitea, callback, /token) for one user."""
    registration = register(client).json()
    verifier, challenge = pkce()
    gitea_url = start_flow(client, registration, challenge)
    mock_gitea_login(httpx_mock, login=login, **gitea)
    callback = client.get(
        "/oauth/callback", params={"code": GITEA_CODE, "state": query_of(gitea_url)["state"]}
    )
    assert callback.status_code == 302, callback.text
    code = query_of(callback.headers["location"])["code"]
    tokens = exchange_code(client, registration, code, verifier)
    assert tokens.status_code == 200, tokens.text
    return SignIn(
        client_id=registration["client_id"],
        client_secret=registration["client_secret"],
        verifier=verifier,
        code=code,
        tokens=tokens.json(),
        gitea_url=gitea_url,
    )


def exchange_code(
    client: TestClient, registration: dict[str, Any], code: str, verifier: str
) -> Any:
    return client.post(
        "/token",
        data={
            "grant_type": "authorization_code",
            "code": code,
            "redirect_uri": REDIRECT,
            "client_id": registration["client_id"],
            "client_secret": registration["client_secret"],
            "code_verifier": verifier,
        },
    )


def refresh(client: TestClient, flow: SignIn, refresh_token: str, **extra: str) -> Any:
    return client.post(
        "/token",
        data={
            "grant_type": "refresh_token",
            "refresh_token": refresh_token,
            "client_id": flow.client_id,
            "client_secret": flow.client_secret,
            **extra,
        },
    )


@asynccontextmanager
async def mcp_client(app: Starlette, token: str) -> AsyncIterator[Client]:
    """The SDK client calling ``app`` over an in-process ASGI transport with ``token``."""
    async with (
        app.router.lifespan_context(app),
        httpx2.AsyncClient(
            transport=httpx2.ASGITransport(app=app),
            headers={"Authorization": f"Bearer {token}"},
            base_url=BASE,
        ) as http_client,
        Client(streamable_http_client(f"{RESOURCE}", http_client=http_client)) as c,
    ):
        yield c


def payload(result: Any) -> dict[str, Any]:
    return json.loads(result.content[0].text)  # type: ignore[no-any-return]


def form_of(request: httpx.Request) -> dict[str, str]:
    """The form-encoded body of a request Gitea received."""
    return {key: values[0] for key, values in parse_qs(request.content.decode()).items()}


def gitea_requests(httpx_mock: HTTPXMock, url_prefix: str) -> list[httpx.Request]:
    return [r for r in httpx_mock.get_requests() if str(r.url).startswith(url_prefix)]


def mcp_status_lifespan(app: Starlette, token: str) -> int:
    """HTTP status of a bare MCP request with ``token``, with the app's lifespan running."""
    with browser(app) as client:
        response = client.post("/mcp", json={}, headers={"Authorization": f"Bearer {token}"})
    return response.status_code