|
| 1 | +# pyright: reportMissingImports=false |
| 2 | +import os |
| 3 | +import logging |
| 4 | +from dotenv import load_dotenv # type: ignore |
| 5 | +from typing import Any, cast |
| 6 | +import base64, json, time |
| 7 | +from starlette.requests import Request # type: ignore |
| 8 | + |
| 9 | +from mcp.server.fastmcp.server import Context |
| 10 | +from mcp.server.auth.proxy.server import build_proxy_server # noqa: E402 |
| 11 | +from mcp.server.auth.providers.transparent_proxy import ProxySettings # type: ignore |
| 12 | + |
| 13 | +# Load environment variables from .env if present |
| 14 | +load_dotenv() |
| 15 | + |
| 16 | +# Configure logging after .env so LOG_LEVEL can come from environment |
| 17 | +LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO").upper() |
| 18 | + |
| 19 | +logging.basicConfig( |
| 20 | + level=LOG_LEVEL, |
| 21 | + format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", |
| 22 | + datefmt="%Y-%m-%d %H:%M:%S", |
| 23 | +) |
| 24 | + |
| 25 | +# Dedicated logger for this server module |
| 26 | +logger = logging.getLogger("proxy_oauth.server") |
| 27 | + |
| 28 | +# Suppress noisy INFO messages from the FastMCP low-level server unless we are |
| 29 | +# explicitly running in DEBUG mode. These logs (e.g. "Processing request of type |
| 30 | +# ListToolsRequest") are helpful for debugging but clutter normal output. |
| 31 | + |
| 32 | +_mcp_lowlevel_logger = logging.getLogger("mcp.server.lowlevel.server") |
| 33 | +if LOG_LEVEL == "DEBUG": |
| 34 | + # In full debug mode, allow the library to emit its detailed logs |
| 35 | + _mcp_lowlevel_logger.setLevel(logging.DEBUG) |
| 36 | +else: |
| 37 | + # Otherwise, only warnings and above |
| 38 | + _mcp_lowlevel_logger.setLevel(logging.WARNING) |
| 39 | + |
| 40 | +# ---------------------------------------------------------------------------- |
| 41 | +# Environment configuration |
| 42 | +# ---------------------------------------------------------------------------- |
| 43 | +# Load and validate settings from the environment (uses .env automatically) |
| 44 | +settings = ProxySettings.load() |
| 45 | + |
| 46 | +# Upstream endpoints (fully-qualified URLs) |
| 47 | +UPSTREAM_AUTHORIZE: str = str(settings.upstream_authorize) |
| 48 | +UPSTREAM_TOKEN: str = str(settings.upstream_token) |
| 49 | +UPSTREAM_JWKS_URI = settings.jwks_uri |
| 50 | +# Derive base URL from the authorize endpoint for convenience / tests |
| 51 | +UPSTREAM_BASE: str = UPSTREAM_AUTHORIZE.rsplit("/", 1)[0] |
| 52 | + |
| 53 | +# Client credentials & defaults |
| 54 | +CLIENT_ID: str = settings.client_id or "demo-client-id" |
| 55 | +CLIENT_SECRET = settings.client_secret |
| 56 | +DEFAULT_SCOPE: str = settings.default_scope |
| 57 | + |
| 58 | +# Optional audience passthrough (not part of ProxySettings yet) |
| 59 | +AUDIENCE = os.getenv("PROXY_AUDIENCE") |
| 60 | + |
| 61 | +# Metadata URL (only used if we need to fetch from upstream) |
| 62 | +UPSTREAM_METADATA = f"{UPSTREAM_BASE}/.well-known/oauth-authorization-server" |
| 63 | + |
| 64 | +# --------------------------------------------------------------------------- |
| 65 | +# Logging helpers |
| 66 | +# --------------------------------------------------------------------------- |
| 67 | + |
| 68 | + |
| 69 | +def _mask_secret(secret: str | None) -> str | None: # noqa: D401 |
| 70 | + """Return a masked version of the given secret. |
| 71 | +
|
| 72 | + The first and last four characters are preserved (if available) and the |
| 73 | + middle section is replaced by asterisks. If the secret is shorter than |
| 74 | + eight characters, the entire value is replaced by ``*``. |
| 75 | + """ |
| 76 | + |
| 77 | + if not secret: |
| 78 | + return None |
| 79 | + |
| 80 | + if len(secret) <= 8: |
| 81 | + return "*" * len(secret) |
| 82 | + |
| 83 | + return f"{secret[:4]}{'*' * (len(secret) - 8)}{secret[-4:]}" |
| 84 | + |
| 85 | + |
| 86 | +# Consolidated configuration (with sensitive data redacted) |
| 87 | +_masked_settings = settings.model_dump(exclude_none=True).copy() |
| 88 | + |
| 89 | +if "client_secret" in _masked_settings: |
| 90 | + _masked_settings["client_secret"] = _mask_secret(_masked_settings["client_secret"]) |
| 91 | + |
| 92 | +# Log configuration at *debug* level only so it can be enabled when needed |
| 93 | +logger.debug("[Proxy Config] %s", _masked_settings) |
| 94 | + |
| 95 | +# Server host/port |
| 96 | +PROXY_PORT = int(os.getenv("PROXY_PORT", "8000")) |
| 97 | + |
| 98 | +# ---------------------------------------------------------------------------- |
| 99 | +# FastMCP server (now created via library helper) |
| 100 | +# ---------------------------------------------------------------------------- |
| 101 | + |
| 102 | +ISSUER_URL = os.getenv("PROXY_ISSUER_URL", "http://localhost:8000") |
| 103 | + |
| 104 | +# Create FastMCP instance using the reusable proxy builder |
| 105 | +mcp = build_proxy_server(port=PROXY_PORT, issuer_url=ISSUER_URL) |
| 106 | + |
| 107 | +# --------------------------------------------------------------------------- |
| 108 | +# Minimal demo tool |
| 109 | +# --------------------------------------------------------------------------- |
| 110 | + |
| 111 | + |
| 112 | +@mcp.tool() |
| 113 | +def echo(message: str) -> str: |
| 114 | + return f"Echo: {message}" |
| 115 | + |
| 116 | + |
| 117 | +@mcp.tool() |
| 118 | +async def user_info(ctx: Context[Any, Any, Request]) -> dict[str, Any]: |
| 119 | + """ |
| 120 | + Get information about the authenticated user. |
| 121 | +
|
| 122 | + This tool demonstrates accessing user information from the OAuth access token. |
| 123 | + The user must be authenticated via OAuth to access this tool. |
| 124 | +
|
| 125 | + Returns: |
| 126 | + Dictionary containing user information from the access token |
| 127 | + """ |
| 128 | + from mcp.server.auth.middleware.auth_context import get_access_token |
| 129 | + |
| 130 | + # Get the access token from the authentication context |
| 131 | + access_token = get_access_token() |
| 132 | + |
| 133 | + if not access_token: |
| 134 | + return { |
| 135 | + "error": "No access token found - user not authenticated", |
| 136 | + "authenticated": False, |
| 137 | + } |
| 138 | + |
| 139 | + # Attempt to decode the access token as JWT to extract useful user claims. |
| 140 | + # Many OAuth providers issue JWT access tokens (or ID tokens) that contain |
| 141 | + # the user's subject (sub) and preferred username. We parse the token |
| 142 | + # *without* signature verification – we only need the public claims for |
| 143 | + # display purposes. If the token is opaque or the decode fails, we simply |
| 144 | + # skip this step. |
| 145 | + |
| 146 | + def _try_decode_jwt(token_str: str) -> dict[str, Any] | None: # noqa: D401 |
| 147 | + """Best-effort JWT decode without verification. |
| 148 | +
|
| 149 | + Returns the payload dictionary if the token *looks* like a JWT and can |
| 150 | + be base64-decoded. If anything fails we return None. |
| 151 | + """ |
| 152 | + |
| 153 | + try: |
| 154 | + parts = token_str.split(".") |
| 155 | + if len(parts) != 3: |
| 156 | + return None # Not a JWT |
| 157 | + |
| 158 | + # JWT parts are URL-safe base64 without padding |
| 159 | + def _b64decode(segment: str) -> bytes: |
| 160 | + padding = "=" * (-len(segment) % 4) |
| 161 | + return base64.urlsafe_b64decode(segment + padding) |
| 162 | + |
| 163 | + payload_bytes = _b64decode(parts[1]) |
| 164 | + return json.loads(payload_bytes) |
| 165 | + except Exception: # noqa: BLE001 |
| 166 | + return None |
| 167 | + |
| 168 | + jwt_claims = _try_decode_jwt(access_token.token) |
| 169 | + |
| 170 | + # Build response with token information plus any extracted claims |
| 171 | + response: dict[str, Any] = { |
| 172 | + "authenticated": True, |
| 173 | + "client_id": access_token.client_id, |
| 174 | + "scopes": access_token.scopes, |
| 175 | + "token_type": "Bearer", |
| 176 | + "expires_at": access_token.expires_at, |
| 177 | + "resource": access_token.resource, |
| 178 | + } |
| 179 | + |
| 180 | + if jwt_claims: |
| 181 | + # Prefer the `userid` claim used in FastMCP examples; fall back to `sub` if absent. |
| 182 | + uid = jwt_claims.get("userid") or jwt_claims.get("sub") |
| 183 | + if uid is not None: |
| 184 | + response["userid"] = uid # camelCase variant used in FastMCP reference |
| 185 | + response["user_id"] = uid # snake_case variant |
| 186 | + response["username"] = ( |
| 187 | + jwt_claims.get("preferred_username") |
| 188 | + or jwt_claims.get("nickname") |
| 189 | + or jwt_claims.get("name") |
| 190 | + ) |
| 191 | + response["issuer"] = jwt_claims.get("iss") |
| 192 | + response["audience"] = jwt_claims.get("aud") |
| 193 | + response["issued_at"] = jwt_claims.get("iat") |
| 194 | + |
| 195 | + # Calculate expiration helpers |
| 196 | + if access_token.expires_at: |
| 197 | + response["expires_at_iso"] = time.strftime( |
| 198 | + "%Y-%m-%dT%H:%M:%S", time.localtime(access_token.expires_at) |
| 199 | + ) |
| 200 | + response["expires_in_seconds"] = max( |
| 201 | + 0, access_token.expires_at - int(time.time()) |
| 202 | + ) |
| 203 | + |
| 204 | + return response |
| 205 | + |
| 206 | + |
| 207 | +@mcp.tool() |
| 208 | +async def test_endpoint(message: str = "Hello from proxy server!") -> dict[str, Any]: |
| 209 | + """ |
| 210 | + Test endpoint for debugging OAuth proxy functionality. |
| 211 | +
|
| 212 | + Args: |
| 213 | + message: Optional message to echo back |
| 214 | +
|
| 215 | + Returns: |
| 216 | + Test response with server information |
| 217 | + """ |
| 218 | + return { |
| 219 | + "message": message, |
| 220 | + "server": "Transparent OAuth Proxy Server", |
| 221 | + "status": "active", |
| 222 | + "oauth_configured": True, |
| 223 | + } |
| 224 | + |
| 225 | + |
| 226 | +if __name__ == "__main__": |
| 227 | + mcp.run(transport="streamable-http") |
0 commit comments