OAuth 2.0 Integration Design for MCP Gateway¶

from oauthlib.oauth2 import BackendApplicationClient, WebApplicationClient
from requests_oauthlib import OAuth2Session
from typing import Optional, Dict, Any

class OAuthManager:
    """Manages OAuth 2.0 authentication flows."""

    async def get_access_token(
        self,
        credentials: Dict[str, Any]
    ) -> str:
        """Get access token based on grant type."""
        if credentials['grant_type'] == 'client_credentials':
            return await self._client_credentials_flow(credentials)
        elif credentials['grant_type'] == 'authorization_code':
            return await self._authorization_code_flow(credentials)
        else:
            raise ValueError(f"Unsupported grant type: {credentials['grant_type']}")

    async def _client_credentials_flow(
        self,
        credentials: Dict[str, Any]
    ) -> str:
        """Machine-to-machine authentication."""
        client = BackendApplicationClient(client_id=credentials['client_id'])
        oauth = OAuth2Session(client=client)

        token = oauth.fetch_token(
            token_url=credentials['token_url'],
            client_id=credentials['client_id'],
            client_secret=credentials['client_secret'],
            scope=credentials.get('scopes', [])
        )

        return token['access_token']

    async def _authorization_code_flow(
        self,
        credentials: Dict[str, Any]
    ) -> Dict[str, str]:
        """User delegation flow - returns authorization URL."""
        oauth = OAuth2Session(
            credentials['client_id'],
            redirect_uri=credentials['redirect_uri'],
            scope=credentials.get('scopes', [])
        )

        authorization_url, state = oauth.authorization_url(
            credentials['authorization_url']
        )

        return {
            'authorization_url': authorization_url,
            'state': state
        }

    async def exchange_code_for_token(
        self,
        credentials: Dict[str, Any],
        code: str,
        state: str
    ) -> str:
        """Exchange authorization code for access token."""
        oauth = OAuth2Session(
            credentials['client_id'],
            state=state,
            redirect_uri=credentials['redirect_uri']
        )

        token = oauth.fetch_token(
            credentials['token_url'],
            client_secret=credentials['client_secret'],
            authorization_response=f"{credentials['redirect_uri']}?code={code}&state={state}"
        )

        return token['access_token']

async def _initialize_gateway(
    self,
    url: str,
    authentication: Optional[Dict[str, str]] = None,
    transport: str = "SSE"
) -> tuple:
    """Initialize gateway with OAuth support."""

    headers = {}

    if authentication and authentication.get('type') == 'oauth':
        # Get OAuth credentials from database
        gateway = await self._get_gateway(authentication['gateway_id'])
        oauth_config = gateway.oauth_config

        # Get access token
        access_token = await self.oauth_manager.get_access_token(oauth_config)
        headers = {'Authorization': f'Bearer {access_token}'}
    else:
        # Existing authentication logic
        headers = decode_auth(authentication)

    # Connect to MCP server
    return await self._connect_to_gateway(url, headers, transport)

async def invoke_tool(
    self,
    db: Session,
    name: str,
    arguments: Dict[str, Any]
) -> ToolResult:
    """Invoke tool with OAuth support."""

    tool = await self.get_tool_by_name(db, name)
    headers = {}

    if tool.gateway and tool.gateway.auth_type == 'oauth':
        # Get fresh access token for each request
        oauth_config = tool.gateway.oauth_config
        access_token = await self.oauth_manager.get_access_token(oauth_config)
        headers = {'Authorization': f'Bearer {access_token}'}
    else:
        # Existing authentication
        headers = self._get_tool_headers(tool)

    # Execute tool
    return await self._execute_tool(tool, arguments, headers)

async def test_client_credentials_flow():
    oauth_manager = OAuthManager()
    credentials = {
        "grant_type": "client_credentials",
        "client_id": "test_client",
        "client_secret": "test_secret",
        "token_url": "https://oauth.example.com/token"
    }

    token = await oauth_manager.get_access_token(credentials)
    assert token is not None
    assert isinstance(token, str)

async def test_tool_invocation_with_oauth():
    tool_service = ToolService(oauth_manager)
    result = await tool_service.invoke_tool(
        db=db,
        name="github_create_issue",
        arguments={"title": "Test Issue"}
    )
    assert result.success