api/cah_ext/websocket_conn.py

#!/usr/bin/env python3.12
# pylint: disable=bare-except, broad-exception-caught, invalid-name

import time
import logging
from fastapi import WebSocket
from cah_ext.constructors import CAHClient

class ConnectionManager:
    """WS Connection Manager"""
    def __init__(self):
        self.active_connections: dict = {}

    def get_connection_by_ws(self, websocket: WebSocket) -> WebSocket:
        """Get Connection by WS"""
        return self.active_connections.get(websocket)
    
    def get_connection_by_csid(self, csid: str) -> WebSocket:
        """Get Connection by CSID"""
        for connection in self.active_connections:
            if connection.get('csid') == csid:
                return connection
            
    def get_connection_by_resource_label(self, resource: str):
        """Get Connection by Resource Label"""
        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):
        """Send Client and Game Lists"""
        clients = []
        
        for _, client in self.active_connections.items():
            logging.debug("Client: %s", 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):
        """Process WS Client Connection"""
        await websocket.accept()
        self.active_connections[websocket] = {
            'client': None,
            'websocket': websocket,
            }
        
    async def handshake_complete(self,
                                 state,
                                 websocket: WebSocket,
                                 csid: str,
                                 handshakedClient: CAHClient):
        """Process Handshake"""
        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): # pylint: disable=unused-argument
        """Process WS Client Disconnection"""
        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):
        """Send WS Client some data"""
        await websocket.send_json(message)

    async def broadcast(self, message: str):
        """Broadcast data to all connected WS clients"""
        for connection in self.active_connections:
            try:
                await connection.send_json(message)
            except:
                continue