Gan, Jimmy
42 lines
1.4 KiB
Python
42 lines
1.4 KiB
Python
"""TOTP (Time-based One-Time Password) 2FA endpoints."""
|
|
from fastapi import APIRouter, Depends, HTTPException
|
|
from pydantic import BaseModel
|
|
import pyotp
|
|
import auth_service as auth
|
|
|
|
router = APIRouter()
|
|
|
|
class Setup2FAResponse(BaseModel):
|
|
secret: str
|
|
uri: str
|
|
|
|
class Verify2FARequest(BaseModel):
|
|
secret: str
|
|
code: str
|
|
|
|
@router.post("/setup-2fa/generate", response_model=Setup2FAResponse)
|
|
async def generate_2fa(current_user = Depends(auth.get_current_user)):
|
|
"""Generate a new TOTP secret and provisioning URI for 2FA setup."""
|
|
secret = pyotp.random_base32()
|
|
uri = pyotp.totp.TOTP(secret).provisioning_uri(
|
|
name=current_user.username,
|
|
issuer_name="NAS Dashboard"
|
|
)
|
|
return {"secret": secret, "uri": uri}
|
|
|
|
@router.post("/setup-2fa/verify")
|
|
async def verify_2fa_setup(request: Verify2FARequest, current_user = Depends(auth.get_current_user)):
|
|
"""Verify TOTP code and enable 2FA for the user."""
|
|
totp = pyotp.TOTP(request.secret)
|
|
if not totp.verify(request.code):
|
|
raise HTTPException(status_code=400, detail="Invalid code")
|
|
|
|
auth.save_totp_secret(request.secret)
|
|
return {"message": "2FA enabled successfully"}
|
|
|
|
@router.post("/setup-2fa/disable")
|
|
async def disable_2fa(current_user = Depends(auth.get_current_user)):
|
|
"""Disable 2FA for the user."""
|
|
auth.save_totp_secret("")
|
|
return {"message": "2FA disabled"}
|