backend/websocket/enhanced_websocket.py

"""
Enhanced WebSocket Implementation for AegisLM SaaS Backend.

Production-ready WebSocket with authentication,
rate limiting, and comprehensive security features.
"""

import json
import asyncio
import logging
from typing import Dict, Set, Optional
from datetime import datetime, timezone

from fastapi import WebSocket, WebSocketDisconnect, HTTPException, status
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
import redis.asyncio as redis

from core.config import settings
from core.security import verify_token
from services.audit_service import audit_service, AuditEventType

logger = logging.getLogger(__name__)


class EnhancedWebSocketManager:
    """Enhanced WebSocket manager with security and rate limiting."""
    
    def __init__(self):
        self.active_connections: Dict[str, WebSocket] = {}
        self.user_connections: Dict[int, Set[str]] = {}  # user_id -> set of connection_ids
        self.connection_metadata: Dict[str, dict] = {}  # connection_id -> metadata
        self.redis_client = None
        self.rate_limiter = {}  # Track message rates per connection
        
    async def get_redis_client(self):
        """Get Redis client connection."""
        if not self.redis_client:
            self.redis_client = redis.from_url(settings.REDIS_URL)
        return self.redis_client
    
    def generate_connection_id(self) -> str:
        """Generate unique connection ID."""
        import uuid
        return str(uuid.uuid4())
    
    async def authenticate_websocket(self, websocket: WebSocket, token: str) -> Optional[int]:
        """Authenticate WebSocket connection using JWT token."""
        try:
            # Verify JWT token
            user_id = verify_token(token)
            if not user_id:
                await websocket.close(code=4001, reason="Invalid authentication token")
                return None
            
            # Check if user is active
            # In production, you'd verify user status from database
            # For now, assume user is valid if token is valid
            
            return int(user_id)
            
        except Exception as e:
            logger.error(f"WebSocket authentication failed: {str(e)}")
            await websocket.close(code=4002, reason="Authentication failed")
            return None
    
    async def check_rate_limit(self, connection_id: str) -> bool:
        """Check if connection has exceeded message rate limit."""
        if connection_id not in self.rate_limiter:
            self.rate_limiter[connection_id] = {
                "messages": 0,
                "window_start": datetime.now(timezone.utc)
            }
        
        current_time = datetime.now(timezone.utc)
        window_start = self.rate_limiter[connection_id]["window_start"]
        
        # Reset window if it's been more than a minute
        if (current_time - window_start).seconds > 60:
            self.rate_limiter[connection_id] = {
                "messages": 0,
                "window_start": current_time
            }
        
        # Check rate limit (e.g., 60 messages per minute)
        if self.rate_limiter[connection_id]["messages"] >= 60:
            return False
        
        self.rate_limiter[connection_id]["messages"] += 1
        return True
    
    async def connect(self, websocket: WebSocket, token: str):
        """Accept and authenticate WebSocket connection."""
        connection_id = self.generate_connection_id()
        
        # Authenticate connection
        user_id = await self.authenticate_websocket(websocket, token)
        if not user_id:
            return
        
        # Accept connection
        await websocket.accept()
        
        # Store connection
        self.active_connections[connection_id] = websocket
        self.connection_metadata[connection_id] = {
            "user_id": user_id,
            "connected_at": datetime.now(timezone.utc).isoformat(),
            "ip_address": websocket.client.host if websocket.client else "unknown",
            "user_agent": websocket.headers.get("user-agent", "unknown")
        }
        
        # Track user connections
        if user_id not in self.user_connections:
            self.user_connections[user_id] = set()
        self.user_connections[user_id].add(connection_id)
        
        # Log connection
        audit_service.log_user_action(
            action="websocket_connect",
            user_id=user_id,
            ip_address=websocket.client.host if websocket.client else None,
            details={"connection_id": connection_id}
        )
        
        logger.info(f"WebSocket connected: {connection_id} for user {user_id}")
        
        # Send welcome message
        await self.send_personal_message(connection_id, {
            "type": "connection_established",
            "data": {
                "connection_id": connection_id,
                "message": "WebSocket connection established successfully",
                "timestamp": datetime.now(timezone.utc).isoformat()
            }
        })
    
    async def disconnect(self, connection_id: str, reason: str = "Normal disconnect"):
        """Handle WebSocket disconnection."""
        if connection_id in self.active_connections:
            websocket = self.active_connections[connection_id]
            metadata = self.connection_metadata.get(connection_id, {})
            user_id = metadata.get("user_id")
            
            # Remove from active connections
            del self.active_connections[connection_id]
            
            # Remove from user connections
            if user_id and user_id in self.user_connections:
                self.user_connections[user_id].discard(connection_id)
                if not self.user_connections[user_id]:
                    del self.user_connections[user_id]
            
            # Remove metadata
            if connection_id in self.connection_metadata:
                del self.connection_metadata[connection_id]
            
            # Remove from rate limiter
            if connection_id in self.rate_limiter:
                del self.rate_limiter[connection_id]
            
            # Log disconnection
            audit_service.log_user_action(
                action="websocket_disconnect",
                user_id=user_id,
                details={"connection_id": connection_id, "reason": reason}
            )
            
            logger.info(f"WebSocket disconnected: {connection_id} for user {user_id} - Reason: {reason}")
    
    async def send_personal_message(self, connection_id: str, message: dict):
        """Send message to specific connection."""
        if connection_id in self.active_connections:
            websocket = self.active_connections[connection_id]
            await websocket.send_text(json.dumps(message))
    
    async def send_user_message(self, user_id: int, message: dict):
        """Send message to all connections for a specific user."""
        if user_id in self.user_connections:
            for connection_id in self.user_connections[user_id]:
                await self.send_personal_message(connection_id, message)
    
    async def broadcast_message(self, message: dict, exclude_connection: Optional[str] = None):
        """Broadcast message to all authenticated connections."""
        for connection_id, websocket in self.active_connections.items():
            if connection_id != exclude_connection:
                try:
                    await websocket.send_text(json.dumps(message))
                except Exception as e:
                    logger.error(f"Failed to send message to {connection_id}: {str(e)}")
    
    async def handle_message(self, connection_id: str, message: str):
        """Handle incoming WebSocket message."""
        # Check rate limit
        if not await self.check_rate_limit(connection_id):
            await self.send_personal_message(connection_id, {
                "type": "error",
                "data": {
                    "message": "Rate limit exceeded. Please slow down.",
                    "code": "RATE_LIMIT_EXCEEDED"
                }
            })
            return
        
        try:
            # Parse message
            data = json.loads(message)
            message_type = data.get("type", "unknown")
            
            # Handle different message types
            if message_type == "ping":
                await self.send_personal_message(connection_id, {
                    "type": "pong",
                    "data": {"timestamp": datetime.now(timezone.utc).isoformat()}
                })
            elif message_type == "subscribe":
                await self.handle_subscription(connection_id, data.get("data", {}))
            elif message_type == "unsubscribe":
                await self.handle_unsubscription(connection_id, data.get("data", {}))
            else:
                # Log unknown message type
                logger.warning(f"Unknown WebSocket message type: {message_type}")
                
        except json.JSONDecodeError:
            await self.send_personal_message(connection_id, {
                "type": "error",
                "data": {
                    "message": "Invalid JSON format",
                    "code": "INVALID_JSON"
                }
            })
        except Exception as e:
            logger.error(f"Error handling WebSocket message: {str(e)}")
            await self.send_personal_message(connection_id, {
                "type": "error",
                "data": {
                    "message": "Internal server error",
                    "code": "INTERNAL_ERROR"
                }
            })
    
    async def handle_subscription(self, connection_id: str, subscription_data: dict):
        """Handle WebSocket subscription to specific channels."""
        metadata = self.connection_metadata.get(connection_id, {})
        if not metadata:
            return
        
        user_id = metadata.get("user_id")
        channel = subscription_data.get("channel")
        
        if not channel:
            await self.send_personal_message(connection_id, {
                "type": "error",
                "data": {
                    "message": "Channel name required",
                    "code": "INVALID_SUBSCRIPTION"
                }
            })
            return
        
        # Store subscription
        if "subscriptions" not in metadata:
            metadata["subscriptions"] = set()
        metadata["subscriptions"].add(channel)
        
        # Log subscription
        audit_service.log_user_action(
            action="websocket_subscribe",
            user_id=user_id,
            details={"channel": channel, "connection_id": connection_id}
        )
        
        await self.send_personal_message(connection_id, {
            "type": "subscription_confirmed",
            "data": {
                "channel": channel,
                "message": f"Subscribed to {channel}"
            }
        })
    
    async def handle_unsubscription(self, connection_id: str, unsubscription_data: dict):
        """Handle WebSocket unsubscription from channels."""
        metadata = self.connection_metadata.get(connection_id, {})
        if not metadata:
            return
        
        user_id = metadata.get("user_id")
        channel = unsubscription_data.get("channel")
        
        if not channel or "subscriptions" not in metadata:
            return
        
        # Remove subscription
        metadata["subscriptions"].discard(channel)
        
        # Log unsubscription
        audit_service.log_user_action(
            action="websocket_unsubscribe",
            user_id=user_id,
            details={"channel": channel, "connection_id": connection_id}
        )
        
        await self.send_personal_message(connection_id, {
            "type": "unsubscription_confirmed",
            "data": {
                "channel": channel,
                "message": f"Unsubscribed from {channel}"
            }
        })
    
    def get_connection_stats(self) -> dict:
        """Get WebSocket connection statistics."""
        return {
            "total_connections": len(self.active_connections),
            "unique_users": len(self.user_connections),
            "connections_per_user": {
                str(user_id): len(connections) 
                for user_id, connections in self.user_connections.items()
            },
            "timestamp": datetime.now(timezone.utc).isoformat()
        }


# Global enhanced WebSocket manager
enhanced_websocket_manager = EnhancedWebSocketManager()