67 lines
2.0 KiB
Python
67 lines
2.0 KiB
Python
from asgiref.sync import sync_to_async
|
||
from rest_framework.exceptions import APIException
|
||
from rest_framework_simplejwt.authentication import JWTAuthentication
|
||
from rest_framework_simplejwt.exceptions import TokenError
|
||
|
||
from mcp.server.auth.provider import AccessToken
|
||
from mcp.server.auth.middleware.bearer_auth import RequireAuthMiddleware
|
||
from starlette.types import ASGIApp, Receive, Scope, Send
|
||
|
||
|
||
def verify_factory_jwt(token: str) -> AccessToken | None:
|
||
"""验证 Factory access token，并生成 MCP 的逐请求身份信息。"""
|
||
authentication = JWTAuthentication()
|
||
try:
|
||
validated_token = authentication.get_validated_token(token)
|
||
user = authentication.get_user(validated_token)
|
||
except (APIException, TokenError):
|
||
return None
|
||
|
||
return AccessToken(
|
||
token=token,
|
||
client_id="factory-mcp",
|
||
scopes=["factory:user"],
|
||
expires_at=validated_token.get("exp"),
|
||
subject=str(user.pk),
|
||
claims={
|
||
"factory_user": {
|
||
"id": str(user.pk),
|
||
"username": user.get_username(),
|
||
"name": user.name,
|
||
"is_superuser": user.is_superuser,
|
||
},
|
||
},
|
||
)
|
||
|
||
|
||
class FactoryJWTVerifier:
|
||
"""让 MCP SDK 复用 Factory SimpleJWT 的验证规则。"""
|
||
|
||
async def verify_token(self, token: str) -> AccessToken | None:
|
||
return await sync_to_async(
|
||
verify_factory_jwt,
|
||
thread_sensitive=True,
|
||
)(token)
|
||
|
||
|
||
class RequireFactoryJWTMiddleware:
|
||
"""仅保护 HTTP 请求，并把 ASGI lifespan 原样交给 MCP SDK。"""
|
||
|
||
def __init__(self, app: ASGIApp):
|
||
self.app = app
|
||
self.protected_app = RequireAuthMiddleware(
|
||
app,
|
||
required_scopes=[],
|
||
)
|
||
|
||
async def __call__(
|
||
self,
|
||
scope: Scope,
|
||
receive: Receive,
|
||
send: Send,
|
||
) -> None:
|
||
if scope["type"] != "http":
|
||
await self.app(scope, receive, send)
|
||
return
|
||
await self.protected_app(scope, receive, send)
|