身分斷言

import time
import uuid

import httpx2
import jwt

from mcp import Client
from mcp.client.auth.extensions.identity_assertion import IdentityAssertionOAuthProvider
from mcp.client.streamable_http import streamable_http_client
from mcp.shared.auth import OAuthClientInformationFull, OAuthToken

IDP_SIGNING_KEY = "the-enterprise-idp-signing-key-for-this-demo"


class InMemoryTokenStorage:
    def __init__(self) -> None:
        self.tokens: OAuthToken | None = None
        self.client_info: OAuthClientInformationFull | None = None

    async def get_tokens(self) -> OAuthToken | None:
        return self.tokens

    async def set_tokens(self, tokens: OAuthToken) -> None:
        self.tokens = tokens

    async def get_client_info(self) -> OAuthClientInformationFull | None:
        return self.client_info

    async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
        self.client_info = client_info


def idp_issue_id_jag(subject: str, audience: str, resource: str) -> str:
    now = int(time.time())
    claims = {
        "iss": "https://idp.example.com",
        "sub": subject,
        "aud": audience,
        "client_id": "finance-agent",
        "resource": resource,
        "scope": "notes:read",
        "jti": str(uuid.uuid4()),
        "iat": now,
        "exp": now + 300,
    }
    return jwt.encode(claims, IDP_SIGNING_KEY, algorithm="HS256", headers={"typ": "oauth-id-jag+jwt"})


async def fetch_id_jag(audience: str, resource: str) -> str:
    return idp_issue_id_jag("alice@example.com", audience, resource)


oauth = IdentityAssertionOAuthProvider(
    server_url="http://localhost:8001/mcp",
    storage=InMemoryTokenStorage(),
    client_id="finance-agent",
    client_secret="finance-agent-secret",
    issuer="https://auth.example.com/",
    assertion_provider=fetch_id_jag,
    scope="notes:read",
)


async def main() -> None:
    async with httpx2.AsyncClient(auth=oauth, follow_redirects=True) as http_client:
        transport = streamable_http_client("http://localhost:8001/mcp", http_client=http_client)
        async with Client(transport) as client:
            result = await client.list_tools()
            print([tool.name for tool in result.tools])

import secrets
import time

import jwt
from pydantic import AnyHttpUrl
from starlette.applications import Starlette

from mcp.server.auth.provider import (
    AccessToken,
    AuthorizationCode,
    AuthorizationParams,
    AuthorizeError,
    IdentityAssertionParams,
    OAuthAuthorizationServerProvider,
    RefreshToken,
    TokenError,
)
from mcp.server.auth.routes import create_auth_routes
from mcp.shared.auth import JWT_BEARER_GRANT_TYPE, OAuthClientInformationFull, OAuthToken

ISSUER = "https://auth.example.com/"
MCP_SERVER = "http://localhost:8001/mcp"
IDP_ISSUER = "https://idp.example.com"
IDP_SIGNING_KEY = "the-enterprise-idp-signing-key-for-this-demo"

REGISTERED_CLIENTS = {
    "finance-agent": OAuthClientInformationFull(
        client_id="finance-agent",
        client_secret="finance-agent-secret",
        redirect_uris=None,
        grant_types=[JWT_BEARER_GRANT_TYPE],
        token_endpoint_auth_method="client_secret_post",
    )
}


class EnterpriseAuthorizationServer(OAuthAuthorizationServerProvider[AuthorizationCode, RefreshToken, AccessToken]):
    def __init__(self) -> None:
        self.access_tokens: dict[str, AccessToken] = {}
        self.seen_jtis: set[str] = set()

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

    async def load_access_token(self, token: str) -> AccessToken | None:
        return self.access_tokens.get(token)

    async def exchange_identity_assertion(
        self, client: OAuthClientInformationFull, params: IdentityAssertionParams
    ) -> OAuthToken:
        try:
            header = jwt.get_unverified_header(params.assertion)
            claims = jwt.decode(
                params.assertion,
                IDP_SIGNING_KEY,
                algorithms=["HS256"],
                issuer=IDP_ISSUER,
                audience=ISSUER,
                options={"require": ["iss", "sub", "aud", "exp", "iat", "jti", "client_id", "resource", "scope"]},
            )
        except jwt.InvalidTokenError as error:
            raise TokenError("invalid_grant", "the assertion did not verify") from error
        if header.get("typ") != "oauth-id-jag+jwt":
            raise TokenError("invalid_grant", "the assertion is not an ID-JAG")
        if claims["client_id"] != client.client_id:
            raise TokenError("invalid_grant", "the assertion was issued to a different client")
        if claims["resource"] != MCP_SERVER:
            raise TokenError("invalid_target", "the assertion is for a resource this server does not serve")
        if claims["jti"] in self.seen_jtis:
            raise TokenError("invalid_grant", "the assertion has already been used")
        self.seen_jtis.add(claims["jti"])
        scopes = claims["scope"].split()
        access_token = f"mcp_{secrets.token_hex(16)}"
        self.access_tokens[access_token] = AccessToken(
            token=access_token,
            client_id=claims["client_id"],
            scopes=scopes,
            expires_at=int(time.time()) + 300,
            resource=claims["resource"],
            subject=claims["sub"],
        )
        return OAuthToken(access_token=access_token, token_type="Bearer", expires_in=300, scope=" ".join(scopes))

    async def authorize(self, client: OAuthClientInformationFull, params: AuthorizationParams) -> str:
        raise AuthorizeError("unauthorized_client", "this authorization server only accepts ID-JAGs")

    async def load_authorization_code(self, client: OAuthClientInformationFull, authorization_code: str) -> None:
        return None

    async def exchange_authorization_code(
        self, client: OAuthClientInformationFull, authorization_code: AuthorizationCode
    ) -> OAuthToken:
        raise TokenError("invalid_grant", "this authorization server only accepts ID-JAGs")

    async def load_refresh_token(self, client: OAuthClientInformationFull, refresh_token: str) -> None:
        return None

    async def exchange_refresh_token(
        self, client: OAuthClientInformationFull, refresh_token: RefreshToken, scopes: list[str]
    ) -> OAuthToken:
        raise TokenError("invalid_grant", "this authorization server only accepts ID-JAGs")


provider = EnterpriseAuthorizationServer()
auth_app = Starlette(
    routes=create_auth_routes(provider, issuer_url=AnyHttpUrl(ISSUER), identity_assertion_enabled=True)
)