"""Authorization server for an MCP server that asks for OAuth, used by mcp_oauth_test.py with fake_mcp_http_server.Server(oauth=...). It publishes RFC 8414 metadata, registers clients (RFC 7591), answers the authorize request with a 302 to the client's loopback redirect (so a test follows it instead of a browser), and checks PKCE S256, the redirect and the resource parameter (RFC 8707) when it trades a code. Access tokens expire after ``expires_in`` seconds; ``revoke_access`` and ``revoke_refresh`` make the tokens it issued stop working. Every request is kept in ``seen``. """ import base64 import hashlib import json import secrets import threading import time from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from typing import cast from urllib.parse import parse_qs, urlencode, urlparse class AuthServer(ThreadingHTTPServer): daemon_threads = True def __init__(self, expires_in=3600): super().__init__(("127.0.0.1", 0), Handler) self.expires_in = expires_in self.clients = {} self.codes = {} self.access = {} self.refresh = {} self.seen = [] self.lock = threading.Lock() @property def issuer(self): return f"http://127.0.0.1:{self.server_address[1]}" def start(self): threading.Thread(target=self.serve_forever, daemon=True).start() return self def stop(self): self.shutdown() self.server_close() def scopes(self, token): """The scopes a live access token grants, or None.""" with self.lock: grant = self.access.get(token) if grant is None or grant["expires"] < time.time(): return None return set(grant["scope"].split()) def revoke_access(self): with self.lock: self.access.clear() def revoke_refresh(self): with self.lock: self.access.clear() self.refresh.clear() def requests(self, path): return [entry for entry in self.seen if entry["path"] == path] def issue(self, client, scope, resource): access, refresh = secrets.token_hex(8), secrets.token_hex(8) with self.lock: self.access[access] = { "scope": scope, "expires": time.time() + self.expires_in, } self.refresh[refresh] = { "client": client, "scope": scope, "resource": resource, } return { "access_token": access, "token_type": "Bearer", "expires_in": self.expires_in, "refresh_token": refresh, "scope": scope, "id_token": jwt({"sub": "u1", "email": "ana@example.com"}), } class Handler(BaseHTTPRequestHandler): protocol_version = "HTTP/1.1" @property def auth_server(self) -> AuthServer: return cast(AuthServer, self.server) def log_message(self, format, *args): pass def do_GET(self): url = urlparse(self.path) query = {key: values[0] for key, values in parse_qs(url.query).items()} self.auth_server.seen.append({"path": url.path, "query": query}) if url.path == "/.well-known/oauth-authorization-server": issuer = self.auth_server.issuer return self.json( 200, { "issuer": issuer, "authorization_endpoint": f"{issuer}/authorize", "token_endpoint": f"{issuer}/token", "registration_endpoint": f"{issuer}/register", "code_challenge_methods_supported": ["S256"], "scopes_supported": ["read", "admin", "offline_access"], "authorization_response_iss_parameter_supported": True, }, ) if url.path == "/authorize": return self.authorize(query) self.json(404, {"error": "not_found"}) def authorize(self, query): client = self.auth_server.clients.get(query.get("client_id")) if ( client is None or query.get("redirect_uri") not in client["redirect_uris"] or query.get("code_challenge_method") != "S256" or not query.get("code_challenge") or not query.get("resource") ): return self.json(400, {"error": "invalid_request"}) code = secrets.token_hex(8) self.auth_server.codes[code] = query location = ( query["redirect_uri"] + "?" + urlencode( {"code": code, "state": query["state"], "iss": self.auth_server.issuer} ) ) self.send_response(302) self.send_header("location", location) self.send_header("content-length", "0") self.end_headers() def do_POST(self): url = urlparse(self.path) body = self.rfile.read(int(self.headers.get("content-length") or 0)) if url.path == "/register": request = json.loads(body) self.auth_server.seen.append({"path": url.path, "body": request}) client = f"client-{len(self.auth_server.clients) + 1}" self.auth_server.clients[client] = request return self.json( 201, {"client_id": client, "token_endpoint_auth_method": "none", **request}, ) if url.path == "/token": form = {key: values[0] for key, values in parse_qs(body.decode()).items()} self.auth_server.seen.append({"path": url.path, "body": form}) return self.token(form) self.json(404, {"error": "not_found"}) def token(self, form): if form.get("grant_type") == "authorization_code": grant = self.auth_server.codes.pop(form.get("code"), None) verifier = form.get("code_verifier", "") challenge = ( base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()) .rstrip(b"=") .decode() ) if ( grant is None or grant["code_challenge"] != challenge or grant["redirect_uri"] != form.get("redirect_uri") or grant["client_id"] != form.get("client_id") or grant["resource"] != form.get("resource") ): return self.json(400, {"error": "invalid_grant"}) scope = grant.get("scope", "") return self.json( 200, self.auth_server.issue(grant["client_id"], scope, grant["resource"]), ) if form.get("grant_type") == "refresh_token": with self.auth_server.lock: grant = self.auth_server.refresh.pop(form.get("refresh_token"), None) if grant is None or grant["resource"] != form.get("resource"): return self.json(400, {"error": "invalid_grant"}) return self.json( 200, self.auth_server.issue( grant["client"], grant["scope"], grant["resource"] ), ) self.json(400, {"error": "unsupported_grant_type"}) def json(self, status, value): body = json.dumps(value).encode() self.send_response(status) self.send_header("content-type", "application/json") self.send_header("content-length", str(len(body))) self.end_headers() self.wfile.write(body) def jwt(claims): def part(value): return ( base64.urlsafe_b64encode(json.dumps(value).encode()).rstrip(b"=").decode() ) return f"{part({'alg': 'none'})}.{part(claims)}.sig"