File size: 1,768 Bytes
9513328 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | """
WebSocket broadcasting — single source of truth.
Extracted from _legacy_main.py to break circular imports.
Usage:
from app.core.websocket import broadcast_alert, broadcast_scan
await broadcast_alert({"severity": "critical", "title": "..."})
"""
from __future__ import annotations
import asyncio
import json
import logging
from typing import Any
logger = logging.getLogger(__name__)
# In-memory set of active WebSocket connections
_connections: set[Any] = set()
async def broadcast_alert(alert_data: dict[str, Any]) -> None:
"""Broadcast an alert to all connected WebSocket clients."""
payload = json.dumps({"event": "alert", "data": alert_data})
dead: list[Any] = []
for ws in _connections:
try:
await ws.send_text(payload)
except Exception:
dead.append(ws)
for ws in dead:
_connections.discard(ws)
logger.debug("alert_broadcast", connections=len(_connections), alert=alert_data.get("title", "")[:50])
async def broadcast_scan(scan_data: dict[str, Any]) -> None:
"""Broadcast a scan result to all connected WebSocket clients."""
payload = json.dumps({"event": "scan", "data": scan_data})
dead: list[Any] = []
for ws in _connections:
try:
await ws.send_text(payload)
except Exception:
dead.append(ws)
for ws in dead:
_connections.discard(ws)
async def register_connection(ws: Any) -> None:
"""Register a new WebSocket connection."""
_connections.add(ws)
async def unregister_connection(ws: Any) -> None:
"""Remove a WebSocket connection."""
_connections.discard(ws)
def active_connections() -> int:
"""Return count of active WebSocket connections."""
return len(_connections)
|