219 lines
8.9 KiB
Python
219 lines
8.9 KiB
Python
"""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.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 == "super_admin":
|
|
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 != "super_admin":
|
|
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)
|