fix(pterodactyl): pylint fixes

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