100 lines
3.3 KiB
Python
100 lines
3.3 KiB
Python
"""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)
|