server.py

import logging
import random
import secrets
import time
from pydantic import AnyHttpUrl
from starlette.requests import Request
from starlette.responses import HTMLResponse, RedirectResponse
import uvicorn

from mcp.server.fastmcp import FastMCP
from mcp.server.auth.provider import (
    AccessToken,
    AuthorizationCode,
    AuthorizationParams,
    OAuthAuthorizationServerProvider,
    RefreshToken,
    construct_redirect_uri,
)
from mcp.server.auth.settings import AuthSettings, ClientRegistrationOptions
from mcp.server.auth.middleware.auth_context import get_access_token
from mcp.shared.auth import OAuthClientInformationFull, OAuthToken

file_handler = logging.FileHandler("oauth_debug.log")
file_handler.setLevel(logging.DEBUG)
file_handler.setFormatter(
    logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
)

console_handler = logging.StreamHandler()
console_handler.setLevel(logging.INFO)
console_handler.setFormatter(
    logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
)

logging.basicConfig(level=logging.DEBUG, handlers=[file_handler, console_handler])
logger = logging.getLogger(__name__)

PORT = 6464


class FakeOAuthProvider(
    OAuthAuthorizationServerProvider[AuthorizationCode, RefreshToken, AccessToken]
):
    def __init__(self):
        self.clients: dict[str, OAuthClientInformationFull] = {}
        self.auth_codes: dict[str, AuthorizationCode] = {}
        self.tokens: dict[str, AccessToken] = {}
        self.refresh_tokens: dict[str, RefreshToken] = {}
        self.refresh_token_resources: dict[str, str | None] = {}
        self.state_mapping: dict[str, dict[str, str | None]] = {}

    async def get_client(self, client_id: str) -> OAuthClientInformationFull | None:
        return self.clients.get(client_id)

    async def register_client(self, client_info: OAuthClientInformationFull):
        logger.info(f"Registering client: {client_info.client_id}")
        self.clients[client_info.client_id] = client_info

    async def authorize(
        self, client: OAuthClientInformationFull, params: AuthorizationParams
    ) -> str:
        state = params.state or secrets.token_hex(16)
        logger.info(
            f"Authorization request from client {client.client_id}, redirecting to approval"
        )
        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,
        }
        return f"http://localhost:{PORT}/approve?state={state}"

    async def load_authorization_code(
        self, client: OAuthClientInformationFull, authorization_code: str
    ) -> AuthorizationCode | None:
        code = self.auth_codes.get(authorization_code)
        if code:
            logger.debug(
                f"Loading auth code: {authorization_code[:20]}... for client {client.client_id}"
            )
        else:
            logger.warning(
                f"Auth code not found: {authorization_code[:20]}... for client {client.client_id}"
            )
        return code

    async def exchange_authorization_code(
        self, client: OAuthClientInformationFull, authorization_code: AuthorizationCode
    ) -> OAuthToken:
        token = f"fake_token_{secrets.token_hex(16)}"
        refresh_token_str = f"refresh_{secrets.token_hex(16)}"

        logger.info(
            f"Generated new access token: {token[:20]}... (length: {len(token)}, expires in 30s)"
        )

        self.tokens[token] = AccessToken(
            token=token,
            client_id=client.client_id,
            scopes=authorization_code.scopes,
            expires_at=int(time.time()) + 30,
            resource=authorization_code.resource,
        )

        self.refresh_tokens[refresh_token_str] = RefreshToken(
            token=refresh_token_str,
            client_id=client.client_id,
            scopes=authorization_code.scopes,
            expires_at=int(time.time()) + 3600,
        )
        self.refresh_token_resources[refresh_token_str] = authorization_code.resource

        del self.auth_codes[authorization_code.code]
        return OAuthToken(
            access_token=token,
            token_type="Bearer",
            expires_in=30,
            refresh_token=refresh_token_str,
            scope=" ".join(authorization_code.scopes),
        )

    async def load_access_token(self, token: str) -> AccessToken | None:
        access_token = self.tokens.get(token)
        if access_token:
            expires_in = (
                access_token.expires_at - int(time.time())
                if access_token.expires_at
                else None
            )
            is_expired = access_token.expires_at and access_token.expires_at < int(
                time.time()
            )
            logger.debug(
                f"Validating token: {token[:20]}... (expires in {expires_in}s, expired={is_expired})"
            )
        else:
            logger.debug(f"Token validation failed: {token[:20]}... (not found)")
        return access_token

    async def load_refresh_token(
        self, client: OAuthClientInformationFull, refresh_token: str
    ) -> RefreshToken | None:
        token = self.refresh_tokens.get(refresh_token)
        if not token:
            logger.warning(f"Refresh token not found: {refresh_token[:20]}...")
            return None
        if token.expires_at and token.expires_at < time.time():
            logger.warning(f"Refresh token expired: {refresh_token[:20]}...")
            del self.refresh_tokens[refresh_token]
            return None
        expires_in = token.expires_at - int(time.time()) if token.expires_at else None
        logger.info(
            f"Loading refresh token: {refresh_token[:20]}... (expires in {expires_in}s)"
        )
        return token

    async def exchange_refresh_token(
        self,
        client: OAuthClientInformationFull,
        refresh_token: RefreshToken,
        scopes: list[str],
    ) -> OAuthToken:
        new_token = f"fake_token_{secrets.token_hex(16)}"
        resource = self.refresh_token_resources.get(refresh_token.token)

        logger.info(
            f"Refreshed access token: {new_token[:20]}... (length: {len(new_token)}, expires in 30s)"
        )

        self.tokens[new_token] = AccessToken(
            token=new_token,
            client_id=client.client_id,
            scopes=scopes,
            expires_at=int(time.time()) + 30,
            resource=resource,
        )

        return OAuthToken(
            access_token=new_token,
            token_type="Bearer",
            expires_in=30,
            refresh_token=refresh_token.token,
            scope=" ".join(scopes),
        )

    async def revoke_token(
        self, token: AccessToken | RefreshToken, token_type_hint: str | None = None
    ) -> None:
        if isinstance(token, AccessToken) and token.token in self.tokens:
            del self.tokens[token.token]
        elif isinstance(token, RefreshToken) and token.token in self.refresh_tokens:
            del self.refresh_tokens[token.token]


