76 lines
2.6 KiB
Python
76 lines
2.6 KiB
Python
from fastapi import APIRouter, Depends, HTTPException, Response, status
|
|
from sqlalchemy.orm import Session
|
|
import base64
|
|
|
|
from app.core.database.db_session import get_db
|
|
from app.core.permissions.RoleChecker import get_current_user
|
|
from app.models.UserModel import User
|
|
from app.utils.Mfa_util import (
|
|
generate_mfa_secret,
|
|
get_mfa_uri,
|
|
verify_mfa_token,
|
|
qr_code_png_base64
|
|
)
|
|
|
|
router = APIRouter(prefix="/api/v1/mfa", tags=["MFA (Two-Factor Authentication)"])
|
|
|
|
@router.post("/setup", response_class=Response)
|
|
def setup_mfa(
|
|
user_id: str = None,
|
|
current_user: User = Depends(get_current_user),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
# 1. Determine target user (defaults to self, Super Admin/Admin can specify another user_id)
|
|
target_id = user_id or current_user.user_id
|
|
if target_id != current_user.user_id and current_user.role.role_name.lower() not in ["super admin", "admin"]:
|
|
raise HTTPException(status_code=403, detail="You do not have permission to setup MFA for other users.")
|
|
|
|
# 2. Fetch target user
|
|
user = db.get(User, target_id)
|
|
if not user or user.deleted_at is not None:
|
|
raise HTTPException(status_code=404, detail="User not found")
|
|
|
|
# 3. Generate TOTP secret
|
|
secret = generate_mfa_secret()
|
|
user.mfa_secret = secret
|
|
user.mfa_enabled = False
|
|
db.commit()
|
|
|
|
# 4. Generate QR code bytes
|
|
uri = get_mfa_uri(user.email, secret)
|
|
qr_b64 = qr_code_png_base64(uri)
|
|
qr_bytes = base64.b64decode(qr_b64.split(",")[1])
|
|
|
|
return Response(
|
|
content=qr_bytes,
|
|
media_type="image/png",
|
|
headers={
|
|
"Content-Disposition": "inline; filename=mfa_qr.png"
|
|
}
|
|
)
|
|
|
|
@router.post("/verify")
|
|
def verify_mfa(
|
|
token: str,
|
|
user_id: str = None,
|
|
current_user: User = Depends(get_current_user),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
target_id = user_id or current_user.user_id
|
|
if target_id != current_user.user_id and current_user.role.role_name.lower() not in ["super admin", "admin"]:
|
|
raise HTTPException(status_code=403, detail="You do not have permission to verify MFA for other users.")
|
|
|
|
user = db.get(User, target_id)
|
|
if not user or user.deleted_at is not None:
|
|
raise HTTPException(status_code=404, detail="User not found")
|
|
|
|
if not user.mfa_secret:
|
|
raise HTTPException(status_code=400, detail="MFA setup has not been initialized. Please call /setup first.")
|
|
|
|
if not verify_mfa_token(user.mfa_secret, token):
|
|
raise HTTPException(status_code=400, detail="Invalid verification code")
|
|
|
|
user.mfa_enabled = True
|
|
db.commit()
|
|
|
|
return {"detail": "MFA enabled successfully"}
|