82 lines
2.8 KiB
Python
82 lines
2.8 KiB
Python
|
"""
|
||
|
|
WebSocket Router - WebSocket 实时推送接口
|
||
|
|
Real-time data push via WebSocket with topic subscriptions.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import logging
|
||
|
|
import secrets
|
||
|
|
|
||
|
|
from fastapi import APIRouter, Query, WebSocket, WebSocketDisconnect
|
||
|
|
|
||
|
|
from app.config import settings
|
||
|
|
from app.websocket_manager import ws_manager, VALID_TOPICS
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
router = APIRouter(tags=["WebSocket / 实时推送"])
|
||
|
|
|
||
|
|
|
||
|
|
@router.websocket("/ws")
|
||
|
|
async def websocket_endpoint(
|
||
|
|
websocket: WebSocket,
|
||
|
|
api_key: str | None = Query(default=None, alias="api_key"),
|
||
|
|
topics: str | None = Query(default=None, description="Comma-separated topics"),
|
||
|
|
):
|
||
|
|
"""
|
||
|
|
WebSocket endpoint for real-time data push.
|
||
|
|
|
||
|
|
Connect: ws://host/ws?api_key=xxx&topics=location,alarm
|
||
|
|
Topics: location, alarm, device_status, attendance, bluetooth
|
||
|
|
If no topics specified, subscribes to all.
|
||
|
|
"""
|
||
|
|
# Authenticate
|
||
|
|
if settings.API_KEY is not None:
|
||
|
|
if api_key is None or not secrets.compare_digest(api_key, settings.API_KEY):
|
||
|
|
# For DB keys, do a simple hash check
|
||
|
|
if api_key is not None:
|
||
|
|
from app.dependencies import _hash_key
|
||
|
|
from app.database import async_session
|
||
|
|
from sqlalchemy import select
|
||
|
|
from app.models import ApiKey
|
||
|
|
|
||
|
|
try:
|
||
|
|
async with async_session() as session:
|
||
|
|
key_hash = _hash_key(api_key)
|
||
|
|
result = await session.execute(
|
||
|
|
select(ApiKey.id).where(
|
||
|
|
ApiKey.key_hash == key_hash,
|
||
|
|
ApiKey.is_active == True, # noqa: E712
|
||
|
|
)
|
||
|
|
)
|
||
|
|
if result.scalar_one_or_none() is None:
|
||
|
|
await websocket.close(code=4001, reason="Invalid API key")
|
||
|
|
return
|
||
|
|
except Exception:
|
||
|
|
await websocket.close(code=4001, reason="Auth error")
|
||
|
|
return
|
||
|
|
else:
|
||
|
|
await websocket.close(code=4001, reason="Missing API key")
|
||
|
|
return
|
||
|
|
|
||
|
|
# Parse topics
|
||
|
|
requested_topics = set()
|
||
|
|
if topics:
|
||
|
|
requested_topics = {t.strip() for t in topics.split(",") if t.strip() in VALID_TOPICS}
|
||
|
|
|
||
|
|
if not await ws_manager.connect(websocket, requested_topics):
|
||
|
|
return
|
||
|
|
|
||
|
|
try:
|
||
|
|
# Keep connection alive, handle pings
|
||
|
|
while True:
|
||
|
|
data = await websocket.receive_text()
|
||
|
|
# Client can send "ping" to keep alive
|
||
|
|
if data.strip().lower() == "ping":
|
||
|
|
await websocket.send_text("pong")
|
||
|
|
except WebSocketDisconnect:
|
||
|
|
pass
|
||
|
|
except Exception:
|
||
|
|
logger.debug("WebSocket connection error", exc_info=True)
|
||
|
|
finally:
|
||
|
|
ws_manager.disconnect(websocket)
|