Gan, Jimmy
b981c06d59
- Fix unit test imports: add env setup in conftest.py before module imports - Add 24 new auth router tests (RBAC, preferences, password validation) - Add 16 new tests for litellm and chat_summary routers - Apply black formatting and ruff linting across codebase - Add pre-commit hooks configuration (black, ruff, file checks) - Increase CI coverage threshold from 40% to 50% Test Results: - 206 tests passing (91 unit + 115 integration) - Coverage: 58.79% on core modules - auth.py: 57% → 85%, litellm.py: 23% → 87%, chat_summary.py: 41% → 100% - auth_service: 96.51%, config: 100%, rbac: 93.48%
46 lines
1.3 KiB
Python
46 lines
1.3 KiB
Python
"""TOTP (Time-based One-Time Password) 2FA endpoints."""
|
|
|
|
import pyotp
|
|
from fastapi import APIRouter, Depends, HTTPException
|
|
from pydantic import BaseModel
|
|
|
|
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"}
|