103 lines
3.2 KiB
Python
103 lines
3.2 KiB
Python
#!/usr/bin/env python3.12
|
|
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
|
|
|
|
import time
|
|
import uuid
|
|
import json
|
|
import random
|
|
from cah.constructors import CAHClient
|
|
from cah.websocket_conn import ConnectionManager
|
|
|
|
class CAH(FastAPI):
|
|
"""CAH Endpoint(s)"""
|
|
def __init__(self, app: FastAPI, util, constants, glob_state): # pylint: disable=super-init-not-called
|
|
self.app = app
|
|
self.util = util
|
|
self.constants = constants
|
|
self.glob_state = glob_state
|
|
|
|
self.ws_endpoints = {
|
|
"cah": self.cah_handler,
|
|
#tbd
|
|
}
|
|
|
|
self.endpoints = {
|
|
#tbd if any non-WS endpoints
|
|
}
|
|
|
|
self.connection_manager = ConnectionManager()
|
|
|
|
|
|
for endpoint, handler in self.ws_endpoints.items():
|
|
print(f"Adding websocket route: {endpoint} @ {handler}")
|
|
app.add_api_websocket_route(f"/{endpoint}/", handler)
|
|
|
|
for endpoint, handler in self.endpoints.items():
|
|
app.add_api_route(f"/{endpoint}/", handler, methods=["POST"])
|
|
|
|
async def cah_handler(self, websocket: WebSocket):
|
|
"""/cah WebSocket"""
|
|
await self.connection_manager.connect(websocket)
|
|
await websocket.send_json({
|
|
"event": "connection_established",
|
|
"ts": int(time.time()),
|
|
})
|
|
|
|
try:
|
|
while True:
|
|
data = await websocket.receive_json()
|
|
if data.get('event') == 'handshake':
|
|
await self.cah_handshake(websocket,
|
|
data)
|
|
except WebSocketDisconnect:
|
|
disconnected = self.connection_manager.get_connection_by_ws(websocket)
|
|
self.connection_manager.disconnect(websocket)
|
|
await self.connection_manager.broadcast({
|
|
"event": "client_disconnected",
|
|
"ts": time.time(),
|
|
"data": {
|
|
"disconnected_resource": disconnected.get('client').resource,
|
|
}
|
|
})
|
|
|
|
|
|
|
|
async def cah_handshake(self, websocket: WebSocket, data):
|
|
"""Handshake"""
|
|
self.connection_manager.connect(websocket)
|
|
data = data.get('data')
|
|
if not data:
|
|
await websocket.send_json({
|
|
"err": "WTF",
|
|
})
|
|
return await websocket.close()
|
|
|
|
csid = str(data.get('csid'))
|
|
resource = data.get('resource')
|
|
platform = data.get('platform')
|
|
|
|
if not csid in self.constants.VALID_CSIDS:
|
|
await websocket.send_json({
|
|
"err": "Unauthorized",
|
|
})
|
|
return await websocket.close()
|
|
|
|
client = CAHClient(
|
|
resource=resource,
|
|
platform=platform,
|
|
csid=csid,
|
|
connected_at=time.time(), # fix
|
|
)
|
|
await self.connection_manager.handshake_complete(websocket, csid, client)
|
|
|
|
await websocket.send_json({
|
|
"event": "handshake_response",
|
|
"ts": int(time.time()),
|
|
"data": {
|
|
"success": True,
|
|
"resource": resource,
|
|
"platform": platform,
|
|
"games": [str(uuid.uuid4()) for x in range(0, 11)],
|
|
},
|
|
})
|