mcp-authflow 0.1.0__py3-none-any.whl

mcp_auth_framework/__init__.py

@@ -0,0 +1,81 @@
+"""MCP Auth Framework - Reusable OAuth Authorization Server components.
+
+Provides building blocks for OAuth 2.0 authorization servers that protect
+MCP (Model Context Protocol) tool access:
+
+- **Token storage** — Pluggable backends (in-memory, PostgreSQL) for access
+  and refresh tokens.
+- **CORS** — Origin validation helpers for OAuth/MCP endpoints.
+- **Rate limiting** — Sliding-window rate limiter for token endpoints.
+- **Responses** — Standardized OAuth 2.0 error responses (RFC 6749).
+- **Validation** — Input sanitization for OAuth identifiers and scopes.
+"""
+
+from mcp_auth_framework.cors import build_cors_headers, get_cors_origin, parse_allowed_origins
+from mcp_auth_framework.rate_limiting import SlidingWindowRateLimiter
+from mcp_auth_framework.responses import (
+    OAUTH_NO_CACHE_HEADERS,
+    backend_connection_error,
+    backend_invalid_response,
+    backend_oauth_error,
+    backend_timeout,
+    invalid_client,
+    invalid_grant,
+    invalid_request,
+    invalid_scope,
+    oauth_error,
+    rate_limit_exceeded,
+    server_error,
+    slow_down,
+)
+from mcp_auth_framework.storage import MemoryTokenStorage, TokenStorage
+from mcp_auth_framework.validation import (
+    VALID_ID_PATTERN,
+    parse_json_field,
+    parse_scope_field,
+    validate_client_id,
+)
+
+__all__ = [
+    # Version
+    "__version__",
+    # CORS
+    "build_cors_headers",
+    "get_cors_origin",
+    "parse_allowed_origins",
+    # Rate limiting
+    "SlidingWindowRateLimiter",
+    # OAuth responses
+    "OAUTH_NO_CACHE_HEADERS",
+    "backend_connection_error",
+    "backend_invalid_response",
+    "backend_oauth_error",
+    "backend_timeout",
+    "invalid_client",
+    "invalid_grant",
+    "invalid_request",
+    "invalid_scope",
+    "oauth_error",
+    "rate_limit_exceeded",
+    "server_error",
+    "slow_down",
+    # Storage
+    "MemoryTokenStorage",
+    "PostgresTokenStorage",
+    "TokenStorage",
+    # Validation
+    "VALID_ID_PATTERN",
+    "parse_json_field",
+    "parse_scope_field",
+    "validate_client_id",
+]
+
+__version__ = "0.1.0"
+
+
+def __getattr__(name: str) -> type:
+    if name == "PostgresTokenStorage":
+        from mcp_auth_framework.storage.postgres import PostgresTokenStorage  # noqa: PLC0415
+
+        return PostgresTokenStorage
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

mcp_auth_framework/cors.py

@@ -0,0 +1,63 @@
+"""CORS origin validation for MCP OAuth endpoints."""
+
+import os
+
+from starlette.requests import Request
+
+
+def parse_allowed_origins(env_var: str = "ALLOWED_MCP_ORIGINS") -> list[str]:
+    """Parse allowed CORS origins from a comma-separated environment variable.
+
+    Args:
+        env_var: Name of the environment variable to read.
+
+    Returns:
+        List of allowed origin strings, stripped of whitespace.
+    """
+    raw = os.getenv(env_var, "")
+    if not raw:
+        return []
+    return [origin.strip() for origin in raw.split(",") if origin.strip()]
+
+
+def get_cors_origin(request: Request, allowed_origins: list[str]) -> str:
+    """Get CORS origin header value based on request origin.
+
+    Only returns the origin if it's in the allowed list, otherwise returns
+    empty string to deny CORS access.
+
+    Args:
+        request: The incoming request.
+        allowed_origins: List of allowed origin strings.
+
+    Returns:
+        Origin value for Access-Control-Allow-Origin header.
+    """
+    request_origin = request.headers.get("origin", "")
+    if request_origin in allowed_origins:
+        return request_origin
+    return ""
+
+
+def build_cors_headers(request: Request, allowed_origins: list[str]) -> dict[str, str]:
+    """Build standard CORS headers for OAuth discovery endpoints.
+
+    Only includes Access-Control-Allow-Origin when the request origin is
+    in the allowlist. Omits it entirely for disallowed origins per the
+    CORS specification.
+
+    Args:
+        request: The incoming request.
+        allowed_origins: List of allowed origin strings.
+
+    Returns:
+        Dict of CORS headers.
+    """
+    headers: dict[str, str] = {
+        "Access-Control-Allow-Methods": "GET, OPTIONS",
+        "Access-Control-Allow-Headers": "*",
+    }
+    origin = get_cors_origin(request, allowed_origins)
+    if origin:
+        headers["Access-Control-Allow-Origin"] = origin
+    return headers

mcp_auth_framework/rate_limiting.py

