a0c57e0f21
- 修复登录接口 response_model 导致 2FA 返回格式验证失败(500错误) - 修复 showProfile 未调用 loadTwofaStatus 导致 2FA 状态一直显示加载中 - 修复 hideOthers 默认值为 false 确保管理员可查看所有数据 - 删除重复的 doLogin 函数 - 删除残留的代码片段修复 JavaScript 语法错误
154 lines
4.4 KiB
Python
154 lines
4.4 KiB
Python
from fastapi import APIRouter, Depends, HTTPException
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
from sqlalchemy import select
|
||
from pydantic import BaseModel
|
||
import pyotp
|
||
import qrcode
|
||
import io
|
||
import base64
|
||
from app.core.database import get_db
|
||
from app.core.deps import get_current_user
|
||
from app.models.models import User
|
||
|
||
router = APIRouter(prefix="/auth/2fa", tags=["双因素认证"])
|
||
|
||
|
||
class TwoFASetupResponse(BaseModel):
|
||
secret: str
|
||
qr_code: str # base64 encoded QR code image
|
||
otpauth_url: str
|
||
|
||
|
||
class TwoFAVerifyRequest(BaseModel):
|
||
code: str
|
||
|
||
|
||
class TwoFALoginRequest(BaseModel):
|
||
temp_token: str
|
||
code: str
|
||
|
||
|
||
@router.post("/setup", response_model=TwoFASetupResponse)
|
||
async def setup_2fa(
|
||
user: User = Depends(get_current_user),
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
"""生成 2FA 密钥和二维码"""
|
||
secret = pyotp.random_base32()
|
||
|
||
totp = pyotp.TOTP(secret)
|
||
otpauth_url = totp.provisioning_uri(
|
||
name=user.email,
|
||
issuer_name="债务管理系统"
|
||
)
|
||
|
||
# 生成 QR 码
|
||
qr = qrcode.QRCode(version=1, box_size=10, border=5)
|
||
qr.add_data(otpauth_url)
|
||
qr.make(fit=True)
|
||
img = qr.make_image(fill_color="black", back_color="white")
|
||
|
||
# 转换为 base64
|
||
buffer = io.BytesIO()
|
||
img.save(buffer, format="PNG")
|
||
qr_base64 = base64.b64encode(buffer.getvalue()).decode()
|
||
|
||
# 临时存储 secret（验证后才会真正保存）
|
||
user.two_factor_secret = secret
|
||
await db.commit()
|
||
|
||
return TwoFASetupResponse(
|
||
secret=secret,
|
||
qr_code=f"data:image/png;base64,{qr_base64}",
|
||
otpauth_url=otpauth_url
|
||
)
|
||
|
||
|
||
@router.post("/verify")
|
||
async def verify_2fa_setup(
|
||
req: TwoFAVerifyRequest,
|
||
user: User = Depends(get_current_user),
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
"""验证 2FA 设置（首次验证后启用）"""
|
||
if not user.two_factor_secret:
|
||
raise HTTPException(status_code=400, detail="请先执行 2FA 设置")
|
||
|
||
totp = pyotp.TOTP(user.two_factor_secret)
|
||
if not totp.verify(req.code, valid_window=1):
|
||
raise HTTPException(status_code=400, detail="验证码错误")
|
||
|
||
# 启用 2FA
|
||
user.two_factor_enabled = 1
|
||
await db.commit()
|
||
|
||
return {"message": "双因素认证已启用"}
|
||
|
||
|
||
@router.post("/disable")
|
||
async def disable_2fa(
|
||
req: TwoFAVerifyRequest,
|
||
user: User = Depends(get_current_user),
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
"""禁用 2FA（需要验证当前验证码）"""
|
||
if not user.two_factor_enabled:
|
||
raise HTTPException(status_code=400, detail="未启用双因素认证")
|
||
|
||
if not user.two_factor_secret:
|
||
raise HTTPException(status_code=400, detail="2FA 配置异常")
|
||
|
||
totp = pyotp.TOTP(user.two_factor_secret)
|
||
if not totp.verify(req.code, valid_window=1):
|
||
raise HTTPException(status_code=400, detail="验证码错误")
|
||
|
||
user.two_factor_secret = None
|
||
user.two_factor_enabled = 0
|
||
await db.commit()
|
||
|
||
return {"message": "双因素认证已禁用"}
|
||
|
||
|
||
@router.post("/verify-login")
|
||
async def verify_2fa_login(
|
||
req: TwoFALoginRequest,
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
"""验证 2FA 登录（密码验证后，2FA 验证前）"""
|
||
from app.core.security import decode_token, create_access_token, create_refresh_token
|
||
|
||
# 验证临时 token
|
||
payload = decode_token(req.temp_token)
|
||
if not payload or payload.get("type") != "2fa_pending":
|
||
raise HTTPException(status_code=401, detail="无效的验证请求")
|
||
|
||
user_id = int(payload["sub"])
|
||
result = await db.execute(select(User).where(User.id == user_id))
|
||
user = result.scalar_one_or_none()
|
||
|
||
if not user or not user.two_factor_secret:
|
||
raise HTTPException(status_code=401, detail="用户状态异常")
|
||
|
||
# 验证 2FA
|
||
totp = pyotp.TOTP(user.two_factor_secret)
|
||
if not totp.verify(req.code, valid_window=1):
|
||
raise HTTPException(status_code=401, detail="验证码错误")
|
||
|
||
# 返回正式 token
|
||
return {
|
||
"access_token": create_access_token(user.id),
|
||
"refresh_token": create_refresh_token(user.id),
|
||
"token_type": "bearer"
|
||
}
|
||
|
||
|
||
@router.get("/status")
|
||
async def get_2fa_status(
|
||
user: User = Depends(get_current_user),
|
||
):
|
||
"""获取 2FA 状态"""
|
||
return {
|
||
"enabled": bool(user.two_factor_enabled),
|
||
"has_secret": bool(user.two_factor_secret)
|
||
}
|