diff --git a/pterodactyl/websocket.py b/pterodactyl/websocket.py index ba7a164..5645a4b 100644 --- a/pterodactyl/websocket.py +++ b/pterodactyl/websocket.py @@ -1,3 +1,4 @@ +# pylint: disable=cyclic-import import json import re from typing import Optional, Union @@ -13,9 +14,65 @@ from pterodactyl.logger import logger from pterodactyl.pterodactyl import Pterodactyl - async def establish_websocket_connection(coginstance: Pterodactyl) -> None: + base_url = await config.base_url() + logger.debug("Establishing WebSocket connection") + + websocket_credentials = await retrieve_websocket_credentials(coginstance) + + async with websockets.connect(websocket_credentials['data']['socket'], origin=base_url, ping_timeout=60) as websocket: + logger.info("WebSocket connection established") + + auth_message = json.dumps({"event": "auth", "args": [websocket_credentials['data']['token']]}) + await websocket.send(auth_message) + logger.debug("Authentication message sent") + + coginstance.websocket = websocket + + await loop(coginstance, websocket) + +async def loop(coginstance: Pterodactyl, websocket: websockets.WebSocketClientProtocol) -> None: + current_status = '' + while True: + message = await websocket.recv() + if json.loads(message)['event'] in ('token expiring', 'token expired'): + logger.debug("Received token expiring/expired event. Refreshing token.") + websocket_credentials = await retrieve_websocket_credentials(coginstance) + auth_message = json.dumps({"event": "auth", "args": [websocket_credentials['data']['token']]}) + await websocket.send(auth_message) + logger.debug("Authentication message sent") + + if json.loads(message)['event'] == 'auth success': + logger.info("WebSocket authentication successful") + + if json.loads(message)['event'] == 'console output' and await config.console_channel() is not None: + if current_status in ('running', 'offline', ''): + content = remove_ansi_escape_codes(json.loads(message)['args'][0]) + + channel = coginstance.bot.get_channel(await config.console_channel()) + if channel is not None: + if content.startswith('['): + pagified_content = pagify(content, delims=[" ", "\n"]) + for page in pagified_content: + await channel.send(content=page) + + chat_message = await check_if_chat_message(content) + if chat_message: + info = await get_info(chat_message['username']) + if info is not None: + await send_chat_discord(coginstance, chat_message['username'], chat_message['message'], info['data']['player']['avatar']) + else: + await send_chat_discord(coginstance, chat_message['username'], chat_message['message'], 'https://seafsh.cc/u/j3AzqQ.png') + + if json.loads(message)['event'] == 'status': + current_status = json.loads(message)['args'][0] + if await config.console_channel() is not None: + console = coginstance.bot.get_channel(await config.console_channel()) + if console is not None: + await console.send(f"Server status changed! `{json.loads(message)['args'][0]}`") + +async def retrieve_websocket_credentials(coginstance: Pterodactyl) -> Optional[dict]: base_url = await config.base_url() api_key = await config.api_key() server_id = await config.server_id() @@ -30,60 +87,13 @@ async def establish_websocket_connection(coginstance: Pterodactyl) -> None: websocket_credentials['data']['socket'], websocket_credentials['data']['token'][:20] ) + return websocket_credentials #NOTE - The token is truncated to prevent it from being logged in its entirety, for security reasons except exceptions.ClientConfigError as e: return logger.error('Failed to initialize Pterodactyl client: %s', e) except exceptions.PterodactylApiError as e: return logger.error('Failed to retrieve Pterodactyl websocket: %s', e) - async with websockets.connect(websocket_credentials['data']['socket'], origin=base_url, ping_timeout=60) as websocket: - logger.info("WebSocket connection established") - - auth_message = json.dumps({"event": "auth", "args": [websocket_credentials['data']['token']]}) - await websocket.send(auth_message) - logger.debug("Authentication message sent") - - coginstance.websocket = websocket - current_status = '' - - while True: - message = await websocket.recv() - if json.loads(message)['event'] in ('token expiring', 'token expired'): - logger.debug("Received token expiring/expired event. Refreshing token.") - websocket_credentials = client.servers.get_websocket(server_id) - auth_message = json.dumps({"event": "auth", "args": [websocket_credentials['data']['token']]}) - await websocket.send(auth_message) - logger.debug("Authentication message sent") - - if json.loads(message)['event'] == 'auth success': - logger.info("WebSocket authentication successful") - - if json.loads(message)['event'] == 'console output' and await config.console_channel() is not None: - if current_status in ('running', 'offline', ''): - content = remove_ansi_escape_codes(json.loads(message)['args'][0]) - - channel = coginstance.bot.get_channel(await config.console_channel()) - if channel is not None: - if content.startswith('['): - pagified_content = pagify(content, delims=[" ", "\n"]) - for page in pagified_content: - await channel.send(content=page) - - chat_message = await check_if_chat_message(content) - if chat_message: - info = await get_info(chat_message['username']) - if info is not None: - await send_chat_discord(coginstance, chat_message['username'], chat_message['message'], info['data']['player']['avatar']) - else: - await send_chat_discord(coginstance, chat_message['username'], chat_message['message'], 'https://seafsh.cc/u/j3AzqQ.png') - - if json.loads(message)['event'] == 'status': - current_status = json.loads(message)['args'][0] - if await config.console_channel() is not None: - console = coginstance.bot.get_channel(await config.console_channel()) - if console is not None: - await console.send(f"Server status changed! `{json.loads(message)['args'][0]}`") - def remove_ansi_escape_codes(text: str) -> str: ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])') #NOTE - https://chat.openai.com/share/d92f9acf-d776-4fd6-a53f-b14ac15dd540