@@ -0,0 +1,59 @@
+"""Rate limiting utilities for OAuth endpoints."""
+
+import time
+from collections import defaultdict
+
+
+class SlidingWindowRateLimiter:
+    """Simple in-memory rate limiter for OAuth endpoints.
+
+    Tracks requests per client within a sliding time window.
+    Thread-safe for async usage within a single process.
+    """
+
+    def __init__(self, requests_per_window: int, window_seconds: int):
+        """Initialize the rate limiter.
+
+        Args:
+            requests_per_window: Maximum number of requests allowed in the window
+            window_seconds: Size of the time window in seconds
+        """
+        self.requests_per_window = requests_per_window
+        self.window_seconds = window_seconds
+        self.clients: dict[str, list[float]] = defaultdict(list)
+
+    def is_allowed(self, client_id: str) -> bool:
+        """Check if the client is allowed to make a request.
+
+        Args:
+            client_id: OAuth client identifier
+
+        Returns:
+            True if the request is allowed, False if rate limited
+        """
+        now = time.time()
+        # Clean old requests outside the window
+        self.clients[client_id] = [
+            req_time for req_time in self.clients[client_id] if now - req_time < self.window_seconds
+        ]
+
+        if len(self.clients[client_id]) >= self.requests_per_window:
+            return False
+
+        self.clients[client_id].append(now)
+        return True
+
+    def get_retry_after(self, client_id: str) -> int:
+        """Get the number of seconds until the client can retry.
+
+        Args:
+            client_id: OAuth client identifier
+
+        Returns:
+            Number of seconds to wait before retrying (minimum 1)
+        """
+        if not self.clients[client_id]:
+            return 0
+        oldest_request = min(self.clients[client_id])
+        retry_after = int(self.window_seconds - (time.time() - oldest_request)) + 1
+        return max(retry_after, 1)

mcp_auth_framework/responses.py

@@ -0,0 +1,140 @@
+"""
+OAuth error response helpers for MCP Auth server.
+
+This module provides standardized error response functions following
+OAuth 2.0 error response format (RFC 6749).
+"""
+
+from starlette.responses import JSONResponse
+
+# Standard headers for OAuth responses
+OAUTH_NO_CACHE_HEADERS: dict[str, str] = {"Cache-Control": "no-store"}
+
+
+def oauth_error(
+    error: str,
+    description: str,
+    status_code: int = 400,
+    extra_headers: dict[str, str] | None = None,
+) -> JSONResponse:
+    """Create a standard OAuth error response.
+
+    Args:
+        error: OAuth error code (e.g., "invalid_request", "server_error")
+        description: Human-readable error description
+        status_code: HTTP status code (default 400)
+        extra_headers: Additional headers to include (e.g., Retry-After)
+
+    Returns:
+        JSONResponse with OAuth error format
+    """
+    headers = OAUTH_NO_CACHE_HEADERS.copy()
+    if extra_headers:
+        headers.update(extra_headers)
+
+    return JSONResponse(
+        {"error": error, "error_description": description},
+        status_code=status_code,
+        headers=headers,
+    )
+
+
+def invalid_request(description: str) -> JSONResponse:
+    """Create an invalid_request error response (400).
+
+    Use for: missing required parameters, invalid parameter format.
+    """
+    return oauth_error("invalid_request", description, 400)
+
+
+def invalid_client(description: str) -> JSONResponse:
+    """Create an invalid_client error response (401).
+
+    Use for: client authentication failed, unknown client.
+    """
+    return oauth_error("invalid_client", description, 401)
+
+
+def slow_down(description: str, retry_after: int | None = None) -> JSONResponse:
+    """Create a slow_down error response (400 or 429).
+
+    Use for: rate limiting during device flow polling.
+
+    Args:
+        description: Error description
+        retry_after: Optional retry-after value in seconds
+    """
+    extra_headers = {"Retry-After": str(retry_after)} if retry_after else None
+    return oauth_error("slow_down", description, 400, extra_headers)
+
+
+def rate_limit_exceeded(description: str, retry_after: int | None = None) -> JSONResponse:
+    """Create a rate limit exceeded error response (429).
+
+    Use for: too many requests.
+    """
+    extra_headers = {"Retry-After": str(retry_after)} if retry_after else None
+    return oauth_error("slow_down", description, 429, extra_headers)
+
+
+def server_error(description: str, status_code: int = 500) -> JSONResponse:
+    """Create a server_error response.
+
+    Use for: internal server errors, backend failures.
+
+    Args:
+        description: Error description
+        status_code: HTTP status code (500, 502, 504, etc.)
+    """
+    return oauth_error("server_error", description, status_code)
+
+
+def backend_timeout() -> JSONResponse:
+    """Create a backend timeout error response (504)."""
+    return server_error("Backend timeout", 504)
+
+
+def backend_connection_error() -> JSONResponse:
+    """Create a backend connection error response (502)."""
+    return server_error("Backend connection error", 502)
+
+
+def backend_invalid_response() -> JSONResponse:
+    """Create an invalid backend response error (502)."""
+    return server_error("Invalid response from backend", 502)
+
+
+def invalid_grant(description: str) -> JSONResponse:
+    """Create an invalid_grant error response (400).
+
+    Use for: invalid or expired authorization codes, refresh tokens, or device codes.
+    """
+    return oauth_error("invalid_grant", description, 400)
+
+
+def invalid_scope(description: str) -> JSONResponse:
+    """Create an invalid_scope error response (400).
+
+    Use for: requested scope is invalid, unknown, or exceeds what was granted.
+    """
+    return oauth_error("invalid_scope", description, 400)
+
+
+def backend_oauth_error(error_dict: dict[str, str], status_code: int) -> JSONResponse:
+    """Create a JSONResponse from an already-formatted OAuth error dict.
+
+    Use for: forwarding transformed backend errors that are already in
+    {"error": "...", "error_description": "..."} format.
+
+    Args:
+        error_dict: Dict with "error" and "error_description" keys
+        status_code: HTTP status code from the backend response
+
+    Returns:
+        JSONResponse with Cache-Control: no-store header
+    """
+    return JSONResponse(
+        error_dict,
+        status_code=status_code,
+        headers=OAUTH_NO_CACHE_HEADERS.copy(),
+    )

