From b7754399677f90c814176e9ee18a209e7a8d2388 Mon Sep 17 00:00:00 2001
From: SeaswimmerTheFsh <seaswimmerthefsh@gmail.com>
Date: Fri, 1 Mar 2024 14:52:14 -0500
Subject: [PATCH] fix(pterodactyl): cleanup

---
 pterodactyl/websocket.py | 37 +++++++++++++++----------------------
 1 file changed, 15 insertions(+), 22 deletions(-)

diff --git a/pterodactyl/websocket.py b/pterodactyl/websocket.py
index e9a62bf..3fb0750 100644
--- a/pterodactyl/websocket.py
+++ b/pterodactyl/websocket.py
@@ -7,7 +7,7 @@ from typing import Optional, Union
 import aiohttp
 import discord
 import websockets
-from pydactyl import PterodactylClient, exceptions
+from pydactyl import PterodactylClient
 from redbot.core.utils.chat_formatting import pagify
 
 from pterodactyl.config import config
@@ -21,8 +21,6 @@ async def establish_websocket_connection(coginstance: Pterodactyl) -> None:
     logger.info("Establishing WebSocket connection")
 
     websocket_credentials = await retrieve_websocket_credentials(coginstance)
-    if not websocket_credentials:
-        return logger.error("Failed to retrieve WebSocket credentials. WebSocket connection not established")
 
     async with websockets.connect(websocket_credentials['data']['socket'], origin=base_url, ping_timeout=60, logger=getLogger("red.sea.pterodactyl.websocket")) as websocket:
         logger.info("WebSocket connection established")
@@ -121,32 +119,27 @@ async def retrieve_websocket_credentials(coginstance: Pterodactyl) -> Optional[d
     api_key = pterodactyl_keys.get("api_key")
     if api_key is None:
         coginstance.task.cancel()
-        return logger.error("Pterodactyl API key not set. Please set it using `[p]set api`.")
+        raise ValueError("Pterodactyl API key not set. Please set it using `[p]set api`.")
     base_url = await config.base_url()
     if base_url is None:
         coginstance.task.cancel()
-        return logger.error("Pterodactyl base URL not set. Please set it using `[p]ptero set url`.")
+        raise ValueError("Pterodactyl base URL not set. Please set it using `[p]pterodactyl config url`.")
     server_id = await config.server_id()
     if server_id is None:
         coginstance.task.cancel()
-        return logger.error("Pterodactyl server ID not set. Please set it using `[p]ptero set serverid`.")
+        raise ValueError("Pterodactyl server ID not set. Please set it using `[p]pterodactyl config serverid`.")
 
-    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)
+    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
 
 def remove_ansi_escape_codes(text: str) -> str:
     ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')