feat: Добавлена авторизация через UDP пакеты
This commit is contained in:
AURUMVORXX
@@ -1,162 +1,143 @@
|
||||
import websockets
|
||||
import asyncio
|
||||
from __future__ import annotations
|
||||
import json
|
||||
import uuid
|
||||
from typing import Optional
|
||||
from .constants import Constant
|
||||
from .functions.event import callEvent
|
||||
from .serialize import _deserialize
|
||||
from loguru import logger
|
||||
import logging
|
||||
import asyncio
|
||||
from weakref import WeakValueDictionary
|
||||
from fastapi import WebSocket, FastAPI, Depends, HTTPException, WebSocketDisconnect, WebSocketException
|
||||
from fastapi.security import HTTPBasic, HTTPBasicCredentials
|
||||
from uuid import uuid4
|
||||
|
||||
class Server:
|
||||
|
||||
_current_server: Server | None = None
|
||||
|
||||
class PythonWebsocketServer:
|
||||
|
||||
_current_server = None
|
||||
|
||||
def __init__(self, host: str, port: int, whitelist: list[str], ping_interval: int = 30):
|
||||
self.host: str = host
|
||||
self.port: int = port
|
||||
self.ping_interval: int = ping_interval
|
||||
self.whitelist = whitelist
|
||||
|
||||
self._messageHandlers: dict[str, callable] = dict()
|
||||
self._requests_list: dict[str, asyncio.Future] = dict()
|
||||
self._stop_event: asyncio.Event = asyncio.Event()
|
||||
self._connected_socket: Optional[websockets.ClientConnection] = None
|
||||
|
||||
self._registerMessage('event', self._message_event)
|
||||
self._registerMessage('init_constants', self._message_init_constants)
|
||||
self._registerMessage('result', self._message_call_result)
|
||||
|
||||
@classmethod
|
||||
async def get_server(cls):
|
||||
def get_current_server(cls) -> Server:
|
||||
if cls._current_server is None:
|
||||
raise ConnectionError('PyG2O сервер не подключен')
|
||||
|
||||
return cls._current_server
|
||||
|
||||
def _registerMessage(self, type: str, handler: callable):
|
||||
if type in self._messageHandlers:
|
||||
return
|
||||
|
||||
def __init__(self, *, app: FastAPI, server_username: str, server_password: str, client_password: str):
|
||||
Server._current_server = self
|
||||
self._security = HTTPBasic()
|
||||
self._server_token: str = ''
|
||||
self._server_username = server_username
|
||||
self._server_password = server_password
|
||||
self._client_password = client_password
|
||||
|
||||
self._messageHandlers[type] = handler
|
||||
self._logger = logging.getLogger(__name__)
|
||||
self._logger.addHandler(logging.NullHandler())
|
||||
|
||||
self._server_connection: WebSocket | None = None
|
||||
self._requests: WeakValueDictionary[str, asyncio.Future] = WeakValueDictionary()
|
||||
self._register_routes(app)
|
||||
|
||||
@classmethod
|
||||
async def server_call(cls, message: str):
|
||||
return await cls.get_current_server()._call(cls.get_current_server()._server_connection, message)
|
||||
|
||||
async def _call(self, socket: WebSocket | None, message: str):
|
||||
if socket is None:
|
||||
raise ConnectionError('PyG2O сервер не подключен')
|
||||
|
||||
request, data = self._make_request()
|
||||
data['data'] = message
|
||||
data = json.dumps(data)
|
||||
# Меняем синтаксис под Squirrel
|
||||
data = data.replace("'", '\\"').replace('True', 'true').replace('False', 'false')
|
||||
|
||||
async def _callMessage(self, type: str, data: dict):
|
||||
if type not in self._messageHandlers:
|
||||
return
|
||||
|
||||
await self._messageHandlers[type](data)
|
||||
|
||||
async def start(self):
|
||||
async with websockets.serve(
|
||||
self.handle_connection,
|
||||
host=self.host,
|
||||
port=self.port,
|
||||
ping_interval=self.ping_interval,
|
||||
):
|
||||
logger.success(f'Server is started at ws://{self.host}:{self.port}')
|
||||
PythonWebsocketServer._current_server = self
|
||||
asyncio.create_task(callEvent('onInit', **{}))
|
||||
await self._stop_event.wait()
|
||||
|
||||
async def stop(self):
|
||||
PythonWebsocketServer._current_server = None
|
||||
self._connected_socket = None
|
||||
self._stop_event.set()
|
||||
|
||||
async def make_request(self, data: str):
|
||||
if (self._connected_socket is None):
|
||||
return None
|
||||
|
||||
request_id = str(uuid.uuid4())
|
||||
self._requests_list[request_id] = asyncio.get_running_loop().create_future()
|
||||
request = {
|
||||
'type': 'call',
|
||||
await socket.send_text(message)
|
||||
return request
|
||||
|
||||
def _make_request(self):
|
||||
request_id = str(uuid4())
|
||||
request = asyncio.Future()
|
||||
self._requests[request_id] = request
|
||||
|
||||
data = {
|
||||
'uuid': request_id,
|
||||
'data': data,
|
||||
'data': None,
|
||||
}
|
||||
request = json.dumps(request)
|
||||
request = request.replace("'", '\\"')
|
||||
request = request.replace('True', 'true')
|
||||
request = request.replace('False', 'false')
|
||||
|
||||
await self._connected_socket.send(request)
|
||||
result = await asyncio.wait_for(
|
||||
self._requests_list[request_id],
|
||||
timeout=30
|
||||
)
|
||||
return result
|
||||
|
||||
return request, data
|
||||
|
||||
async def _message_event(self, data: dict):
|
||||
if (not isinstance(data['data'], dict) or
|
||||
'event' not in data['data']):
|
||||
def _register_routes(self, app):
|
||||
@app.websocket('/pyg2o/server')
|
||||
async def pyg2o_main(websocket: WebSocket):
|
||||
await self._handle_server_connection(websocket)
|
||||
|
||||
@app.websocket('/pyg2o/client/{playerid}')
|
||||
async def pyg2o_client(websocket: WebSocket, playerid: int):
|
||||
await self._handle_client_connection(websocket, playerid)
|
||||
|
||||
# Я потратил примерно 2ч чтобы понять, почему pyright игнорирует type: ignore
|
||||
# Я сдаюсь, мне пришлось это добавить
|
||||
_ = pyg2o_main
|
||||
_ = pyg2o_client
|
||||
|
||||
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:
|
||||
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:
|
||||
# Закрытие до принятия подключения выбрасывает 403 (Forbidden) код, так что не нужны доп сообщения
|
||||
await websocket.close()
|
||||
return
|
||||
|
||||
eventName = data['data']['event']
|
||||
del data['data']['event']
|
||||
|
||||
if 'desc' in data['data']:
|
||||
obj_name = data['data']['desc']['obj_name']
|
||||
obj_data = data['data']['desc']['obj_data']
|
||||
data['data']['desc'] = _deserialize(obj_name, obj_data)
|
||||
elif 'itemGround' in data['data']:
|
||||
obj_name = data['data']['itemGround']['obj_name']
|
||||
obj_data = data['data']['itemGround']['obj_data']
|
||||
data['data']['itemGround'] = _deserialize(obj_name, obj_data)
|
||||
|
||||
asyncio.create_task(callEvent(eventName, **data['data']))
|
||||
|
||||
async def _message_init_constants(self, data: dict):
|
||||
if data['data'] is not dict:
|
||||
return
|
||||
|
||||
Constant._update(data['data'])
|
||||
|
||||
async def _message_call_result(self, data: dict):
|
||||
if data['uuid'] not in self._requests_list:
|
||||
return
|
||||
|
||||
result = data['data']
|
||||
if (isinstance(data['data'], dict) and
|
||||
'obj_name' in data['data'] and
|
||||
'obj_data' in data['data']):
|
||||
result = _deserialize(result['obj_name'], result['obj_data'])
|
||||
|
||||
self._requests_list[data['uuid']].set_result(result)
|
||||
del self._requests_list[data['uuid']]
|
||||
|
||||
async def handle_connection(self, websocket: websockets.ClientConnection):
|
||||
|
||||
if len(self.whitelist) != 0 and websocket.remote_address[0] not in self.whitelist:
|
||||
await websocket.close(4000, 'Connection denied (whitelist)')
|
||||
return
|
||||
|
||||
if self._connected_socket is not None:
|
||||
await websocket.close(4000, 'Connection denied (already_connected)')
|
||||
return
|
||||
|
||||
self._connected_socket = websocket
|
||||
self.is_connected = websocket
|
||||
logger.info(f'Client connected: {websocket.remote_address}')
|
||||
|
||||
asyncio.create_task(callEvent('onWebsocketConnect', **{}))
|
||||
|
||||
|
||||
if self._server_connection is not None:
|
||||
await self._server_connection.close()
|
||||
|
||||
await websocket.accept()
|
||||
self._server_connection = websocket
|
||||
self._logger.info('PyG2O сервер подключился')
|
||||
|
||||
try:
|
||||
async for message in websocket:
|
||||
while True:
|
||||
try:
|
||||
message_json = json.loads(message)
|
||||
if not all(key in message_json for key in ('type', 'uuid', 'data')):
|
||||
logger.error(f'Expected message with (type, uuid, data) fields, got: {message_json}')
|
||||
continue
|
||||
|
||||
await self._callMessage(message_json['type'], message_json)
|
||||
|
||||
data = await websocket.receive_text()
|
||||
message_data = json.loads(data)
|
||||
self._logger.info(f'Сообщение сервера: {message_data}')
|
||||
except json.JSONDecodeError as e:
|
||||
logger.exception(f'JSON Exception: {e}')
|
||||
continue
|
||||
except Exception as e:
|
||||
logger.exception(f'Exception: {e}')
|
||||
continue
|
||||
except websockets.exceptions.ConnectionClosedError:
|
||||
pass
|
||||
finally:
|
||||
logger.info('Client disconnected')
|
||||
self.is_connected = None
|
||||
self._connected_socket = None
|
||||
asyncio.create_task(callEvent('onWebsocketDisconnect', **{}))
|
||||
self._logger.exception(f'Ошибка декодирования JSON: {e}')
|
||||
except WebSocketDisconnect:
|
||||
self._logger.info('PyG2O сервер отключился')
|
||||
except WebSocketException as e:
|
||||
self._logger.exception(f'Ошибка подключения PyG2O сервера: {e}')
|
||||
|
||||
async def _process_server_message(self, message: dict):
|
||||
match message:
|
||||
case {'uuid': id, 'data': data}:
|
||||
...
|
||||
case {'data': data}:
|
||||
...
|
||||
case _:
|
||||
raise ValueError(f'Неподдерживаемый тип PyG2O Server сообщения: {message}')
|
||||
|
||||
async def _handle_client_connection(self, websocket: WebSocket, playerid: int):
|
||||
...
|
||||
|
||||
Reference in New Issue
Block a user