"""WebSocket API endpoints""" import asyncio from datetime import UTC, datetime from typing import Optional from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query from jose import jwt, JWTError from sqlalchemy import text from app.core.config import settings from app.core.enums import UserRole from app.core.logging import get_logger from app.core.time import to_iso8601_utc from app.core.websocket.manager import manager from app.db.session import async_session_factory from app.services.log_tail import LOG_TAIL_CHANNEL, log_tail_manager logger = get_logger(__name__, service="api") router = APIRouter() EARTH_UPDATES_CHANNEL = "earth_updates" async def authenticate_token(token: str) -> Optional[dict]: """Authenticate WebSocket connection via token""" try: payload = jwt.decode(token, settings.SECRET_KEY, algorithms=[settings.ALGORITHM]) if payload.get("type") != "access": logger.warning_event( "WebSocket auth failed: wrong token type", event="auth.websocket.invalid_token_type", ) return None return payload except JWTError as e: logger.warning_event( "WebSocket auth failed", event="auth.websocket.decode_failed", context={"error": str(e)}, ) return None async def load_websocket_user_role(user_id: str | None) -> str | None: if not user_id: return None try: async with async_session_factory() as db: result = await db.execute( text("SELECT role, is_active FROM users WHERE id = :id"), {"id": int(user_id)}, ) row = result.fetchone() except Exception as exc: logger.warning_event( "WebSocket user role lookup failed", event="auth.websocket.role_lookup_failed", context={"user_id": user_id, "error": str(exc)}, ) return None if row is None or not row[1]: return None return str(row[0] or "") @router.websocket("/ws") async def websocket_endpoint( websocket: WebSocket, token: str | None = Query(None), ): """WebSocket endpoint for real-time data""" logger.info_event( "WebSocket connection attempt", event="auth.websocket.connection_attempt", context={"token_preview": f"{token[:8]}..." if token else "anonymous"}, ) payload = await authenticate_token(token) if token else None if token and payload is None: logger.warning_event( "WebSocket authentication failed, closing connection", event="auth.websocket.connection_rejected", ) await websocket.close(code=4001) return is_anonymous = payload is None user_id = str(payload.get("sub")) if payload else f"anonymous:{id(websocket)}" user_role = await load_websocket_user_role(user_id) if payload else None supported_channels = ["vessels", "earth_news", EARTH_UPDATES_CHANNEL] if is_anonymous else [ "gpu_clusters", "submarine_cables", "ixp_nodes", "alerts", "dashboard", "datasource_tasks", "vessels", "earth_news", EARTH_UPDATES_CHANNEL, ] if user_role == UserRole.SUPER_ADMIN.value: supported_channels = [*supported_channels, LOG_TAIL_CHANNEL] await manager.connect(websocket, user_id) try: await websocket.send_json( { "type": "connection_established", "data": { "connection_id": f"conn_{user_id}", "server_version": settings.VERSION, "heartbeat_interval": 30, "supported_channels": supported_channels, }, } ) while True: try: data = await asyncio.wait_for(websocket.receive_json(), timeout=30) if data.get("type") == "heartbeat": await websocket.send_json( { "type": "heartbeat", "data": {"action": "pong", "timestamp": to_iso8601_utc(datetime.now(UTC))}, } ) elif data.get("type") == "subscribe": payload_data = data.get("data", {}) if not isinstance(payload_data, dict): payload_data = {} log_tail_config = None channels = payload_data.get("channels", []) if isinstance(channels, str): channels = [channels] elif not isinstance(channels, list): channels = [] channel = payload_data.get("channel") if channel and channel not in channels: channels = [*channels, channel] if LOG_TAIL_CHANNEL in channels: if user_role != UserRole.SUPER_ADMIN.value: await websocket.send_json( { "type": "subscription_error", "data": {"channel": LOG_TAIL_CHANNEL, "detail": "Only super_admin can subscribe logs"}, } ) channels = [item for item in channels if item != LOG_TAIL_CHANNEL] else: try: log_tail_config = await log_tail_manager.subscribe(websocket, payload_data) except ValueError as exc: await websocket.send_json( { "type": "subscription_error", "data": {"channel": LOG_TAIL_CHANNEL, "detail": str(exc)}, } ) channels = [item for item in channels if item != LOG_TAIL_CHANNEL] if is_anonymous: channels = [channel for channel in channels if channel in supported_channels] vessel_subscription = None if "vessels" in channels and "bbox" in payload_data: try: vessel_subscription = manager.subscribe_vessels(websocket, payload_data) except ValueError as exc: await websocket.send_json( { "type": "subscription_error", "data": {"channel": "vessels", "detail": str(exc)}, } ) continue channels = [channel for channel in channels if channel != "vessels"] manager.subscribe(websocket, channels) await websocket.send_json( { "type": "subscription_confirmed", "data": { "action": "subscribe", "channels": [ *channels, *([LOG_TAIL_CHANNEL] if log_tail_config else []), *(["vessels"] if vessel_subscription else []), ], "vessels": vessel_subscription, "logs_tail": log_tail_config.__dict__ if log_tail_config else None, }, } ) elif data.get("type") == "unsubscribe": channels = data.get("data", {}).get("channels", []) if isinstance(channels, str): channels = [channels] if LOG_TAIL_CHANNEL in channels: await log_tail_manager.unsubscribe(websocket) manager.unsubscribe(websocket, channels) await websocket.send_json( { "type": "subscription_confirmed", "data": {"action": "unsubscribe", "channels": channels}, } ) elif data.get("type") == "control_frame": await websocket.send_json( {"type": "control_acknowledged", "data": {"received": True}} ) else: await websocket.send_json({"type": "ack", "data": {"received": True}}) except asyncio.TimeoutError: await websocket.send_json({"type": "heartbeat", "data": {"action": "ping"}}) except WebSocketDisconnect: pass finally: await log_tail_manager.disconnect(websocket) manager.disconnect(websocket, user_id)