From afc5868e9e96f608e4f5ab258160bf27cbb22863 Mon Sep 17 00:00:00 2001
From: SeaswimmerTheFsh <seaswimmerthefsh@gmail.com>
Date: Fri, 1 Mar 2024 01:17:31 -0500
Subject: [PATCH] feat(pterodactyl): use red's api key storage system for api
 key storage and not a config key

---
 pterodactyl/pterodactyl.py | 18 ++++++++----------
 pterodactyl/websocket.py   | 12 +++++++++++-
 2 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/pterodactyl/pterodactyl.py b/pterodactyl/pterodactyl.py
index c6e40b3..11a449a 100644
--- a/pterodactyl/pterodactyl.py
+++ b/pterodactyl/pterodactyl.py
@@ -1,6 +1,6 @@
 import asyncio
 import json
-from typing import Optional
+from typing import Mapping, Optional
 
 import discord
 import websockets
@@ -76,6 +76,13 @@ class Pterodactyl(commands.Cog):
         command = command.replace(".$U", username).replace(".$M", message).replace(".$C", str(color))
         return command
 
+    @commands.Cog.listener()
+    async def on_red_api_tokens_update(self, service_name: str, api_tokens: Mapping[str,str]):
+        if service_name == "pterodactyl":
+            logger.info("Configuration value set: api_key\nRestarting task...")
+            self.task.cancel()
+            self.task = self.get_task()
+
     @commands.group(autohelp = True, name = "pterodactyl", aliases = ["ptero"])
     async def pterodactyl(self, ctx: commands.Context) -> None:
         """Pterodactyl allows you to manage your Pterodactyl Panel from Discord."""
@@ -100,15 +107,6 @@ class Pterodactyl(commands.Cog):
         self.task.cancel()
         self.task = self.get_task()
 
-    @pterodactyl_config.command(name = "apikey")
-    async def pterodactyl_config_api_key(self, ctx: commands.Context, *, api_key: str) -> None:
-        """Set the API key for your Pterodactyl Panel."""
-        await config.api_key.set(api_key)
-        await ctx.send(f"API key set to `{api_key[:5]}...{api_key[-4:]}`")
-        logger.info("Configuration value set: api_key = %s\nRestarting task...", api_key)
-        self.task.cancel()
-        self.task = self.get_task()
-
     @pterodactyl_config.command(name = "serverid")
     async def pterodactyl_config_server_id(self, ctx: commands.Context, *, server_id: str) -> None:
         """Set the server ID for your Pterodactyl Panel."""
diff --git a/pterodactyl/websocket.py b/pterodactyl/websocket.py
index e23651c..4bf3695 100644
--- a/pterodactyl/websocket.py
+++ b/pterodactyl/websocket.py
@@ -114,9 +114,19 @@ async def establish_websocket_connection(coginstance: Pterodactyl) -> None:
                                 await chat.send(await config.shutdown_msg())
 
 async def retrieve_websocket_credentials(coginstance: Pterodactyl) -> Optional[dict]:
+    pterodactyl_keys = await coginstance.bot.get_shared_api_tokens("pterodactyl")
+    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`.")
     base_url = await config.base_url()
-    api_key = await config.api_key()
+    if base_url is None:
+        coginstance.task.cancel()
+        return logger.error("Pterodactyl base URL not set. Please set it using `[p]ptero set 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`.")
 
     try:
         client = PterodactylClient(base_url, api_key, debug=True).client