""" FastAPI MFA Integration Router Multi-factor Authentication for ATOM Chat Interface """ import base64 from datetime import datetime, timedelta import json import logging import os import secrets import time from typing import Any, Dict, List, Optional from fastapi import APIRouter, HTTPException from pydantic import BaseModel, Field logger = logging.getLogger(__name__) # MFA integration models class MFAConfig(BaseModel): """MFA configuration model""" enabled: bool = Field(True, description="Enable MFA") method: str = Field("totp", description="MFA method (totp, sms, email)") required_for_all_users: bool = Field(False, description="Require MFA for all users") backup_codes_count: int = Field(10, description="Number of backup codes") class MFAEnrollment(BaseModel): """MFA enrollment model""" user_id: str = Field(..., description="User ID") method: str = Field("totp", description="MFA method") phone_number: Optional[str] = Field(None, description="Phone number for SMS") email: Optional[str] = Field(None, description="Email for email MFA") class MFAAuthentication(BaseModel): """MFA authentication model""" user_id: str = Field(..., description="User ID") code: str = Field(..., description="MFA code") method: str = Field("totp", description="MFA method") class MFARecovery(BaseModel): """MFA recovery model""" user_id: str = Field(..., description="User ID") backup_code: str = Field(..., description="Backup code") # Create FastAPI router mfa_router = APIRouter() # Mock MFA service for demonstration class MFAService: def __init__(self): self.enabled = False self.config = MFAConfig() self.user_mfa_data = {} # user_id -> MFA data self.backup_codes = {} # user_id -> backup codes self.sessions = {} # session_id -> session data async def enable_mfa(self, config: MFAConfig) -> bool: """Enable MFA system""" try: self.config = config self.enabled = True logger.info("MFA system enabled") return True except Exception as e: logger.error(f"Failed to enable MFA: {e}") return False async def enroll_user(self, enrollment: MFAEnrollment) -> Dict[str, Any]: """Enroll user in MFA""" if not self.enabled: raise HTTPException(status_code=400, detail="MFA system not enabled") user_id = enrollment.user_id # Generate MFA data based on method if enrollment.method == "totp": # Generate TOTP secret secret = secrets.token_hex(16) qr_code_url = f"otpauth://totp/ATOM:{user_id}?secret={secret}&issuer=ATOM" self.user_mfa_data[user_id] = { "method": "totp", "secret": secret, "enrolled_at": datetime.utcnow().isoformat(), "status": "active", } elif enrollment.method == "sms": if not enrollment.phone_number: raise HTTPException( status_code=400, detail="Phone number required for SMS MFA" ) self.user_mfa_data[user_id] = { "method": "sms", "phone_number": enrollment.phone_number, "enrolled_at": datetime.utcnow().isoformat(), "status": "active", } elif enrollment.method == "email": if not enrollment.email: raise HTTPException( status_code=400, detail="Email required for email MFA" ) self.user_mfa_data[user_id] = { "method": "email", "email": enrollment.email, "enrolled_at": datetime.utcnow().isoformat(), "status": "active", } else: raise HTTPException( status_code=400, detail=f"Unsupported MFA method: {enrollment.method}" ) # Generate backup codes backup_codes = await self._generate_backup_codes(user_id) return { "user_id": user_id, "method": enrollment.method, "qr_code_url": qr_code_url if enrollment.method == "totp" else None, "backup_codes": backup_codes, "enrolled_at": datetime.utcnow().isoformat(), } async def verify_mfa(self, auth: MFAAuthentication) -> Dict[str, Any]: """Verify MFA code""" if not self.enabled: raise HTTPException(status_code=400, detail="MFA system not enabled") user_id = auth.user_id if user_id not in self.user_mfa_data: raise HTTPException(status_code=400, detail="User not enrolled in MFA") user_data = self.user_mfa_data[user_id] # Mock verification - in production, implement actual verification # For TOTP, verify against secret # For SMS/Email, verify against sent code if auth.method == "totp": # Mock TOTP verification is_valid = len(auth.code) == 6 and auth.code.isdigit() elif auth.method == "sms": # Mock SMS verification is_valid = len(auth.code) == 6 and auth.code.isdigit() elif auth.method == "email": # Mock email verification is_valid = len(auth.code) == 6 and auth.code.isdigit() else: is_valid = False if is_valid: # Create session session_id = secrets.token_urlsafe(32) self.sessions[session_id] = { "user_id": user_id, "created_at": datetime.utcnow().isoformat(), "expires_at": (datetime.utcnow() + timedelta(hours=24)).isoformat(), } return { "success": True, "session_id": session_id, "message": "MFA verification successful", } else: return {"success": False, "message": "Invalid MFA code"} async def verify_backup_code(self, recovery: MFARecovery) -> Dict[str, Any]: """Verify backup code""" user_id = recovery.user_id if user_id not in self.backup_codes: raise HTTPException( status_code=400, detail="No backup codes found for user" ) backup_codes = self.backup_codes[user_id] # Check if backup code is valid and not used for code_data in backup_codes: if code_data["code"] == recovery.backup_code and not code_data["used"]: # Mark code as used code_data["used"] = True code_data["used_at"] = datetime.utcnow().isoformat() # Create session session_id = secrets.token_urlsafe(32) self.sessions[session_id] = { "user_id": user_id, "created_at": datetime.utcnow().isoformat(), "expires_at": (datetime.utcnow() + timedelta(hours=24)).isoformat(), "via_backup_code": True, } return { "success": True, "session_id": session_id, "message": "Backup code verification successful", } return {"success": False, "message": "Invalid or already used backup code"} async def generate_new_backup_codes(self, user_id: str) -> List[str]: """Generate new backup codes for user""" backup_codes = await self._generate_backup_codes(user_id) return backup_codes async def get_user_mfa_status(self, user_id: str) -> Dict[str, Any]: """Get MFA status for user""" if user_id in self.user_mfa_data: user_data = self.user_mfa_data[user_id] return { "enrolled": True, "method": user_data["method"], "enrolled_at": user_data["enrolled_at"], "status": user_data["status"], } else: return { "enrolled": False, "method": None, "enrolled_at": None, "status": "not_enrolled", } async def _generate_backup_codes(self, user_id: str) -> List[str]: """Generate backup codes for user""" backup_codes = [] for i in range(self.config.backup_codes_count): code = f"{secrets.randbelow(10000):04d}-{secrets.randbelow(10000):04d}" backup_codes.append( { "code": code, "used": False, "generated_at": datetime.utcnow().isoformat(), } ) self.backup_codes[user_id] = backup_codes return [code["code"] for code in backup_codes] # Initialize MFA service mfa_service = MFAService() # MFA API endpoints @mfa_router.post("/mfa/enable") async def enable_mfa(config: MFAConfig): """Enable MFA system""" try: success = await mfa_service.enable_mfa(config) if success: return { "success": True, "message": "MFA system enabled successfully", "config": config.dict(), } else: raise HTTPException(status_code=500, detail="Failed to enable MFA system") except Exception as e: raise HTTPException(status_code=500, detail=f"Failed to enable MFA: {str(e)}") @mfa_router.post("/mfa/enroll") async def enroll_user(enrollment: MFAEnrollment): """Enroll user in MFA""" try: result = await mfa_service.enroll_user(enrollment) return { "success": True, "message": "User enrolled in MFA successfully", "data": result, } except Exception as e: raise HTTPException(status_code=500, detail=f"Failed to enroll user: {str(e)}") @mfa_router.post("/mfa/verify") async def verify_mfa(auth: MFAAuthentication): """Verify MFA code""" try: result = await mfa_service.verify_mfa(auth) return result except Exception as e: raise HTTPException( status_code=500, detail=f"MFA verification failed: {str(e)}" ) @mfa_router.post("/mfa/recover") async def recover_with_backup_code(recovery: MFARecovery): """Recover access using backup code""" try: result = await mfa_service.verify_backup_code(recovery) return result except Exception as e: raise HTTPException(status_code=500, detail=f"Recovery failed: {str(e)}") @mfa_router.post("/mfa/users/{user_id}/backup-codes") async def generate_backup_codes(user_id: str): """Generate new backup codes for user""" try: backup_codes = await mfa_service.generate_new_backup_codes(user_id) return { "success": True, "user_id": user_id, "backup_codes": backup_codes, "message": "New backup codes generated successfully", } except Exception as e: raise HTTPException( status_code=500, detail=f"Failed to generate backup codes: {str(e)}" ) @mfa_router.get("/mfa/users/{user_id}/status") async def get_user_mfa_status(user_id: str): """Get MFA status for user""" try: status = await mfa_service.get_user_mfa_status(user_id) return {"user_id": user_id, "status": status} except Exception as e: raise HTTPException( status_code=500, detail=f"Failed to get MFA status: {str(e)}" ) @mfa_router.get("/mfa/health") async def mfa_health_check(): """MFA system health check""" return { "status": "healthy" if mfa_service.enabled else "disabled", "service": "mfa_system", "enabled": mfa_service.enabled, "config": mfa_service.config.dict() if mfa_service.enabled else None, "enrolled_users": len(mfa_service.user_mfa_data), "active_sessions": len(mfa_service.sessions), "timestamp": datetime.utcnow().isoformat(), } @mfa_router.get("/mfa/stats") async def get_mfa_stats(): """Get MFA system statistics""" try: enrolled_users = len(mfa_service.user_mfa_data) active_sessions = len(mfa_service.sessions) # Calculate method distribution method_distribution = {} for user_data in mfa_service.user_mfa_data.values(): method = user_data["method"] method_distribution[method] = method_distribution.get(method, 0) + 1 return { "stats": { "enrolled_users": enrolled_users, "active_sessions": active_sessions, "method_distribution": method_distribution, "backup_codes_generated": sum( len(codes) for codes in mfa_service.backup_codes.values() ), }, "timestamp": datetime.utcnow().isoformat(), } except Exception as e: raise HTTPException( status_code=500, detail=f"Failed to get MFA stats: {str(e)}" ) logger.info("MFA FastAPI router initialized") # Export router for main application integration router = mfa_router