diff --git a/src/pyg2o/server_v3.py b/src/pyg2o/server_v3.py index 0f28afe..06d82b2 100644 --- a/src/pyg2o/server_v3.py +++ b/src/pyg2o/server_v3.py @@ -1,16 +1,25 @@ -from fastapi import WebSocket, FastAPI, Depends +from fastapi import WebSocket, FastAPI, Depends, HTTPException from fastapi.security import HTTPBasic, HTTPBasicCredentials +from uuid import uuid4 class Server: - def __init__(self, app: FastAPI): + def __init__(self, *, app: FastAPI, server_username: str, server_password: str, client_password: str): self._security = HTTPBasic() + self._server_token: str = '' + self._server_username = server_username + self._server_password = server_password + self._client_password = client_password self._register_routes(app) def _register_routes(self, app): @app.get('/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 {'token': response} @app.websocket('/pyg2o') async def pyg2o_main(websocket: WebSocket): @@ -20,13 +29,33 @@ class Server: async def pyg2o_client(websocket: WebSocket, playerid: int): await self._handle_client_connection(websocket, playerid) - # I have to do this crap, because I've spent like 2hrs to figure out why pyright ignores type: ignore stuff - # I give up + # Я потратил примерно 2ч чтобы понять, почему pyright игнорирует type: ignore + # Я сдаюсь, мне пришлось это добавить + _ = pyg2o_auth _ = pyg2o_main _ = pyg2o_client - async def _verify_server_token(self, headers): - ... + async def _verify_token(self, credentials: HTTPBasicCredentials): + username = credentials.username + password = credentials.password + + if username == self._server_username and password == self._server_password: + token = self._create_server_token() + if token is None: + raise HTTPException(status_code=403) + return token + elif password == self._client_password: + ... + + 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_server_connection(self, websocket: WebSocket): await websocket.accept()