55 lines
1.7 KiB
Python
55 lines
1.7 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 connect(self, websocket: WebSocket):
|
||
|
|
await websocket.accept()
|
||
|
|
self.active_connections[websocket] = {
|
||
|
|
'client': None,
|
||
|
|
'websocket': websocket,
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
async def handshake_complete(self,
|
||
|
|
websocket: WebSocket,
|
||
|
|
csid: str,
|
||
|
|
handshakedClient: CAHClient):
|
||
|
|
self.active_connections[websocket] = {
|
||
|
|
'websocket': websocket,
|
||
|
|
'csid': csid,
|
||
|
|
'client': handshakedClient,
|
||
|
|
}
|
||
|
|
|
||
|
|
await self.broadcast({
|
||
|
|
"event": "client_connected",
|
||
|
|
"ts": str(time.time()),
|
||
|
|
"data": {
|
||
|
|
"connected_resource": handshakedClient.resource,
|
||
|
|
}
|
||
|
|
})
|
||
|
|
|
||
|
|
def disconnect(self, websocket: WebSocket, csid: str = None):
|
||
|
|
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:
|
||
|
|
await connection.send_json(message)
|