""" app/auth.py ─────────── Security layer: · Argon2id password hashing via passlib · HS256 JWT creation / verification (2-hour expiry) · TOTP MFA setup, verification, and activation via pyotp · Login, MFA-verify, MFA-setup, /me endpoints · get_current_user FastAPI dependency (used as auth guard on every route) · seed_admin() helper called at app startup to bootstrap the first admin """ from __future__ import annotations from datetime import datetime, timedelta from typing import Optional import pyotp from fastapi import APIRouter, Depends, HTTPException, status from fastapi.security import OAuth2PasswordBearer from jose import JWTError, jwt from passlib.context import CryptContext from sqlalchemy.orm import Session from app.config import settings from app.database import get_db from app.models import ( AdminUser, LoginRequest, MFASetupResponse, MFAVerifyRequest, TokenResponse, UserRead, ) router = APIRouter(prefix="/api/auth", tags=["auth"]) # ───────────────────────────────────────────────────────────────────────────── # Passlib CryptContext — Argon2id # ───────────────────────────────────────────────────────────────────────────── pwd_context = CryptContext(schemes=["argon2"], deprecated="auto") # ───────────────────────────────────────────────────────────────────────────── # OAuth2 Bearer token extraction # ───────────────────────────────────────────────────────────────────────────── oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/api/auth/login") # ═══════════════════════════════════════════════════════════════════════════════ # Password helpers # ═══════════════════════════════════════════════════════════════════════════════ def hash_password(plain: str) -> str: """Hash a plaintext password with Argon2id.""" return pwd_context.hash(plain) def verify_password(plain: str, hashed: str) -> bool: """Constant-time comparison using Argon2id verify.""" return pwd_context.verify(plain, hashed) # ═══════════════════════════════════════════════════════════════════════════════ # JWT helpers # ═══════════════════════════════════════════════════════════════════════════════ def create_access_token(data: dict, expires_delta: Optional[timedelta] = None) -> str: payload = data.copy() expire = datetime.utcnow() + ( expires_delta or timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES) ) payload["exp"] = expire return jwt.encode(payload, settings.SECRET_KEY, algorithm=settings.ALGORITHM) def _decode_token(token: str) -> dict: try: return jwt.decode(token, settings.SECRET_KEY, algorithms=[settings.ALGORITHM]) except JWTError: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Token is invalid or has expired.", headers={"WWW-Authenticate": "Bearer"}, ) # ═══════════════════════════════════════════════════════════════════════════════ # DB helpers # ═══════════════════════════════════════════════════════════════════════════════ def _get_user(db: Session, username: str) -> Optional[AdminUser]: return db.query(AdminUser).filter(AdminUser.username == username).first() # ═══════════════════════════════════════════════════════════════════════════════ # FastAPI dependency: authenticated current user # ═══════════════════════════════════════════════════════════════════════════════ def get_current_user( token: str = Depends(oauth2_scheme), db: Session = Depends(get_db), ) -> AdminUser: payload = _decode_token(token) username = payload.get("sub") # Reject partial MFA-pending tokens on protected routes if payload.get("mfa_pending"): raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="MFA verification required.") if not username: raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid token payload.") user = _get_user(db, username) if not user: raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="User account not found.") return user # ═══════════════════════════════════════════════════════════════════════════════ # First-boot seed # ═══════════════════════════════════════════════════════════════════════════════ def seed_admin(db: Session) -> None: """Create a default admin account if the admin_users table is empty.""" if db.query(AdminUser).first(): return admin = AdminUser( username=settings.ADMIN_USERNAME, hashed_password=hash_password(settings.ADMIN_PASSWORD), ) db.add(admin) db.commit() # ═══════════════════════════════════════════════════════════════════════════════ # Auth routes # ═══════════════════════════════════════════════════════════════════════════════ @router.post("/login", response_model=TokenResponse, summary="Authenticate with username + password") def login(payload: LoginRequest, db: Session = Depends(get_db)): user = _get_user(db, payload.username) if not user or not verify_password(payload.password, user.hashed_password): raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Incorrect username or password.", ) if user.mfa_enabled and user.totp_secret: # Issue a short-lived, MFA-pending partial token; client must call /mfa/verify partial = create_access_token( {"sub": user.username, "mfa_pending": True}, expires_delta=timedelta(minutes=5), ) return TokenResponse(access_token=partial, mfa_required=True) access_token = create_access_token({"sub": user.username}) return TokenResponse(access_token=access_token) @router.post("/mfa/verify", response_model=TokenResponse, summary="Complete MFA login with TOTP code") def verify_mfa(payload: MFAVerifyRequest, db: Session = Depends(get_db)): user = _get_user(db, payload.username) if not user or not user.totp_secret: raise HTTPException(status_code=400, detail="MFA not configured for this account.") totp = pyotp.TOTP(user.totp_secret) if not totp.verify(payload.totp_code, valid_window=1): raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid TOTP code.") access_token = create_access_token({"sub": user.username}) return TokenResponse(access_token=access_token) @router.post( "/mfa/setup", response_model=MFASetupResponse, summary="Generate a new TOTP secret and QR provisioning URI", ) def setup_mfa( current_user: AdminUser = Depends(get_current_user), db: Session = Depends(get_db), ): """ Generates a 32-character base32 TOTP secret and stores it on the account (MFA is NOT enabled yet until /mfa/enable is called with a valid code). Returns the secret and an otpauth:// URI compatible with Google Authenticator. """ secret = pyotp.random_base32(length=32) totp = pyotp.TOTP(secret) qr_uri = totp.provisioning_uri(name=current_user.username, issuer_name="PMDI") # Persist the tentative secret; activation requires /mfa/enable current_user.totp_secret = secret db.commit() return MFASetupResponse(secret=secret, qr_uri=qr_uri) @router.post("/mfa/enable", summary="Confirm TOTP code and activate MFA on the account") def enable_mfa( payload: MFAVerifyRequest, db: Session = Depends(get_db), ): user = _get_user(db, payload.username) if not user or not user.totp_secret: raise HTTPException(status_code=400, detail="Call /mfa/setup first.") totp = pyotp.TOTP(user.totp_secret) if not totp.verify(payload.totp_code, valid_window=1): raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid TOTP code.") user.mfa_enabled = True db.commit() return {"detail": "MFA enabled successfully."} @router.post("/mfa/disable", summary="Disable MFA on the account (requires valid TOTP)") def disable_mfa( payload: MFAVerifyRequest, current_user: AdminUser = Depends(get_current_user), db: Session = Depends(get_db), ): if not current_user.totp_secret: raise HTTPException(status_code=400, detail="MFA not configured.") totp = pyotp.TOTP(current_user.totp_secret) if not totp.verify(payload.totp_code, valid_window=1): raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid TOTP code.") current_user.mfa_enabled = False current_user.totp_secret = None db.commit() return {"detail": "MFA disabled."} @router.get("/me", response_model=UserRead, summary="Return the current authenticated user") def me(current_user: AdminUser = Depends(get_current_user)): return current_user