feat(session-verify): 添加 TOTP 支持 (#34)

* chore(deps): add pyotp * feat(auth): implement TOTP verification feat(auth): implement TOTP verification and email verification services - Added TOTP keys management with a new database model `TotpKeys`. - Introduced `EmailVerification` and `LoginSession` models for email verification. - Created `verification_service` to handle email verification logic and TOTP processes. - Updated user response models to include session verification methods. - Implemented routes for TOTP creation, verification, and fallback to email verification. - Enhanced login session management to support new location checks and verification methods. - Added migration script to create `totp_keys` table in the database. * feat(config): update config example * docs(totp): complete creating TOTP flow * refactor(totp): resolve review * feat(api): forbid unverified request * fix(totp): trace session by token id to avoid other sessions are forbidden * chore(linter): make pyright happy * fix(totp): only mark sessions with a specified token id
2025-09-21 19:50:11 +08:00
parent 7b4ff1224d
commit 1527e23b43
25 changed files with 684 additions and 235 deletions
--- a/app/router/auth.py
+++ b/app/router/auth.py
@@ -17,6 +17,7 @@ from app.auth import (
 from app.config import settings
 from app.const import BANCHOBOT_ID
 from app.database import DailyChallengeStats, OAuthClient, User
+from app.database.auth import TotpKeys
 from app.database.statistics import UserStatistics
 from app.dependencies.database import Database, get_redis
 from app.dependencies.geoip import get_client_ip, get_geoip_helper
@@ -30,12 +31,12 @@ from app.models.oauth import (
    UserRegistrationErrors,
 )
 from app.models.score import GameMode
-from app.service.email_verification_service import (
+from app.service.login_log_service import LoginLogService
+from app.service.password_reset_service import password_reset_service
+from app.service.verification_service import (
    EmailVerificationService,
    LoginSessionService,
 )
-from app.service.login_log_service import LoginLogService
-from app.service.password_reset_service import password_reset_service
 from app.utils import utcnow

 from fastapi import APIRouter, Depends, Form, Request
@@ -287,8 +288,23 @@ async def oauth_token(
        # 确保用户对象与当前会话关联
        await db.refresh(user)

-        # 获取用户信息和客户端信息
        user_id = user.id
+        totp_key: TotpKeys | None = await user.awaitable_attrs.totp_key
+
+        # 生成令牌
+        access_token_expires = timedelta(minutes=settings.access_token_expire_minutes)
+        access_token = create_access_token(data={"sub": str(user_id)}, expires_delta=access_token_expires)
+        refresh_token_str = generate_refresh_token()
+        token = await store_token(
+            db,
+            user_id,
+            client_id,
+            scopes,
+            access_token,
+            refresh_token_str,
+            settings.access_token_expire_minutes * 60,
+        )
+        token_id = token.id

        ip_address = get_client_ip(request)
        user_agent = request.headers.get("User-Agent", "")
@@ -300,15 +316,22 @@ async def oauth_token(
        # 检查是否为新位置登录
        is_new_location = await LoginSessionService.check_new_location(db, user_id, ip_address, country_code)

-        # 创建登录会话记录
-        login_session = await LoginSessionService.create_session(  # noqa: F841
-            db, redis, user_id, ip_address, user_agent, country_code, is_new_location
-        )
-
-        # 如果是新位置登录，需要邮件验证
-        if is_new_location and settings.enable_email_verification:
+        session_verification_method = None
+        if settings.enable_totp_verification and totp_key is not None:
+            session_verification_method = "totp"
+            await LoginLogService.record_login(
+                db=db,
+                user_id=user_id,
+                request=request,
+                login_success=True,
+                login_method="password_pending_verification",
+                notes="需要 TOTP 验证",
+            )
+        elif is_new_location and settings.enable_email_verification:
+            # 如果是新位置登录，需要邮件验证
            # 刷新用户对象以确保属性已加载
            await db.refresh(user)
+            session_verification_method = "mail"

            # 发送邮件验证码
            verification_sent = await EmailVerificationService.send_verification_email(
@@ -328,9 +351,9 @@ async def oauth_token(
            if not verification_sent:
                # 邮件发送失败，记录错误
                logger.error(f"[Auth] Failed to send email verification code for user {user_id}")
-        elif is_new_location and not settings.enable_email_verification:
+        elif is_new_location:
            # 新位置登录但邮件验证功能被禁用，直接标记会话为已验证
-            await LoginSessionService.mark_session_verified(db, user_id)
+            await LoginSessionService.mark_session_verified(db, redis, user_id, token_id)
            logger.debug(
                f"[Auth] New location login detected but email verification disabled, auto-verifying user {user_id}"
            )
@@ -345,25 +368,16 @@ async def oauth_token(
                notes=f"正常登录 - IP: {ip_address}, 国家: {country_code}",
            )

-        # 无论是否新位置登录，都返回正常的token
-        # session_verified状态通过/me接口的session_verified字段来体现
+        if session_verification_method:
+            await LoginSessionService.create_session(
+                db, redis, user_id, token_id, ip_address, user_agent, country_code, is_new_location, False
+            )
+            await LoginSessionService.set_login_method(user_id, token_id, session_verification_method, redis)
+        else:
+            await LoginSessionService.create_session(
+                db, redis, user_id, token_id, ip_address, user_agent, country_code, is_new_location, True
+            )

-        # 生成令牌
-        access_token_expires = timedelta(minutes=settings.access_token_expire_minutes)
-        # 获取用户ID，避免触发延迟加载
-        access_token = create_access_token(data={"sub": str(user_id)}, expires_delta=access_token_expires)
-        refresh_token_str = generate_refresh_token()
-
-        # 存储令牌
-        await store_token(
-            db,
-            user_id,
-            client_id,
-            scopes,
-            access_token,
-            refresh_token_str,
-            settings.access_token_expire_minutes * 60,
-        )
        return TokenResponse(
            access_token=access_token,
            token_type="Bearer",
--- a/app/router/notification/server.py
+++ b/app/router/notification/server.py
@@ -12,7 +12,7 @@ from app.dependencies.database import (
    get_redis,
    with_db,
 )
-from app.dependencies.user import get_current_user
+from app.dependencies.user import get_current_user_and_token
 from app.log import logger
 from app.models.chat import ChatEvent
 from app.models.notification import NotificationDetail
@@ -311,7 +311,11 @@ async def chat_websocket(
            await websocket.close(code=1008, reason="Missing authentication token")
            return

-        if (user := await get_current_user(session, SecurityScopes(scopes=["chat.read"]), token_pw=auth_token)) is None:
+        if (
+            user_and_token := await get_current_user_and_token(
+                session, SecurityScopes(scopes=["chat.read"]), token_pw=auth_token
+            )
+        ) is None:
            await websocket.close(code=1008, reason="Invalid or expired token")
            return

@@ -320,6 +324,7 @@ async def chat_websocket(
        if login.get("event") != "chat.start":
            await websocket.close(code=1008)
            return
+        user = user_and_token[0]
        user_id = user.id
        server.connect(user_id, websocket)
        # 使用明确的查询避免延迟加载
--- a/app/router/private/init.py
+++ b/app/router/private/init.py
@@ -1,8 +1,13 @@
 from __future__ import annotations

+from app.config import settings
+
 from . import avatar, beatmapset_ratings, cover, oauth, relationship, team, username  # noqa: F401
 from .router import router as private_router

+if settings.enable_totp_verification:
+    from . import totp  # noqa: F401
+
 __all__ = [
    "private_router",
 ]
--- a/app/router/private/totp.py
+++ b/app/router/private/totp.py
@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+from app.auth import (
+    check_totp_backup_code,
+    finish_create_totp_key,
+    start_create_totp_key,
+    totp_redis_key,
+    verify_totp_key,
+)
+from app.config import settings
+from app.const import BACKUP_CODE_LENGTH
+from app.database.auth import TotpKeys
+from app.database.lazer_user import User
+from app.dependencies.database import Database, get_redis
+from app.dependencies.user import get_client_user
+from app.models.totp import FinishStatus, StartCreateTotpKeyResp
+
+from .router import router
+
+from fastapi import Body, Depends, HTTPException, Security
+import pyotp
+from redis.asyncio import Redis
+
+
+@router.post(
+    "/totp/create",
+    name="开始 TOTP 创建流程",
+    description=(
+        "开始 TOTP 创建流程\n\n"
+        "返回 TOTP 密钥和 URI，供用户在身份验证器应用中添加账户。\n\n"
+        "然后将身份验证器应用提供的 TOTP 代码请求 PUT `/api/private/totp/create` 来完成 TOTP 创建流程。\n\n"
+        "若 5 分钟内未完成或错误 3 次以上则创建流程需要重新开始。"
+    ),
+    tags=["验证", "g0v0 API"],
+    response_model=StartCreateTotpKeyResp,
+    status_code=201,
+)
+async def start_create_totp(
+    redis: Redis = Depends(get_redis),
+    current_user: User = Security(get_client_user),
+):
+    if await current_user.awaitable_attrs.totp_key:
+        raise HTTPException(status_code=400, detail="TOTP is already enabled for this user")
+
+    previous = await redis.hgetall(totp_redis_key(current_user))  # pyright: ignore[reportGeneralTypeIssues]
+    if previous:  # pyright: ignore[reportGeneralTypeIssues]
+        return StartCreateTotpKeyResp(
+            secret=previous["secret"],
+            uri=pyotp.totp.TOTP(previous["secret"]).provisioning_uri(
+                name=current_user.email,
+                issuer_name=settings.totp_issuer,
+            ),
+        )
+    return await start_create_totp_key(current_user, redis)
+
+
+@router.put(
+    "/totp/create",
+    name="完成 TOTP 创建流程",
+    description=(
+        "完成 TOTP 创建流程，验证用户提供的 TOTP 代码。\n\n"
+        "- 如果验证成功，启用用户的 TOTP 双因素验证，并返回备份码。\n- 如果验证失败，返回错误信息。"
+    ),
+    tags=["验证", "g0v0 API"],
+    response_model=list[str],
+    status_code=201,
+)
+async def finish_create_totp(
+    session: Database,
+    code: str = Body(..., embed=True, description="用户提供的 TOTP 代码"),
+    redis: Redis = Depends(get_redis),
+    current_user: User = Security(get_client_user),
+):
+    status, backup_codes = await finish_create_totp_key(current_user, code, redis, session)
+    if status == FinishStatus.SUCCESS:
+        return backup_codes
+    elif status == FinishStatus.INVALID:
+        raise HTTPException(status_code=400, detail="No TOTP setup in progress or invalid data")
+    elif status == FinishStatus.TOO_MANY_ATTEMPTS:
+        raise HTTPException(status_code=400, detail="Too many failed attempts. Please start over.")
+    else:
+        raise HTTPException(status_code=400, detail="Invalid TOTP code")
+
+
+@router.delete(
+    "/totp",
+    name="禁用 TOTP 双因素验证",
+    description="禁用当前用户的 TOTP 双因素验证",
+    tags=["验证", "g0v0 API"],
+    status_code=204,
+)
+async def disable_totp(
+    session: Database,
+    code: str = Body(..., embed=True, description="用户提供的 TOTP 代码或备份码"),
+    current_user: User = Security(get_client_user),
+):
+    totp = await session.get(TotpKeys, current_user.id)
+    if not totp:
+        raise HTTPException(status_code=400, detail="TOTP is not enabled for this user")
+    if verify_totp_key(totp.secret, code) or (len(code) == BACKUP_CODE_LENGTH and check_totp_backup_code(totp, code)):
+        await session.delete(totp)
+        await session.commit()
+    else:
+        raise HTTPException(status_code=400, detail="Invalid TOTP code or backup code")
--- a/app/router/v2/me.py
+++ b/app/router/v2/me.py
@@ -1,11 +1,11 @@
 from __future__ import annotations

-from app.database import User
+from app.database import MeResp, User
 from app.database.lazer_user import ALL_INCLUDED
 from app.dependencies import get_current_user
 from app.dependencies.database import Database
+from app.dependencies.user import UserAndToken, get_current_user_and_token
 from app.exceptions.userpage import UserpageError
-from app.models.api_me import APIMe
 from app.models.score import GameMode
 from app.models.user import Page
 from app.models.userpage import (
@@ -23,7 +23,7 @@ from fastapi import HTTPException, Path, Security

@router.get(
    "/me/{ruleset}",
-    response_model=APIMe,
+    response_model=MeResp,
    name="获取当前用户信息 (指定 ruleset)",
    description="获取当前登录用户信息 （含指定 ruleset 统计）。",
    tags=["用户"],
@@ -31,34 +31,24 @@ from fastapi import HTTPException, Path, Security
 async def get_user_info_with_ruleset(
    session: Database,
    ruleset: GameMode = Path(description="指定 ruleset"),
-    current_user: User = Security(get_current_user, scopes=["identify"]),
+    user_and_token: UserAndToken = Security(get_current_user_and_token, scopes=["identify"]),
 ):
-    user_resp = await APIMe.from_db(
-        current_user,
-        session,
-        ALL_INCLUDED,
-        ruleset,
-    )
+    user_resp = await MeResp.from_db(user_and_token[0], session, ALL_INCLUDED, ruleset, token_id=user_and_token[1].id)
    return user_resp


@router.get(
    "/me/",
-    response_model=APIMe,
+    response_model=MeResp,
    name="获取当前用户信息",
    description="获取当前登录用户信息。",
    tags=["用户"],
 )
 async def get_user_info_default(
    session: Database,
-    current_user: User = Security(get_current_user, scopes=["identify"]),
+    user_and_token: UserAndToken = Security(get_current_user_and_token, scopes=["identify"]),
 ):
-    user_resp = await APIMe.from_db(
-        current_user,
-        session,
-        ALL_INCLUDED,
-        None,
-    )
+    user_resp = await MeResp.from_db(user_and_token[0], session, ALL_INCLUDED, None, token_id=user_and_token[1].id)
    return user_resp


--- a/app/router/v2/session_verify.py
+++ b/app/router/v2/session_verify.py
@@ -4,24 +4,35 @@

 from __future__ import annotations

-from typing import Annotated
+from typing import Annotated, Literal

-from app.database import User
-from app.dependencies import get_current_user
+from app.auth import check_totp_backup_code, verify_totp_key
+from app.config import settings
+from app.const import BACKUP_CODE_LENGTH
+from app.database.auth import TotpKeys
+from app.dependencies.api_version import APIVersion
 from app.dependencies.database import Database, get_redis
-from app.service.email_verification_service import (
-    EmailVerificationService,
-)
+from app.dependencies.geoip import get_client_ip
+from app.dependencies.user import UserAndToken, get_client_user_and_token
+from app.log import logger
 from app.service.login_log_service import LoginLogService
+from app.service.verification_service import (
+    EmailVerificationService,
+    LoginSessionService,
+)

 from .router import router

 from fastapi import Depends, Form, HTTPException, Request, Security, status
-from fastapi.responses import Response
+from fastapi.responses import JSONResponse, Response
 from pydantic import BaseModel
 from redis.asyncio import Redis


+class VerifyMethod(BaseModel):
+    method: Literal["totp", "mail"] = "mail"
+
+
 class SessionReissueResponse(BaseModel):
    """重新发送验证码响应"""

@@ -29,66 +40,94 @@ class SessionReissueResponse(BaseModel):
    message: str


+class VerifyFailed(Exception): ...
+
+
@router.post(
-    "/session/verify", name="验证会话", description="验证邮件验证码并完成会话认证", status_code=204, tags=["验证"]
+    "/session/verify",
+    name="验证会话",
+    description="验证邮件验证码并完成会话认证",
+    status_code=204,
+    tags=["验证"],
+    responses={
+        401: {"model": VerifyMethod, "description": "验证失败，返回当前使用的验证方法"},
+        204: {"description": "验证成功，无内容返回"},
+    },
 )
 async def verify_session(
    request: Request,
    db: Database,
+    api_version: APIVersion,
    redis: Annotated[Redis, Depends(get_redis)],
-    verification_key: str = Form(..., description="8位邮件验证码"),
-    current_user: User = Security(get_current_user),
+    verification_key: str = Form(..., description="8 位邮件验证码或者 6 位 TOTP 代码或 10 位备份码 （g0v0 扩展支持）"),
+    user_and_token: UserAndToken = Security(get_client_user_and_token),
 ) -> Response:
-    """
-    验证邮件验证码并完成会话认证
+    current_user = user_and_token[0]
+    token_id = user_and_token[1].id
+    user_id = current_user.id
+
+    if not await LoginSessionService.check_is_need_verification(db, user_id, token_id):
+        return Response(status_code=status.HTTP_204_NO_CONTENT)
+
+    verify_method: str | None = (
+        "mail" if api_version < 20250913 else await LoginSessionService.get_login_method(user_id, token_id, redis)
+    )
+
+    ip_address = get_client_ip(request)
+    user_agent = request.headers.get("User-Agent", "Unknown")
+    login_method = "password"

-    对应 osu! 的 session/verify 接口
-    成功时返回 204 No Content，失败时返回 401 Unauthorized
-    """
    try:
-        from app.dependencies.geoip import get_client_ip
+        totp_key: TotpKeys | None = await current_user.awaitable_attrs.totp_key
+        if verify_method is None:
+            verify_method = "totp" if totp_key else "mail"
+            await LoginSessionService.set_login_method(user_id, token_id, verify_method, redis)
+        login_method = verify_method

-        ip_address = get_client_ip(request)  # noqa: F841
-        user_agent = request.headers.get("User-Agent", "Unknown")  # noqa: F841
+        if verify_method == "totp":
+            if not totp_key:
+                if settings.enable_email_verification:
+                    await LoginSessionService.set_login_method(user_id, token_id, "mail", redis)
+                    await EmailVerificationService.send_verification_email(
+                        db, redis, user_id, current_user.username, current_user.email, ip_address, user_agent
+                    )
+                    verify_method = "mail"
+                    raise VerifyFailed("用户未设置 TOTP，已发送邮件验证码")
+                # 如果未开启邮箱验证，则直接认为认证通过
+                # 正常不会进入到这里

-        # 从当前认证用户获取信息
-        user_id = current_user.id
-        if not user_id:
-            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="用户未认证")
-
-        # 验证邮件验证码
-        success, message = await EmailVerificationService.verify_code(db, redis, user_id, verification_key)
-
-        if success:
-            # 记录成功的邮件验证
-            await LoginLogService.record_login(
-                db=db,
-                user_id=user_id,
-                request=request,
-                login_method="email_verification",
-                login_success=True,
-                notes="邮件验证成功",
-            )
-
-            # 返回 204 No Content 表示验证成功
-            return Response(status_code=status.HTTP_204_NO_CONTENT)
+            elif verify_totp_key(totp_key.secret, verification_key):
+                pass
+            elif len(verification_key) == BACKUP_CODE_LENGTH and check_totp_backup_code(totp_key, verification_key):
+                login_method = "totp_backup_code"
+            else:
+                raise VerifyFailed("TOTP 验证失败")
        else:
-            # 记录失败的邮件验证尝试
-            await LoginLogService.record_failed_login(
-                db=db,
-                request=request,
-                attempted_username=current_user.username,
-                login_method="email_verification",
-                notes=f"邮件验证失败: {message}",
-            )
+            success, message = await EmailVerificationService.verify_email_code(db, redis, user_id, verification_key)
+            if not success:
+                raise VerifyFailed(f"邮件验证失败: {message}")

-            # 返回 401 Unauthorized 表示验证失败
-            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail=message)
+        await LoginLogService.record_login(
+            db=db,
+            user_id=user_id,
+            request=request,
+            login_method=login_method,
+            login_success=True,
+            notes=f"{login_method} 验证成功",
+        )
+        await LoginSessionService.mark_session_verified(db, redis, user_id, token_id)
+        await db.commit()
+        return Response(status_code=status.HTTP_204_NO_CONTENT)

-    except ValueError:
-        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="无效的用户会话")
-    except Exception:
-        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="验证过程中发生错误")
+    except VerifyFailed as e:
+        await LoginLogService.record_failed_login(
+            db=db,
+            request=request,
+            attempted_username=current_user.username,
+            login_method=login_method,
+            notes=str(e),
+        )
+        return JSONResponse(status_code=status.HTTP_401_UNAUTHORIZED, content={"method": verify_method})


@router.post(
@@ -101,26 +140,27 @@ async def verify_session(
 async def reissue_verification_code(
    request: Request,
    db: Database,
+    api_version: APIVersion,
    redis: Annotated[Redis, Depends(get_redis)],
-    current_user: User = Security(get_current_user),
+    user_and_token: UserAndToken = Security(get_client_user_and_token),
 ) -> SessionReissueResponse:
-    """
-    重新发送邮件验证码
+    current_user = user_and_token[0]
+    token_id = user_and_token[1].id
+    user_id = current_user.id
+
+    if not await LoginSessionService.check_is_need_verification(db, user_id, token_id):
+        return SessionReissueResponse(success=False, message="当前会话不需要验证")
+
+    verify_method: str | None = (
+        "mail" if api_version < 20250913 else await LoginSessionService.get_login_method(user_id, token_id, redis)
+    )
+    if verify_method != "mail":
+        return SessionReissueResponse(success=False, message="当前会话不支持重新发送验证码")

-    对应 osu! 的 session/verify/reissue 接口
-    """
    try:
-        from app.dependencies.geoip import get_client_ip
-
        ip_address = get_client_ip(request)
        user_agent = request.headers.get("User-Agent", "Unknown")
-
-        # 从当前认证用户获取信息
        user_id = current_user.id
-        if not user_id:
-            return SessionReissueResponse(success=False, message="用户未认证")
-
-        # 重新发送验证码
        success, message = await EmailVerificationService.resend_verification_code(
            db,
            redis,
@@ -137,3 +177,41 @@ async def reissue_verification_code(
        return SessionReissueResponse(success=False, message="无效的用户会话")
    except Exception:
        return SessionReissueResponse(success=False, message="重新发送过程中发生错误")
+
+
+@router.post(
+    "/session/verify/mail-fallback",
+    name="邮件验证码回退",
+    description="当 TOTP 验证不可用时，使用邮件验证码进行回退验证",
+    response_model=VerifyMethod,
+    tags=["验证"],
+)
+async def fallback_email(
+    db: Database,
+    request: Request,
+    redis: Annotated[Redis, Depends(get_redis)],
+    user_and_token: UserAndToken = Security(get_client_user_and_token),
+) -> VerifyMethod:
+    current_user = user_and_token[0]
+    token_id = user_and_token[1].id
+    if not await LoginSessionService.get_login_method(current_user.id, token_id, redis):
+        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="当前会话不需要回退")
+
+    ip_address = get_client_ip(request)
+    user_agent = request.headers.get("User-Agent", "Unknown")
+
+    await LoginSessionService.set_login_method(current_user.id, token_id, "mail", redis)
+    success, message = await EmailVerificationService.resend_verification_code(
+        db,
+        redis,
+        current_user.id,
+        current_user.username,
+        current_user.email,
+        ip_address,
+        user_agent,
+    )
+    if not success:
+        logger.error(
+            f"[Email Fallback] Failed to send fallback email to user {current_user.id} (token: {token_id}): {message}"
+        )
+    return VerifyMethod()