diff --git a/src/pyg2o/server_v3.py b/src/pyg2o/server_v3.py index 06d82b2..c00f0c0 100644 --- a/src/pyg2o/server_v3.py +++ b/src/pyg2o/server_v3.py @@ -1,4 +1,5 @@ -from fastapi import WebSocket, FastAPI, Depends, HTTPException +import json +from fastapi import WebSocket, FastAPI, Depends, HTTPException, WebSocketDisconnect, WebSocketException from fastapi.security import HTTPBasic, HTTPBasicCredentials from uuid import uuid4 @@ -10,22 +11,19 @@ class Server: self._server_username = server_username self._server_password = server_password self._client_password = client_password + self._server_connection: WebSocket | None = None self._register_routes(app) def _register_routes(self, app): - @app.get('/auth') + @app.get('/pyg2o/auth') async def pyg2o_auth(credentials: HTTPBasicCredentials = Depends(self._security)): - response: str | None = await self._verify_token(credentials) - if response is None: - raise HTTPException(status_code=401) + return await self._handle_auth_connection(credentials) - return {'token': response} - - @app.websocket('/pyg2o') + @app.websocket('/pyg2o/server') async def pyg2o_main(websocket: WebSocket): await self._handle_server_connection(websocket) - @app.websocket('/pyg2o/{playerid}') + @app.websocket('/pyg2o/client/{playerid}') async def pyg2o_client(websocket: WebSocket, playerid: int): await self._handle_client_connection(websocket, playerid) @@ -50,15 +48,51 @@ class Server: return None def _create_server_token(self) -> str | None: - # TODO: Добавить поддержку несколько токенов, и запрет на подключение при уже активном - if self._server_token != '': - return None - self._server_token = str(uuid4()) return self._server_token + async def _handle_auth_connection(self, credentials: HTTPBasicCredentials): + response: str | None = await self._verify_token(credentials) + if response is None: + raise HTTPException(status_code=401) + + return {'token': response} + async def _handle_server_connection(self, websocket: WebSocket): + headers = websocket.headers + uuid = headers.get('Authorization') + + if uuid != self._server_token: + await websocket.close() + return + + if self._server_connection is not None: + await self._server_connection.close() + self._server_connection = None + await websocket.accept() + #TODO: Заменить принты на логирование + try: + while True: + try: + data = await websocket.receive_text() + message_data = json.loads(data) + print('Server message:', message_data) + except json.JSONDecodeError as e: + print('JSON Decode exception:', e) + except WebSocketDisconnect: + print('Server socket disconnected') + except WebSocketException as e: + print('Server socket exception:', e) + + async def _process_server_message(self, message: dict): + match message: + case {'type': 'message', 'uuid': id, 'data': data}: + ... + case {'type': 'event', 'uuid': id, 'data': data}: + ... + case _: + raise ValueError(f'Неподдерживаемый тип PyG2O Server сообщения: {message}') async def _handle_client_connection(self, websocket: WebSocket, playerid: int): ...