pterodactyl/ptero.py

@@ -20,16 +20,23 @@ class Pterodactyl(commands.Cog):
             startup_arguments=None
         )
 
-    async def get_headers(self, guild: discord.Guild):
-        """Returns the headers used to access the Pterodactyl API."""
-        if await self.config.guild(guild).api_key() is None:
+    async def get_session(self, guild, endpoint = None):
+        if await self.config.guild(guild).server_id() is None:
+            raise LookupError("Server ID not set.")
+        elif await self.config.guild(guild).base_url() is None:
+            raise LookupError("Base URL not set.")
+        elif await self.config.guild(guild).api_key() is None:
             raise LookupError("API Key not set.")
         headers = {
             "Authorization": f"Bearer {await self.config.guild(guild).api_key()}",
             "Content-Type": "application/json",
             "Accept": "application/json"
         }
-        return headers
+        url = f"https://{await self.config.guild(guild).base_url}/api/client/servers/{await self.config.guild(guild).server_id}"
+        if endpoint:
+            url += '/' + endpoint
+        async with aiohttp.ClientSession(url=url, headers=headers) as session:
+            return session
 
     async def get_url(self, guild, endpoint = None):
         """Returns the base url for the server's API, or the url for a specific API endpoint if one is provided."""
@@ -54,12 +61,10 @@ class Pterodactyl(commands.Cog):
     @app_commands.guild_only()
     async def update(self, interaction: discord.Interaction):
         """Updates the server using the arguments provided in the server's configuration."""
+        session = await self.get_session(interaction.guild)
         await interaction.response.defer(ephemeral=True, thinking=True)
         interaction_message = await interaction.original_response()
-        if await self.config.guild(interaction.guild).api_key() is None:
-            await interaction_message.edit(f"Something went wrong.\nError: `API Key not set.`", ephemeral=True)
-            raise LookupError("API Key not set.")
-        elif await self.config.guild(interaction.guild).startup_jar() is None:
+        if await self.config.guild(interaction.guild).startup_jar() is None:
             await interaction_message.edit(f"Something went wrong.\nError: `Startup jar not set.`", ephemeral=True)
             raise LookupError("Startup jar not set.")
         elif await self.config.guild(interaction.guild).startup_arguments() is None:
@@ -68,11 +73,10 @@ class Pterodactyl(commands.Cog):
         else:
             startup_jar = await self.config.guild(interaction.guild).startup_jar()
             startup_commands = await self.config.guild(interaction.guild).startup_arguments()
-            headers = await self.get_headers(interaction.guild)
-        response = requests.get(await self.get_url(interaction.guild, "resources"), headers=headers)
+        async with session.get(await self.get_url(interaction.guild, "resources")) as response:
             response_dict = response.json()
