pterodactyl/websocket.py

@@ -1,3 +1,4 @@
+# pylint: disable=cyclic-import
 import json
 import re
 from typing import Optional, Union
@@ -13,28 +14,12 @@ from pterodactyl.logger import logger
 from pterodactyl.pterodactyl import Pterodactyl
 
 
-
 async def establish_websocket_connection(coginstance: Pterodactyl) -> None:
-    logger.debug("Establishing WebSocket connection")
     base_url = await config.base_url()
-    api_key = await config.api_key()
-    server_id = await config.server_id()
 
-    try:
+    logger.debug("Establishing WebSocket connection")
-        client = PterodactylClient(base_url, api_key, debug=True).client
+
-        coginstance.client = client
+    websocket_credentials = await retrieve_websocket_credentials(coginstance)
-        websocket_credentials = client.servers.get_websocket(server_id)
-        logger.debug("""Websocket connection details retrieved:
-                            Socket: %s
-                            Token: %s...""",
-                            websocket_credentials['data']['socket'],
-                            websocket_credentials['data']['token'][:20]
-                        )
-        #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")
@@ -44,13 +29,16 @@ async def establish_websocket_connection(coginstance: Pterodactyl) -> None:
         logger.debug("Authentication message sent")
 
         coginstance.websocket = websocket
-        current_status = ''
 
+        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 = client.servers.get_websocket(server_id)
+            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")
@@ -84,6 +72,28 @@ async def establish_websocket_connection(coginstance: Pterodactyl) -> None:
                 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()
+
+    try:
+        client = PterodactylClient(base_url, api_key, debug=True).client
+        coginstance.client = client
+        websocket_credentials = client.servers.get_websocket(server_id)
+        logger.debug("""Websocket connection details retrieved:
+                            Socket: %s
+                            Token: %s...""",
+                            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)
+
 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