ETB/ETB-API/security/mfa/totp.py

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