-        list_var = requests.get(await self.get_url(interaction.guild, "startup"), headers=headers)
-        list_var_response_dict = list_var.json()
+        async with session.get(await self.get_url(interaction.guild, "startup")) as response:
+            list_var_response_dict = response.json()
         updater_startup_vars = [
             {
                 "key": "FLAGS",
@@ -95,13 +99,12 @@ class Pterodactyl(commands.Cog):
         ]
         if response_dict['attributes']['current_state'] == "offline":
             for data in updater_startup_vars:
-                await self.put(await self.get_url(interaction.guild, "startup/variable"), headers, data)
-            requests.post(await self.get_url(interaction.guild, "power"), headers=headers, json={"signal": "start"})
+                await session.put(await self.get_url(interaction.guild, "startup/variable"), data)
+            await session.post(await self.get_url(interaction.guild, "power"), json={"signal": "start"})
             await interaction_message.edit(content="Updater started...")
             await asyncio.sleep(1)
             while True:
-                async with aiohttp.ClientSession() as session:
-                    async with session.get(await self.get_url(interaction.guild, "resources"), headers=headers) as response:
+                async with session.get(await self.get_url(interaction.guild, "resources")) as response:
                     response_dict = await response.json()
                     if response_dict['attributes']['current_state'] == "offline":
                         await interaction_message.edit(content="Updater finished.")
@@ -110,17 +113,17 @@ class Pterodactyl(commands.Cog):
                         await asyncio.sleep(1)
                         continue
             for data in old_startup_vars:
-                await self.put(await self.get_url(interaction.guild, "startup/variable"), headers, data)
+                await session.put(await self.get_url(interaction.guild, "startup/variable"), data)
+            session.close()
             await interaction_message.edit(content="Updater finished.\nUpdate process completed!")
         elif response_dict['attributes']['current_state'] == "running" or response_dict['attributes']['current_state'] == "starting":
             passed_info = {
-                "headers": headers,
                 "updater_startup_vars": updater_startup_vars,
                 "old_startup_vars": old_startup_vars,
                 "interaction": interaction,
                 "guild": interaction.guild,
             }
-            await interaction_message.edit(content="The server is already running! Are you sure you'd like to stop the server for updates?", view=self.UpdateButtons(timeout=180, passed_info=passed_info))
+            await interaction_message.edit(content="The server is already running! Are you sure you'd like to stop the server for updates?", view=self.UpdateButtons(timeout=180, passed_info=passed_info, session=session))
 
     power = app_commands.Group(name='power', description="Controls the server's power state.")
 
@@ -203,19 +206,19 @@ class Pterodactyl(commands.Cog):
             await message.delete(delay=3)
 
     class UpdateButtons(ui.View):
-        def __init__(self, timeout, passed_info):
+        def __init__(self, timeout, passed_info, session: aiohttp.ClientSession):
             super().__init__()
             self.passed_info = passed_info
+            self.session = session
             self.config = Config.get_conf(None, cog_name='Pterodactyl', identifier=457581387213637448123567)
 
         @ui.button(label="Yes", style=discord.ButtonStyle.success)
         async def yes_button(self, button:ui.Button, interaction:discord.Interaction):
-            headers = self.passed_info['headers']
-            requests.post(await Pterodactyl.get_url(self, self.passed_info['guild'], "power"), headers=headers, json={"signal": "stop"})
+            session = self.session
+            session.post(await Pterodactyl.get_url(self, self.passed_info['guild'], "power"), json={"signal": "stop"})
             await self.passed_info['interaction'].edit_original_response(content="Server stopping...", view=None)
             while True:
-                async with aiohttp.ClientSession() as session:
-                    async with session.get(await Pterodactyl.get_url(self, self.passed_info['guild'], "resources"), headers=headers) as response:
+                async with session.get(await Pterodactyl.get_url(self, self.passed_info['guild'], "resources")) as response:
                     response_dict = await response.json()
                     if response_dict['attributes']['current_state'] == "offline":
                         await self.passed_info['interaction'].edit_original_response(content="\nServer stopped!")
@@ -224,13 +227,13 @@ class Pterodactyl(commands.Cog):
                         await asyncio.sleep(2)
                         continue
             for data in self.passed_info['updater_startup_vars']:
-                await Pterodactyl.put(self, url=await Pterodactyl.get_url(self, self.passed_info['guild'], "startup/variable"), headers=headers, data=data)
-            requests.post(url=await Pterodactyl.get_url(self, self.passed_info['guild'], "power"), headers=headers, json={"signal": "start"})
+                async with session.put(self, url=await Pterodactyl.get_url(self, self.passed_info['guild'], "startup/variable"), data=data) as response:
+                    pass
+            await session.post(url=await Pterodactyl.get_url(self, self.passed_info['guild'], "power"), json={"signal": "start"})
             await self.passed_info['interaction'].edit_original_response(content="Updater started...")
             await asyncio.sleep(2)
             while True:
-                async with aiohttp.ClientSession() as session:
-                    async with session.get(await Pterodactyl.get_url(self, self.passed_info['guild'], "resources"), headers=headers) as response:
+                async with session.get(await Pterodactyl.get_url(self, self.passed_info['guild'], "resources")) as response:
                     response_dict = await response.json()
                     if response_dict['attributes']['current_state'] == "offline":
                         await self.passed_info['interaction'].edit_original_response(content="Updater finished!")
@@ -239,13 +242,12 @@ class Pterodactyl(commands.Cog):
                         await asyncio.sleep(1)
                         continue
             for data in self.passed_info['old_startup_vars']:
-                await Pterodactyl.put(self, await Pterodactyl.get_url(self, self.passed_info['guild'], "startup/variable"), headers, data)
+                await session.put(self, await Pterodactyl.get_url(self, self.passed_info['guild'], "startup/variable"), data)
             await asyncio.sleep(2)
-            requests.post(url=await Pterodactyl.get_url(self, self.passed_info['guild'], "power"), headers=headers, json={"signal": "start"})
+            await session.post(url=await Pterodactyl.get_url(self, self.passed_info['guild'], "power"), json={"signal": "start"})
             await self.passed_info['interaction'].edit_original_response(content="Server starting...")
             while True:
-                async with aiohttp.ClientSession() as session:
-                    async with session.get(await Pterodactyl.get_url(self, self.passed_info['guild'], "resources"), headers=headers) as response:
+                async with session.get(await Pterodactyl.get_url(self, self.passed_info['guild'], "resources")) as response:
                     response_dict = await response.json()
                     if response_dict['attributes']['current_state'] == "running":
                         await self.passed_info['interaction'].edit_original_response(content="Server started!\nUpdate process completed!")
@@ -253,9 +255,11 @@ class Pterodactyl(commands.Cog):
                     else:
                         await asyncio.sleep(1)
                         continue
+            session.close()
 
         @ui.button(label="No", style=discord.ButtonStyle.danger)
         async def no_button(self, button:ui.Button, interaction:discord.Interaction):
+            self.session.close()
             message = await self.passed_info['interaction'].edit_original_response(content=f"Command cancelled.", view=None)
             await message.delete(delay=3)