pmdi/backend/app/auth.py

"""
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