mcp_auth_framework/storage/__init__.py

@@ -0,0 +1,30 @@
+"""Token storage abstractions and implementations.
+
+``PostgresTokenStorage`` is importable from this module but loaded lazily
+so that ``asyncpg`` is only required when actually used.  Install the
+``postgres`` extra to enable it: ``pip install mcp-auth-framework[postgres]``
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from mcp_auth_framework.storage.base import TokenStorage
+from mcp_auth_framework.storage.memory import MemoryTokenStorage
+
+if TYPE_CHECKING:
+    from mcp_auth_framework.storage.postgres import PostgresTokenStorage
+
+__all__ = [
+    "TokenStorage",
+    "MemoryTokenStorage",
+    "PostgresTokenStorage",
+]
+
+
+def __getattr__(name: str) -> type:
+    if name == "PostgresTokenStorage":
+        from mcp_auth_framework.storage.postgres import PostgresTokenStorage  # noqa: PLC0415
+
+        return PostgresTokenStorage
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

mcp_auth_framework/storage/base.py

@@ -0,0 +1,131 @@
+"""Abstract base class for token storage implementations."""
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class TokenStorage(ABC):
+    """Abstract interface for MCP token storage."""
+
+    @abstractmethod
+    async def initialize(self) -> None:
+        """Initialize the storage backend."""
+        ...
+
+    @abstractmethod
+    async def close(self) -> None:
+        """Close the storage backend and clean up resources."""
+        ...
+
+    @abstractmethod
+    async def store_token(
+        self,
+        token: str,
+        client_id: str,
+        scopes: list[str],
+        expires_at: int,
+        resource: str | None = None,
+        user_id: int | None = None,
+    ) -> None:
+        """Store an access token.
+
+        Args:
+            token: The access token string
+            client_id: OAuth client ID
+            scopes: List of granted scopes
+            expires_at: Unix timestamp when token expires
+            resource: Optional RFC 8707 resource binding
+            user_id: Optional ID of the user who authorized the token
+        """
+        ...
+
+    @abstractmethod
+    async def load_token(self, token: str) -> dict[str, Any] | None:
+        """Load an access token.
+
+        Args:
+            token: The access token string to look up
+
+        Returns:
+            Token data dict if found and not expired, None otherwise
+        """
+        ...
+
+    @abstractmethod
+    async def delete_token(self, token: str) -> None:
+        """Delete a token.
+
+        Args:
+            token: The access token string to delete
+        """
+        ...
+
+    @abstractmethod
+    async def store_refresh_token(
+        self,
+        refresh_token: str,
+        client_id: str,
+        scopes: list[str],
+        expires_at: int,
+        resource: str | None = None,
+        user_id: int | None = None,
+    ) -> None:
+        """Store a refresh token.
+
+        Args:
+            refresh_token: The refresh token string
+            client_id: OAuth client ID
+            scopes: List of granted scopes
+            expires_at: Unix timestamp when token expires
+            resource: Optional RFC 8707 resource binding
+            user_id: Optional ID of the user who authorized the token
+        """
+        ...
+
+    @abstractmethod
+    async def load_refresh_token(self, refresh_token: str) -> dict[str, Any] | None:
+        """Load a refresh token.
+
+        Args:
+            refresh_token: The refresh token string to look up
+
+        Returns:
+            Token data dict if found and not expired, None otherwise
+        """
+        ...
+
+    @abstractmethod
+    async def delete_refresh_token(self, refresh_token: str) -> None:
+        """Delete a refresh token.
+
+        Args:
+            refresh_token: The refresh token string to delete
+        """
+        ...
+
+    @abstractmethod
+    async def cleanup_expired_tokens(self) -> int:
+        """Remove all expired access tokens.
+
+        Returns:
+            Number of tokens removed
+        """
+        ...
+
+    @abstractmethod
+    async def cleanup_expired_refresh_tokens(self) -> int:
+        """Remove all expired refresh tokens.
+
+        Returns:
+            Number of tokens removed
+        """
+        ...
+
+    @abstractmethod
+    async def get_token_count(self) -> int:
+        """Get the total number of access tokens in storage.
+
+        Returns:
+            Number of tokens stored
+        """
+        ...