""" 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()