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/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()