api/cah/websocket_conn.py

106 lines
3.4 KiB
Python

#!/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
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
await self.broadcast({
"event": "client_disconnected",
"ts": int(time.time()),
"data": {
"disconnected_resource": 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