first commit

backend/app/api/v1/websocket.py

@@ -0,0 +1,99 @@
+"""WebSocket API endpoints"""
+
+import asyncio
+import json
+import logging
+from datetime import datetime
+from typing import Optional
+
+from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query
+from jose import jwt, JWTError
+
+from app.core.config import settings
+from app.core.websocket.manager import manager
+
+logger = logging.getLogger(__name__)
+router = APIRouter()
+
+
+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(f"WebSocket auth failed: wrong token type")
+            return None
+        return payload
+    except JWTError as e:
+        logger.warning(f"WebSocket auth failed: {e}")
+        return None
+
+
+@router.websocket("/ws")
+async def websocket_endpoint(
+    websocket: WebSocket,
+    token: str = Query(...),
+):
+    """WebSocket endpoint for real-time data"""
+    logger.info(f"WebSocket connection attempt with token: {token[:20]}...")
+    payload = await authenticate_token(token)
+    if payload is None:
+        logger.warning("WebSocket authentication failed, closing connection")
+        await websocket.close(code=4001)
+        return
+
+    user_id = str(payload.get("sub"))
+    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": [
+                        "gpu_clusters",
+                        "submarine_cables",
+                        "ixp_nodes",
+                        "alerts",
+                        "dashboard",
+                    ],
+                },
+            }
+        )
+
+        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": datetime.utcnow().isoformat()},
+                        }
+                    )
+                elif data.get("type") == "subscribe":
+                    channels = data.get("data", {}).get("channels", [])
+                    await websocket.send_json(
+                        {
+                            "type": "subscription_confirmed",
+                            "data": {"action": "subscribe", "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:
+        manager.disconnect(websocket, user_id)