""" Time-based One-Time Password (TOTP) implementation for MFA """ import base64 import hashlib import hmac import time import qrcode import io from typing import Optional, Tuple from django.conf import settings from django.contrib.auth import get_user_model from django.core.exceptions import ValidationError from ..models import MFADevice, AuditLog User = get_user_model() class TOTPGenerator: """TOTP generator and validator""" def __init__(self, secret: str, time_step: int = 30, digits: int = 6): self.secret = secret self.time_step = time_step self.digits = digits def generate_token(self, timestamp: Optional[int] = None) -> str: """Generate TOTP token for given timestamp""" if timestamp is None: timestamp = int(time.time()) time_counter = timestamp // self.time_step # Convert secret to bytes secret_bytes = base64.b32decode(self.secret.upper() + '=' * (8 - len(self.secret) % 8)) # Generate HMAC-SHA1 time_bytes = time_counter.to_bytes(8, byteorder='big') hmac_digest = hmac.new(secret_bytes, time_bytes, hashlib.sha1).digest() # Extract dynamic binary code offset = hmac_digest[-1] & 0x0f code = int.from_bytes(hmac_digest[offset:offset+4], byteorder='big') & 0x7fffffff # Convert to string with leading zeros return str(code % (10 ** self.digits)).zfill(self.digits) def verify_token(self, token: str, window: int = 1) -> bool: """Verify TOTP token with time window tolerance""" current_time = int(time.time()) for i in range(-window, window + 1): expected_token = self.generate_token(current_time + i * self.time_step) if token == expected_token: return True return False class MFAProvider: """MFA provider for managing TOTP devices""" @staticmethod def generate_secret() -> str: """Generate a new TOTP secret""" import secrets return base64.b32encode(secrets.token_bytes(20)).decode('ascii') @staticmethod def create_totp_device(user: User, device_name: str) -> Tuple[MFADevice, str]: """Create a new TOTP device for user""" secret = MFAProvider.generate_secret() device = MFADevice.objects.create( user=user, device_type='TOTP', name=device_name, secret_key=secret, is_active=True ) # Generate QR code data qr_data = MFAProvider.generate_qr_code_data(user, secret) return device, qr_data @staticmethod def generate_qr_code_data(user: User, secret: str) -> str: """Generate QR code data for TOTP setup""" issuer = getattr(settings, 'MFA_ISSUER_NAME', 'ETB Incident Management') account_name = f"{user.username}@{issuer}" # TOTP URI format uri = f"otpauth://totp/{account_name}?secret={secret}&issuer={issuer}" return uri @staticmethod def generate_qr_code_image(qr_data: str) -> bytes: """Generate QR code image as bytes""" qr = qrcode.QRCode( version=1, error_correction=qrcode.constants.ERROR_CORRECT_L, box_size=10, border=4, ) qr.add_data(qr_data) qr.make(fit=True) img = qr.make_image(fill_color="black", back_color="white") # Convert to bytes img_buffer = io.BytesIO() img.save(img_buffer, format='PNG') img_buffer.seek(0) return img_buffer.getvalue() @staticmethod def verify_totp_token(user: User, token: str, device_id: Optional[str] = None) -> bool: """Verify TOTP token for user""" try: # Get user's TOTP devices devices = MFADevice.objects.filter( user=user, device_type='TOTP', is_active=True ) if device_id: devices = devices.filter(id=device_id) for device in devices: totp = TOTPGenerator(device.secret_key) if totp.verify_token(token): # Update last used timestamp device.last_used = timezone.now() device.save(update_fields=['last_used']) # Log successful MFA verification AuditLog.objects.create( user=user, action_type='MFA_VERIFIED', details={ 'device_id': str(device.id), 'device_name': device.name, 'device_type': device.device_type } ) return True # Log failed MFA attempt AuditLog.objects.create( user=user, action_type='MFA_FAILED', severity='MEDIUM', details={ 'device_id': device_id, 'token_provided': bool(token) } ) return False except Exception as e: # Log MFA error AuditLog.objects.create( user=user, action_type='MFA_ERROR', severity='HIGH', details={'error': str(e)} ) return False @staticmethod def enable_mfa_for_user(user: User) -> bool: """Enable MFA for user if they have at least one active device""" active_devices = MFADevice.objects.filter( user=user, is_active=True ).count() if active_devices > 0: user.mfa_enabled = True user.save(update_fields=['mfa_enabled']) # Log MFA enablement AuditLog.objects.create( user=user, action_type='MFA_ENABLED', details={'device_count': active_devices} ) return True return False @staticmethod def disable_mfa_for_user(user: User) -> bool: """Disable MFA for user""" user.mfa_enabled = False user.save(update_fields=['mfa_enabled']) # Deactivate all MFA devices MFADevice.objects.filter(user=user).update(is_active=False) # Log MFA disablement AuditLog.objects.create( user=user, action_type='MFA_DISABLED', details={'reason': 'user_requested'} ) return True @staticmethod def get_user_mfa_devices(user: User) -> list: """Get all MFA devices for user""" devices = MFADevice.objects.filter(user=user).order_by('-is_primary', 'name') return [ { 'id': str(device.id), 'name': device.name, 'device_type': device.device_type, 'is_active': device.is_active, 'is_primary': device.is_primary, 'last_used': device.last_used, 'created_at': device.created_at } for device in devices ]