import asyncio import logging from typing import Set from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends from starlette.websockets import WebSocketState from utils.i18n import get_translations_ws, _, _load_translations from babel.support import Translations from .network import update_container_position, get_network_data router = APIRouter() logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) class WebSocketManager: """Класс для управления подключениями к WebSocket-серверу. Обрабатывает подключение, отключение и трансляцию сообщений для клиентов, подключённых через WebSocket. """ def __init__(self): """Инициализирует менеджер подключений к WebSocket-серверу. Создаёт пустое множество для хранения активных соединений. """ self.active_connections: Set[WebSocket] = set() async def connect(self, websocket: WebSocket): """Подключает клиента к WebSocket-серверу. Принимает подключение и добавляет WebSocket в список активных соединений. Args: websocket (WebSocket): Объект WebSocket для подключения. Raises: Exception: Если произошла ошибка при принятии подключения. """ await websocket.accept() self.active_connections.add(websocket) def disconnect(self, websocket: WebSocket): """Отключает клиента от WebSocket-серверу. Удаляет WebSocket из списка активных соединений. Args: websocket (WebSocket): Объект WebSocket для отключения. """ self.active_connections.remove(websocket) async def broadcast(self, message: dict): """Транслирует сообщения всем подключённым клиентам. Отправляет JSON-данные всем активным соединениям, обрабатывая отключения клиентов. Args: message (dict): Сообщение для трансляции в формате JSON. Raises: WebSocketDisconnect: Если клиент неожиданно отключился. Exception: Если произошла другая ошибка при отправке. """ for connection in list(self.active_connections): try: if connection.client_state == WebSocketState.CONNECTED: await connection.send_json(message) except WebSocketDisconnect: self.disconnect(connection) except Exception as e: logger.error(f"Broadcast error: {e}") self.disconnect(connection) manager = WebSocketManager() @router.websocket("/api/v1") async def websocket_endpoint(websocket: WebSocket, translations: Translations = Depends(get_translations_ws)): """Обработчик WebSocket-соединения. Управляет соединением клиента, принимает данные для обновления позиций контейнеров и транслирует сетевые данные. Args: websocket (WebSocket): Объект WebSocket для обработки соединения. translations: Объект переводов для текущего языка. Raises: WebSocketDisconnect: Если клиент отключился. Exception: Если произошла ошибка при обработке данных. """ await manager.connect(websocket) try: while True: data = await websocket.receive_json() logger.info(f"Received data: {data}") if all(key in data for key in ["container_id", "container_x", "container_y"]): await update_container_position(data, translations=translations) # Для внутреннего вызова используем переводы из WebSocket зависимостии network_data = await get_network_data(translations=translations) logger.info(f"Current network data: {network_data}") await manager.broadcast(network_data) await asyncio.sleep(2) except WebSocketDisconnect: logger.info("Client disconnected") except Exception as e: logger.error(f"Error in websocket_endpoint: {e}") finally: manager.disconnect(websocket) if websocket.client_state != WebSocketState.DISCONNECTED: await websocket.close() @router.on_event("startup") async def startup_event(): """Запуск таска для трансляции сетевых данных при старте приложения. Создаёт асинхронную задачу для периодической трансляции данных. Raises: Exception: Если произошла ошибка при запуске таска. """ asyncio.create_task(broadcast_data()) async def broadcast_data(translations: Translations = None): """Таск для периодической трансляции сетевых данных. Выполняет сканирование сети и отправляет данные всем подключённым клиентам с интервалом в 2 секунды. Args: translations: Объект переводов для текущего языка. Raises: Exception: Если произошла ошибка при получении данных или трансляции. """ translations = translations or _load_translations("ru") while True: try: network_data = await get_network_data(translations=translations) if "activation" not in network_data: logger.error(translations.gettext("Ошибка: Ключ 'activation' отсутствует в полученных данных сети")) await asyncio.sleep(2) continue await manager.broadcast(network_data) except Exception as e: logger.error(f"Error in broadcast_data: {e}") await asyncio.sleep(2)