Add TOTP for Admin login

src/routers/admins.py

@@ -1,12 +1,17 @@
+import base64
+import io
 from typing import Optional
 
+import pyotp
+import qrcode
 from fastapi import APIRouter, Depends, HTTPException
+from pydantic import BaseModel
 from sqlalchemy.orm import Session
 
-from core.auth import hash_password
+from core.auth import decrypt_secret, encrypt_secret, hash_password
 from core.database import AdminUser, get_db
 from core.dependencies import require_admin
-from core.schemas import AdminUserCreate, AdminUserResponse, AdminPasswordChange
+from core.schemas import AdminPasswordChange, AdminUserCreate, AdminUserResponse, TotpSetupResponse
 
 router = APIRouter(prefix="/api/admin/admins", tags=["admin-admins"])
 
@@ -15,7 +20,7 @@ router = APIRouter(prefix="/api/admin/admins", tags=["admin-admins"])
 async def list_admins(
     db: Session = Depends(get_db), _: dict = Depends(require_admin)
 ):
-    return [AdminUserResponse(id=u.id, username=u.username, role=u.role) for u in db.query(AdminUser).order_by(AdminUser.id).all()]
+    return [AdminUserResponse(id=u.id, username=u.username, role=u.role, totp_enabled=u.totp_enabled) for u in db.query(AdminUser).order_by(AdminUser.id).all()]
 
 
 @router.post("", response_model=AdminUserResponse, status_code=201)
@@ -73,3 +78,70 @@ async def change_password(
         raise HTTPException(404, "Admin not found")
     user.password_hash = hash_password(req.new_password)
     db.commit()
+
+
+# ── TOTP (2FA) ────────────────────────────────────────────────────────────────
+
+@router.post("/{username}/totp/setup", response_model=TotpSetupResponse)
+async def totp_setup(
+    username: str,
+    db: Session = Depends(get_db),
+    caller: dict = Depends(require_admin),
+):
+    """Generate a new TOTP secret and return the provisioning URI for QR scanning.
+    Does NOT enable 2FA yet — call /enable with a valid code to activate."""
+    if caller["sub"] != username:
+        raise HTTPException(403, "You can only configure 2FA for your own account")
+    user: Optional[AdminUser] = db.query(AdminUser).filter_by(username=username).first()
+    if not user:
+        raise HTTPException(404, "Admin not found")
+    secret = pyotp.random_base32()
+    user.totp_secret_enc = encrypt_secret(secret)
+    user.totp_enabled = False   # pending confirmation
+    db.commit()
+    uri = pyotp.TOTP(secret).provisioning_uri(name=username, issuer_name="Lagomare Gates")
+    buf = io.BytesIO()
+    qrcode.make(uri).save(buf)
+    qr_b64 = base64.b64encode(buf.getvalue()).decode()
+    return TotpSetupResponse(provisioning_uri=uri, qr_image_b64=qr_b64)
+
+
+class TotpConfirm(BaseModel):
+    otp_code: str
+
+
+@router.post("/{username}/totp/enable", status_code=204)
+async def totp_enable(
+    username: str,
+    req: TotpConfirm,
+    db: Session = Depends(get_db),
+    caller: dict = Depends(require_admin),
+):
+    """Verify the OTP code and activate 2FA for the account."""
+    if caller["sub"] != username:
+        raise HTTPException(403, "You can only configure 2FA for your own account")
+    user: Optional[AdminUser] = db.query(AdminUser).filter_by(username=username).first()
+    if not user or not user.totp_secret_enc:
+        raise HTTPException(422, "Run /setup first to generate a secret")
+    secret = decrypt_secret(user.totp_secret_enc)
+    if not pyotp.TOTP(secret).verify(req.otp_code.strip(), valid_window=1):
+        raise HTTPException(422, "Invalid OTP code — make sure your authenticator is in sync")
+    user.totp_enabled = True
+    db.commit()
+
+
+@router.delete("/{username}/totp", status_code=204)
+async def totp_disable(
+    username: str,
+    db: Session = Depends(get_db),
+    caller: dict = Depends(require_admin),
+):
+    """Disable 2FA and discard the stored secret."""
+    if caller["sub"] != username:
+        raise HTTPException(403, "You can only configure 2FA for your own account")
+    user: Optional[AdminUser] = db.query(AdminUser).filter_by(username=username).first()
+    if not user:
+        raise HTTPException(404, "Admin not found")
+    user.totp_enabled = False
+    user.totp_secret_enc = None
+    db.commit()