312 lines
10 KiB
Python
312 lines
10 KiB
Python
"""
|
||
会话验证中间件和状态管理
|
||
|
||
基于osu-web的会话验证系统实现
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from collections.abc import Awaitable, Callable
|
||
from typing import ClassVar, Literal, cast
|
||
|
||
from app.database.lazer_user import User
|
||
from app.database.verification import LoginSession
|
||
from app.dependencies.database import get_redis, with_db
|
||
from app.log import logger
|
||
from app.service.verification_service import LoginSessionService
|
||
from app.utils import bg_tasks
|
||
|
||
from fastapi import HTTPException, Request, Response, status
|
||
from fastapi.responses import JSONResponse
|
||
from redis.asyncio import Redis
|
||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||
|
||
|
||
class SessionVerificationState:
|
||
"""会话验证状态管理类
|
||
|
||
参考osu-web的State类实现
|
||
"""
|
||
|
||
def __init__(self, session: LoginSession, user: User, redis: Redis):
|
||
self.session = session
|
||
self.user = user
|
||
self.redis = redis
|
||
|
||
@classmethod
|
||
async def get_current(
|
||
cls,
|
||
request: Request,
|
||
db: AsyncSession,
|
||
redis: Redis,
|
||
user: User,
|
||
) -> SessionVerificationState | None:
|
||
"""获取当前会话验证状态"""
|
||
try:
|
||
# 从请求头或token中获取会话信息
|
||
session_token = cls._extract_session_token(request)
|
||
if not session_token:
|
||
return None
|
||
|
||
# 查找会话
|
||
session = await LoginSessionService.find_for_verification(db, session_token)
|
||
if not session or session.user_id != user.id:
|
||
return None
|
||
|
||
return cls(session, user, redis)
|
||
except Exception as e:
|
||
logger.error(f"[Session Verification] Error getting current state: {e}")
|
||
return None
|
||
|
||
@staticmethod
|
||
def _extract_session_token(request: Request) -> str | None:
|
||
"""从请求中提取会话token"""
|
||
# 尝试从Authorization header提取
|
||
auth_header = request.headers.get("Authorization", "")
|
||
if auth_header.startswith("Bearer "):
|
||
return auth_header[7:] # 移除"Bearer "前缀
|
||
|
||
# 可以扩展其他提取方式
|
||
return None
|
||
|
||
def get_method(self) -> str:
|
||
"""获取验证方法
|
||
|
||
参考osu-web的逻辑，智能选择验证方法
|
||
"""
|
||
current_method = self.session.verification_method
|
||
|
||
if current_method is None:
|
||
# 智能选择验证方法
|
||
# 参考osu-web: API版本 < 20250913 或用户没有TOTP时使用邮件验证
|
||
# 这里简化为检查用户是否有TOTP
|
||
totp_key = getattr(self.user, "totp_key", None)
|
||
current_method = "totp" if totp_key else "mail"
|
||
|
||
# 设置验证方法
|
||
bg_tasks.add_task(self._set_verification_method, current_method)
|
||
|
||
return current_method
|
||
|
||
async def _set_verification_method(self, method: str) -> None:
|
||
"""内部方法：设置验证方法"""
|
||
try:
|
||
token_id = self.session.token_id
|
||
if token_id is not None and method in ["totp", "mail"]:
|
||
# 类型检查确保method是正确的字面量类型
|
||
verification_method = method if method in ["totp", "mail"] else "totp"
|
||
await LoginSessionService.set_login_method(
|
||
self.user.id,
|
||
token_id,
|
||
cast(Literal["totp", "mail"], verification_method),
|
||
self.redis,
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"[Session Verification] Error setting verification method: {e}")
|
||
|
||
def is_verified(self) -> bool:
|
||
"""检查会话是否已验证"""
|
||
return self.session.is_verified
|
||
|
||
async def mark_verified(self) -> None:
|
||
"""标记会话为已验证"""
|
||
try:
|
||
# 创建专用数据库会话
|
||
db = with_db()
|
||
try:
|
||
token_id = self.session.token_id
|
||
if token_id is not None:
|
||
await LoginSessionService.mark_session_verified(db, self.redis, self.user.id, token_id)
|
||
finally:
|
||
await db.close()
|
||
except Exception as e:
|
||
logger.error(f"[Session Verification] Error marking session verified: {e}")
|
||
|
||
def get_key(self) -> str:
|
||
"""获取会话密钥"""
|
||
return str(self.session.id) if self.session.id else ""
|
||
|
||
def get_key_for_event(self) -> str:
|
||
"""获取用于事件广播的会话密钥"""
|
||
return LoginSessionService.get_key_for_event(self.get_key())
|
||
|
||
def user_id(self) -> int:
|
||
"""获取用户ID"""
|
||
return self.user.id
|
||
|
||
async def issue_mail_if_needed(self) -> None:
|
||
"""如果需要，发送验证邮件"""
|
||
try:
|
||
if self.get_method() == "mail":
|
||
from app.service.verification_service import EmailVerificationService
|
||
|
||
# 创建专用数据库会话发送邮件
|
||
db = with_db()
|
||
try:
|
||
await EmailVerificationService.send_verification_email(
|
||
db, self.redis, self.user.id, self.user.username, self.user.email, None, None
|
||
)
|
||
finally:
|
||
await db.close()
|
||
except Exception as e:
|
||
logger.error(f"[Session Verification] Error issuing mail: {e}")
|
||
|
||
|
||
class SessionVerificationController:
|
||
"""会话验证控制器
|
||
|
||
参考osu-web的Controller类实现
|
||
"""
|
||
|
||
# 需要跳过验证的路由（参考osu-web的SKIP_VERIFICATION_ROUTES）
|
||
SKIP_VERIFICATION_ROUTES: ClassVar[set[str]] = {
|
||
"/api/v2/session/verify",
|
||
"/api/v2/session/verify/reissue",
|
||
"/api/v2/me",
|
||
"/api/v2/logout",
|
||
"/oauth/token",
|
||
}
|
||
|
||
@staticmethod
|
||
def should_skip_verification(request: Request) -> bool:
|
||
"""检查是否应该跳过验证"""
|
||
path = request.url.path
|
||
return path in SessionVerificationController.SKIP_VERIFICATION_ROUTES
|
||
|
||
@staticmethod
|
||
async def initiate_verification(
|
||
state: SessionVerificationState,
|
||
request: Request,
|
||
) -> Response:
|
||
"""启动会话验证流程
|
||
|
||
参考osu-web的initiate方法
|
||
"""
|
||
try:
|
||
method = state.get_method()
|
||
|
||
# 如果是邮件验证，发送验证邮件
|
||
if method == "mail":
|
||
await state.issue_mail_if_needed()
|
||
|
||
# API请求返回JSON响应
|
||
if request.url.path.startswith("/api/"):
|
||
return JSONResponse(status_code=status.HTTP_401_UNAUTHORIZED, content={"method": method})
|
||
|
||
# 其他情况可以扩展支持HTML响应
|
||
return JSONResponse(
|
||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||
content={"authentication": "verify", "method": method, "message": "Session verification required"},
|
||
)
|
||
|
||
except Exception as e:
|
||
logger.error(f"[Session Verification] Error initiating verification: {e}")
|
||
raise HTTPException(
|
||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Verification initiation failed"
|
||
)
|
||
|
||
|
||
class SessionVerificationMiddleware:
|
||
"""会话验证中间件
|
||
|
||
参考osu-web的VerifyUser中间件实现
|
||
"""
|
||
|
||
def __init__(self, app: Callable[[Request], Awaitable[Response]]):
|
||
self.app = app
|
||
|
||
async def __call__(self, request: Request, call_next: Callable[[Request], Awaitable[Response]]) -> Response:
|
||
"""中间件主要逻辑"""
|
||
try:
|
||
# 检查是否需要跳过验证
|
||
if SessionVerificationController.should_skip_verification(request):
|
||
return await call_next(request)
|
||
|
||
# 获取依赖项
|
||
user = await self._get_user(request)
|
||
if not user:
|
||
# 未认证用户跳过验证
|
||
return await call_next(request)
|
||
|
||
# 获取数据库和Redis连接
|
||
db = await self._get_db()
|
||
redis = await self._get_redis()
|
||
|
||
# 获取会话验证状态
|
||
state = await SessionVerificationState.get_current(request, db, redis, user)
|
||
if not state:
|
||
# 无法获取会话状态，继续请求
|
||
return await call_next(request)
|
||
|
||
# 检查是否已验证
|
||
if state.is_verified():
|
||
# 已验证，继续请求
|
||
return await call_next(request)
|
||
|
||
# 检查是否需要验证
|
||
if not self._requires_verification(request):
|
||
return await call_next(request)
|
||
|
||
# 启动验证流程
|
||
return await SessionVerificationController.initiate_verification(state, request)
|
||
|
||
except Exception as e:
|
||
logger.error(f"[Session Verification Middleware] Unexpected error: {e}")
|
||
# 出错时允许请求继续，避免阻塞正常流程
|
||
return await call_next(request)
|
||
|
||
async def _get_user(self, request: Request) -> User | None:
|
||
"""获取当前用户"""
|
||
try:
|
||
# 这里需要手动获取用户，因为在中间件中无法直接使用依赖注入
|
||
# 简化实现，实际应该从token中解析用户
|
||
return None # 暂时返回None，需要实际实现
|
||
except Exception:
|
||
return None
|
||
|
||
async def _get_db(self) -> AsyncSession:
|
||
"""获取数据库连接"""
|
||
return with_db()
|
||
|
||
async def _get_redis(self) -> Redis:
|
||
"""获取Redis连接"""
|
||
return get_redis()
|
||
|
||
def _requires_verification(self, request: Request) -> bool:
|
||
"""检查是否需要验证
|
||
|
||
参考osu-web的requiresVerification方法
|
||
"""
|
||
method = request.method
|
||
|
||
# GET/HEAD/OPTIONS请求一般不需要验证
|
||
safe_methods = {"GET", "HEAD", "OPTIONS"}
|
||
if method in safe_methods:
|
||
return False
|
||
|
||
# POST/PUT/DELETE等修改操作需要验证
|
||
return True
|
||
|
||
|
||
# FastAPI中间件包装器
|
||
class FastAPISessionVerificationMiddleware:
|
||
"""FastAPI会话验证中间件包装器"""
|
||
|
||
def __init__(self, app):
|
||
self.app = app
|
||
|
||
async def __call__(self, scope, receive, send):
|
||
if scope["type"] != "http":
|
||
return await self.app(scope, receive, send)
|
||
|
||
request = Request(scope, receive)
|
||
|
||
async def call_next(req: Request) -> Response:
|
||
# 这里需要调用FastAPI应用
|
||
return Response("OK") # 占位符实现
|
||
|
||
middleware = SessionVerificationMiddleware(call_next)
|
||
response = await middleware(request, call_next)
|
||
|
||
await response(scope, receive, send)
|