oauth_provider = FakeOAuthProvider()

mcp = FastMCP(
    "dice",
    auth_server_provider=oauth_provider,
    auth=AuthSettings(
        issuer_url=AnyHttpUrl(f"http://localhost:{PORT}"),
        client_registration_options=ClientRegistrationOptions(
            enabled=True,
            valid_scopes=["dice"],
            default_scopes=["dice"],
        ),
        required_scopes=["dice"],
        resource_server_url=None,
    ),
)

@mcp.custom_route("/approve", methods=["GET"])
async def auto_approve(request: Request):
    state = request.query_params.get("state")
    logger.debug(f"Auto-approve endpoint hit with state: {state}")

    if not state:
        logger.warning("Missing state parameter in approval")
        return HTMLResponse("Missing state parameter", status_code=400)

    state_data = oauth_provider.state_mapping.get(state)
    if not state_data:
        logger.warning(f"Invalid state in approval: {state}")
        return HTMLResponse("Invalid state", status_code=400)

    client_id = state_data["client_id"]
    redirect_uri = state_data["redirect_uri"]
    code_challenge = state_data["code_challenge"]
    redirect_uri_explicit = state_data["redirect_uri_provided_explicitly"]

    if not client_id or not redirect_uri or not code_challenge:
        logger.warning("Missing required state data")
        return HTMLResponse("Invalid state data", status_code=400)

    code = f"code_{secrets.token_hex(16)}"
    logger.info(f"Generated auth code: {code[:20]}... for client {client_id}")

    oauth_provider.auth_codes[code] = AuthorizationCode(
        code=code,
        client_id=client_id,
        redirect_uri=AnyHttpUrl(redirect_uri),
        redirect_uri_provided_explicitly=redirect_uri_explicit == "True",
        expires_at=time.time() + 300,
        scopes=["dice"],
        code_challenge=code_challenge,
        resource=state_data["resource"],
    )
    del oauth_provider.state_mapping[state]

    redirect_url = construct_redirect_uri(redirect_uri, code=code, state=state)
    logger.debug(f"Redirecting to: {redirect_url}")
    return RedirectResponse(url=redirect_url, status_code=302)


@mcp.tool()
def roll_dice(sides: int = 6) -> int:
    access_token = get_access_token()
    if access_token:
        token_str = access_token.token
        token_preview = token_str[:20] + "..."
        expires_in = (
            access_token.expires_at - int(time.time())
            if access_token.expires_at
            else None
        )
        logger.info(
            f"Tool called with token: {token_preview} (expires in {expires_in}s)"
        )
    else:
        logger.warning("Tool called without authentication")

    result = random.randint(1, sides)
    logger.info(f"Rolled {sides}-sided dice: {result}")
    return result


if __name__ == "__main__":
    app = mcp.streamable_http_app()
    uvicorn.run(app, port=PORT)