python-sdk/examples/servers/simple-auth/mcp_simple_auth/github_oauth_provider.py

266 lines
9.5 KiB
Python

"""
Shared GitHub OAuth provider for MCP servers.
This module contains the common GitHub OAuth functionality used by both
the standalone authorization server and the legacy combined server.
NOTE: this is a simplified example for demonstration purposes.
This is not a production-ready implementation.
"""
import logging
import secrets
import time
from typing import Any
from pydantic import AnyHttpUrl
from pydantic_settings import BaseSettings, SettingsConfigDict
from starlette.exceptions import HTTPException
from mcp.server.auth.provider import (
AccessToken,
AuthorizationCode,
AuthorizationParams,
OAuthAuthorizationServerProvider,
RefreshToken,
construct_redirect_uri,
)
from mcp.shared._httpx_utils import create_mcp_http_client
from mcp.shared.auth import OAuthClientInformationFull, OAuthToken
logger = logging.getLogger(__name__)
class GitHubOAuthSettings(BaseSettings):
"""Common GitHub OAuth settings."""
model_config = SettingsConfigDict(env_prefix="MCP_")
# GitHub OAuth settings - MUST be provided via environment variables
github_client_id: str | None = None
github_client_secret: str | None = None
# GitHub OAuth URLs
github_auth_url: str = "https://github.com/login/oauth/authorize"
github_token_url: str = "https://github.com/login/oauth/access_token"
mcp_scope: str = "user"
github_scope: str = "read:user"
class GitHubOAuthProvider(OAuthAuthorizationServerProvider):
"""
OAuth provider that uses GitHub as the identity provider.
This provider handles the OAuth flow by:
1. Redirecting users to GitHub for authentication
2. Exchanging GitHub tokens for MCP tokens
3. Maintaining token mappings for API access
"""
def __init__(self, settings: GitHubOAuthSettings, github_callback_url: str):
self.settings = settings
self.github_callback_url = github_callback_url
self.clients: dict[str, OAuthClientInformationFull] = {}
self.auth_codes: dict[str, AuthorizationCode] = {}
self.tokens: dict[str, AccessToken] = {}
self.state_mapping: dict[str, dict[str, str | None]] = {}
# Maps MCP tokens to GitHub tokens
self.token_mapping: dict[str, str] = {}
async def get_client(self, client_id: str) -> OAuthClientInformationFull | None:
"""Get OAuth client information."""
return self.clients.get(client_id)
async def register_client(self, client_info: OAuthClientInformationFull):
"""Register a new OAuth client."""
self.clients[client_info.client_id] = client_info
async def authorize(self, client: OAuthClientInformationFull, params: AuthorizationParams) -> str:
"""Generate an authorization URL for GitHub OAuth flow."""
state = params.state or secrets.token_hex(16)
# Store state mapping for callback
self.state_mapping[state] = {
"redirect_uri": str(params.redirect_uri),
"code_challenge": params.code_challenge,
"redirect_uri_provided_explicitly": str(params.redirect_uri_provided_explicitly),
"client_id": client.client_id,
"resource": params.resource, # RFC 8707
}
# Build GitHub authorization URL
auth_url = (
f"{self.settings.github_auth_url}"
f"?client_id={self.settings.github_client_id}"
f"&redirect_uri={self.github_callback_url}"
f"&scope={self.settings.github_scope}"
f"&state={state}"
)
return auth_url
async def handle_github_callback(self, code: str, state: str) -> str:
"""Handle GitHub OAuth callback and return redirect URI."""
state_data = self.state_mapping.get(state)
if not state_data:
raise HTTPException(400, "Invalid state parameter")
redirect_uri = state_data["redirect_uri"]
code_challenge = state_data["code_challenge"]
redirect_uri_provided_explicitly = state_data["redirect_uri_provided_explicitly"] == "True"
client_id = state_data["client_id"]
resource = state_data.get("resource") # RFC 8707
# These are required values from our own state mapping
assert redirect_uri is not None
assert code_challenge is not None
assert client_id is not None
# Exchange code for token with GitHub
async with create_mcp_http_client() as client:
response = await client.post(
self.settings.github_token_url,
data={
"client_id": self.settings.github_client_id,
"client_secret": self.settings.github_client_secret,
"code": code,
"redirect_uri": self.github_callback_url,
},
headers={"Accept": "application/json"},
)
if response.status_code != 200:
raise HTTPException(400, "Failed to exchange code for token")
data = response.json()
if "error" in data:
raise HTTPException(400, data.get("error_description", data["error"]))
github_token = data["access_token"]
# Create MCP authorization code
new_code = f"mcp_{secrets.token_hex(16)}"
auth_code = AuthorizationCode(
code=new_code,
client_id=client_id,
redirect_uri=AnyHttpUrl(redirect_uri),
redirect_uri_provided_explicitly=redirect_uri_provided_explicitly,
expires_at=time.time() + 300,
scopes=[self.settings.mcp_scope],
code_challenge=code_challenge,
resource=resource, # RFC 8707
)
self.auth_codes[new_code] = auth_code
# Store GitHub token with MCP client_id
self.tokens[github_token] = AccessToken(
token=github_token,
client_id=client_id,
scopes=[self.settings.github_scope],
expires_at=None,
)
del self.state_mapping[state]
return construct_redirect_uri(redirect_uri, code=new_code, state=state)
async def load_authorization_code(
self, client: OAuthClientInformationFull, authorization_code: str
) -> AuthorizationCode | None:
"""Load an authorization code."""
return self.auth_codes.get(authorization_code)
async def exchange_authorization_code(
self, client: OAuthClientInformationFull, authorization_code: AuthorizationCode
) -> OAuthToken:
"""Exchange authorization code for tokens."""
if authorization_code.code not in self.auth_codes:
raise ValueError("Invalid authorization code")
# Generate MCP access token
mcp_token = f"mcp_{secrets.token_hex(32)}"
# Store MCP token
self.tokens[mcp_token] = AccessToken(
token=mcp_token,
client_id=client.client_id,
scopes=authorization_code.scopes,
expires_at=int(time.time()) + 3600,
resource=authorization_code.resource, # RFC 8707
)
# Find GitHub token for this client
github_token = next(
(
token
for token, data in self.tokens.items()
if (token.startswith("ghu_") or token.startswith("gho_")) and data.client_id == client.client_id
),
None,
)
# Store mapping between MCP token and GitHub token
if github_token:
self.token_mapping[mcp_token] = github_token
del self.auth_codes[authorization_code.code]
return OAuthToken(
access_token=mcp_token,
token_type="Bearer",
expires_in=3600,
scope=" ".join(authorization_code.scopes),
)
async def load_access_token(self, token: str) -> AccessToken | None:
"""Load and validate an access token."""
access_token = self.tokens.get(token)
if not access_token:
return None
# Check if expired
if access_token.expires_at and access_token.expires_at < time.time():
del self.tokens[token]
return None
return access_token
async def load_refresh_token(self, client: OAuthClientInformationFull, refresh_token: str) -> RefreshToken | None:
"""Load a refresh token - not supported in this example."""
return None
async def exchange_refresh_token(
self,
client: OAuthClientInformationFull,
refresh_token: RefreshToken,
scopes: list[str],
) -> OAuthToken:
"""Exchange refresh token - not supported in this example."""
raise NotImplementedError("Refresh tokens not supported")
async def revoke_token(self, token: str, token_type_hint: str | None = None) -> None:
"""Revoke a token."""
if token in self.tokens:
del self.tokens[token]
async def get_github_user_info(self, mcp_token: str) -> dict[str, Any]:
"""Get GitHub user info using MCP token."""
github_token = self.token_mapping.get(mcp_token)
if not github_token:
raise ValueError("No GitHub token found for MCP token")
async with create_mcp_http_client() as client:
response = await client.get(
"https://api.github.com/user",
headers={
"Authorization": f"Bearer {github_token}",
"Accept": "application/vnd.github.v3+json",
},
)
if response.status_code != 200:
raise ValueError(f"GitHub API error: {response.status_code}")
return response.json()