#!/usr/bin/env python3.12 import json import time from fastapi import WebSocket from cah.constructors import CAHClient class ConnectionManager: def __init__(self): self.active_connections: dict = {} def get_connection_by_ws(self, websocket: WebSocket) -> WebSocket: return self.active_connections.get(websocket) def get_connection_by_csid(self, csid: str) -> WebSocket: for connection in self.active_connections: if connection.get('csid') == csid: return connection def get_connection_by_resource_label(self, resource: str): for connection in self.active_connections: try: if connection.get('client').get('resource') == resource: return connection except: continue async def send_client_and_game_lists(self, state, websocket: WebSocket): clients = [] games = [game.__dict__ for game in state.games] for ws, client in self.active_connections.items(): print(f"Client: {client}") _client = client.get('client') clients.append({ 'resource': _client.resource, 'platform': _client.platform, 'connected_at': _client.connected_at, }) await websocket.send_json({ "event": "client_list", "ts": int(time.time()), "data": { "clients": clients } }) await websocket.send_json({ "event": "game_list", "ts": int(time.time()), "data": { "games": state.get_games() } }) async def connect(self, websocket: WebSocket): await websocket.accept() self.active_connections[websocket] = { 'client': None, 'websocket': websocket, } async def handshake_complete(self, state, websocket: WebSocket, csid: str, handshakedClient: CAHClient): if websocket in self.active_connections: self.active_connections.pop(websocket) self.active_connections[websocket] = { 'websocket': websocket, 'csid': csid, 'client': handshakedClient, } await self.broadcast({ "event": "client_connected", "ts": int(time.time()), "data": { "connected_resource": handshakedClient.resource, "connected_platform": handshakedClient.platform, } }) await self.send_client_and_game_lists(state, websocket) async def disconnect(self, state, websocket: WebSocket, csid: str = None): disconnected = self.get_connection_by_ws(websocket) disconnected_client = disconnected.get('client') disconnected_resource = disconnected_client.resource disconnected_games = [str(game.id) for game in disconnected_client.games] await self.broadcast({ "event": "client_disconnected", "ts": int(time.time()), "data": { "disconnected_resource": disconnected_resource, } }) await state.remove_resource(disconnected_games, disconnected_resource) self.active_connections.pop(websocket) async def send(self, message: str, websocket: WebSocket): await websocket.send_json(message) async def broadcast(self, message: str): for connection in self.active_connections: try: await connection.send_json(message) except: continue