Add Okta OAuth to a Calculator MCP Server

import time
import requests
import base64

class TokenManager:
    """Manages OAuth access tokens for machine-to-machine authentication."""

    def __init__(self, token_endpoint, client_id, client_secret, audience, scope=None, **kwargs):
        """
        Initialize the TokenManager.
        
        Args:
            token_endpoint: Your OAuth provider's token endpoint URL
            client_id: Your OAuth client ID
            client_secret: Your OAuth client secret
            audience: The audience value configured in your authorization server
            scope: The scopes for the token (optional, provider-specific)
            **kwargs: Additional provider-specific parameters (e.g., auth_server_id for Okta)
        """
        self.token_endpoint = token_endpoint
        self.client_id = client_id
        self.client_secret = client_secret
        self.audience = audience
        self.scope = scope
        self.extra_params = kwargs
        self.token = None
        self.token_expiry = 0

    def get_token(self):
        """
        Get a valid access token, requesting a new one if necessary.
        
        Returns:
            str: A valid access token
        """
        # Return cached token if still valid (with 5 minute buffer)
        if self.token and time.time() < self.token_expiry:
            return self.token

        # Fetch new token
        credentials = base64.b64encode(
            f"{self.client_id}:{self.client_secret}".encode()
        ).decode()

        data = {
            "grant_type": "client_credentials",
            "audience": self.audience
        }

        if self.scope:
            data["scope"] = " ".join(self.scope) if isinstance(self.scope, list) else self.scope

        # Add provider-specific parameters
        data.update(self.extra_params)

        response = requests.post(
            self.token_endpoint,
            headers={
                "Authorization": f"Basic {credentials}",
                "Content-Type": "application/x-www-form-urlencoded"
            },
            data=data
        )

        response.raise_for_status()
        token_data = response.json()
        self.token = token_data["access_token"]

        # Request new token 5 minutes before actual expiry
        expires_in = token_data.get("expires_in", 3600)
        self.token_expiry = time.time() + expires_in - 300

        return self.token

# Example usage with Okta
token_manager = TokenManager(
    token_endpoint="https://dev-12345678.okta.com/oauth2/aus123abc/v1/token",
    client_id="0oa123abc...",
    client_secret="secret123...",
    audience="https://your-mcp-server.example.com",
    scope=["my_scope"]
)

# Get a token (will be cached until near expiry)
access_token = token_manager.get_token()

# Use the token to make authenticated requests
import requests
response = requests.post(
    "http://localhost:8000/mcp",
    headers={
        "Authorization": f"Bearer {access_token}",
        "Content-Type": "application/json"
    },
    json={
        "jsonrpc": "2.0",
        "id": 1,
        "method": "tools/list"
    }
)