Updates

ETB-API/security/mfa/totp.py

@@ -0,0 +1,235 @@
+"""
+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
+        ]