Merge pull request 'Refactor Pterodactyl cog to use aiohttp' (#2) from aiohttp-refactor into main

Reviewed-on: https://git.seaswimmer.cc/GalacticFactory/GalacticCogs/pulls/2

pterodactyl/ptero.py

@@ -1,10 +1,10 @@
 import asyncio
 import aiohttp
 import discord
-import requests
 from discord import ui
 from discord.ext import commands
-from redbot.core import commands, app_commands, Config
+from redbot.core import Config, app_commands, commands
+
 
 class Pterodactyl(commands.Cog):
     """Pterodactyl allows you to manage your Pterodactyl Panel from Discord."""
@@ -17,8 +17,16 @@ class Pterodactyl(commands.Cog):
             api_key=None,
             server_id=None,
             startup_jar=None,
-            startup_arguments=None
+            startup_arguments=None,
+            power_action_in_progress=False
         )
+        self.session: aiohttp.ClientSession = None
+
+    async def cog_load(self):
+        self.session = aiohttp.ClientSession()
+
+    async def cog_unload(self):
+       await self.session.close()
 
     async def get_headers(self, guild: discord.Guild):
         """Returns the headers used to access the Pterodactyl API."""
@@ -44,22 +52,17 @@ class Pterodactyl(commands.Cog):
             url += '/' + endpoint
         return url
 
-    async def put(self, url: str, headers: dict, data: dict):
-        """Sends an asyncio PUT request to the specified URL with the specified headers and data."""
-        async with aiohttp.ClientSession() as session:
-            async with session.put(url, headers=headers, json=data) as response:
-                return response
-
     @app_commands.command(name="update", description="Updates the server.")
     @app_commands.guild_only()
     async def update(self, interaction: discord.Interaction):
         """Updates the server using the arguments provided in the server's configuration."""
+        if await self.config.guild(interaction.guild).power_action_in_progress() is True:
+            await interaction.response.send_message(ephemeral=True, content="Power action already in progress!\nTry again later.")
+            return
+        session = self.session
         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 +71,15 @@ 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()
+        try:
             headers = await self.get_headers(interaction.guild)
-        response = requests.get(await self.get_url(interaction.guild, "resources"), headers=headers)
-        response_dict = response.json()
-        list_var = requests.get(await self.get_url(interaction.guild, "startup"), headers=headers)
-        list_var_response_dict = list_var.json()
+        except LookupError as e:
+            await interaction_message.edit(f"Something went wrong.\nError: `{e}`", ephemeral=True)
+            return
+        async with session.get(await self.get_url(interaction.guild, "resources"), headers=headers) as response:
+            response_dict = await response.json()
+        async with session.get(await self.get_url(interaction.guild, "startup"), headers=headers) as response:
+            list_var_response_dict = await response.json()
         updater_startup_vars = [
             {
                 "key": "FLAGS",
@@ -94,24 +101,25 @@ class Pterodactyl(commands.Cog):
             }
         ]
         if response_dict['attributes']['current_state'] == "offline":
+            await self.config.guild(interaction.guild).power_action_in_progress.set(True)
             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"), headers=headers, json=data)
+            await session.post(await self.get_url(interaction.guild, "power"), headers=headers, 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:
-                        response_dict = await response.json()
-                        if response_dict['attributes']['current_state'] == "offline":
-                            await interaction_message.edit(content="Updater finished.")
-                            break
-                        else:
-                            await asyncio.sleep(1)
-                            continue
+                async with session.get(await self.get_url(interaction.guild, "resources"), headers=headers) as response:
+                    response_dict = await response.json()
+                    if response_dict['attributes']['current_state'] == "offline":
+                        await interaction_message.edit(content="Updater finished.")
+                        break
+                    else:
+                        await asyncio.sleep(0.5)
+                        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"), headers=headers, json=data)
             await interaction_message.edit(content="Updater finished.\nUpdate process completed!")
+            await self.config.guild(interaction.guild).power_action_in_progress.set(False)
         elif response_dict['attributes']['current_state'] == "running" or response_dict['attributes']['current_state'] == "starting":
             passed_info = {
                 "headers": headers,
@@ -120,20 +128,23 @@ class Pterodactyl(commands.Cog):
                 "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.")
 
     @power.command(name='start', description="Starts the server.")
-    @power.guild_only()
+    @app_commands.guild_only()
     async def start(self, interaction: discord.Interaction):
         """Starts the server."""
+        if await self.config.guild(interaction.guild).power_action_in_progress() is True:
+            await interaction.response.send_message(ephemeral=True, content="Power action already in progress!\nTry again later.")
+            return
         await interaction.response.defer(ephemeral=True, thinking=True)
         interaction_message = await interaction.original_response()
         headers = await self.get_headers(interaction.guild)
-        response = requests.get(await self.get_url(interaction.guild, "resources"), headers=headers)
-        requests_json = response.json()
-        current_status = requests_json['attributes']['current_state']
+        async with self.session.get(await self.get_url(interaction.guild, "resources"), headers=headers) as response:
+            response_json = await response.json()
+        current_status = response_json['attributes']['current_state']
         if current_status == "offline":
             passed_info = {
                 "headers": headers,
@@ -144,21 +155,24 @@ class Pterodactyl(commands.Cog):
                 "message": "Server starting...",
                 "completed_message": "Server started!"
             }
-            await interaction_message.edit(content="Are you sure you'd like to start the server?", view=self.PowerButtons(timeout=180, passed_info=passed_info))
+            await interaction_message.edit(content="Are you sure you'd like to start the server?", view=self.PowerButtons(timeout=180, passed_info=passed_info, session=self.session))
         else:
             message = await interaction_message.edit(content="The server is already running!")
             await message.delete(delay=3)
 
     @power.command(name='restart', description="Restarts the server.")
-    @power.guild_only()
+    @app_commands.guild_only()
     async def restart(self, interaction: discord.Interaction):
         """Restarts the server."""
+        if await self.config.guild(interaction.guild).power_action_in_progress() is True:
+            await interaction.response.send_message(ephemeral=True, content="Power action already in progress!\nTry again later.")
+            return
         await interaction.response.defer(ephemeral=True, thinking=True)
         interaction_message = await interaction.original_response()
         headers = await self.get_headers(interaction.guild)
-        response = requests.get(await self.get_url(interaction.guild, "resources"), headers=headers)
-        requests_json = response.json()
-        current_status = requests_json['attributes']['current_state']
+        async with self.session.get(await self.get_url(interaction.guild, "resources"), headers=headers) as response:
+            response_json = await response.json()
+        current_status = response_json['attributes']['current_state']
         if current_status == "running":
             passed_info = {
                 "headers": headers,
@@ -169,7 +183,7 @@ class Pterodactyl(commands.Cog):
                 "message": "Server restarting...",
                 "completed_message": "Server restarted!"
             }
-            await interaction_message.edit(content="Are you sure you'd like to restart the server?", view=self.PowerButtons(timeout=180, passed_info=passed_info))
+            await interaction_message.edit(content="Are you sure you'd like to restart the server?", view=self.PowerButtons(timeout=180, passed_info=passed_info, session=self.session))
         elif current_status == "offline":
             message = await interaction_message.edit(content="The server is offline!")
             await message.delete(delay=3)
@@ -178,15 +192,18 @@ class Pterodactyl(commands.Cog):
             await message.delete(delay=3)
 
     @power.command(name='stop', description="Stops the server.")
-    @power.guild_only()
+    @app_commands.guild_only()
     async def stop(self, interaction: discord.Interaction):
         """Stops the server."""
+        if await self.config.guild(interaction.guild).power_action_in_progress() is True:
+            await interaction.response.send_message(ephemeral=True, content="Power action already in progress!\nTry again later.")
+            return
         await interaction.response.defer(ephemeral=True, thinking=True)
         interaction_message = await interaction.original_response()
         headers = await self.get_headers(interaction.guild)
-        response = requests.get(await self.get_url(interaction.guild, "resources"), headers=headers)
-        requests_json = response.json()
-        current_status = requests_json['attributes']['current_state']
+        async with self.session.get(await self.get_url(interaction.guild, "resources"), headers=headers) as response:
+            response_json = await response.json()
+        current_status = response_json['attributes']['current_state']
         if current_status == "running" or current_status == "starting":
             passed_info = {
                 "headers": headers,
@@ -197,62 +214,62 @@ class Pterodactyl(commands.Cog):
                 "message": "Server stopping...",
                 "completed_message": "Server stopped!"
             }
-            await interaction_message.edit(content="Are you sure you'd like to stop the server?", view=self.PowerButtons(timeout=180, passed_info=passed_info))
+            await interaction_message.edit(content="Are you sure you'd like to stop the server?", view=self.PowerButtons(timeout=180, passed_info=passed_info, session=self.session))
         elif current_status == "offline":
             message = await interaction_message.edit(content="The server is already offline!")
             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"})
+            await self.config.guild(self.passed_info['guild']).power_action_in_progress.set(True)
+            session = self.session
+            await session.post(await Pterodactyl.get_url(self, self.passed_info['guild'], "power"), headers=self.passed_info['headers'], 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:
-                        response_dict = await response.json()
-                        if response_dict['attributes']['current_state'] == "offline":
-                            await self.passed_info['interaction'].edit_original_response(content="\nServer stopped!")
-                            break
-                        else:
-                            await asyncio.sleep(2)
-                            continue
+                async with session.get(await Pterodactyl.get_url(self, self.passed_info['guild'], "resources"), headers=self.passed_info['headers']) 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!")
+                        break
+                    else:
+                        await asyncio.sleep(0.5)
+                        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"})
+                await session.put(url=await Pterodactyl.get_url(self, self.passed_info['guild'], "startup/variable"), headers=self.passed_info['headers'], json=data)
+            await session.post(url=await Pterodactyl.get_url(self, self.passed_info['guild'], "power"), headers=self.passed_info['headers'], 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:
-                        response_dict = await response.json()
-                        if response_dict['attributes']['current_state'] == "offline":
-                            await self.passed_info['interaction'].edit_original_response(content="Updater finished!")
-                            break
-                        else:
-                            await asyncio.sleep(1)
-                            continue
+                async with session.get(await Pterodactyl.get_url(self, self.passed_info['guild'], "resources"), headers=self.passed_info['headers']) 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!")
+                        break
+                    else:
+                        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(await Pterodactyl.get_url(self, self.passed_info['guild'], "startup/variable"), headers=self.passed_info['headers'], json=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"), headers=self.passed_info['headers'], 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:
-                        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!")
-                            break
-                        else:
-                            await asyncio.sleep(1)
-                            continue
+                async with session.get(await Pterodactyl.get_url(self, self.passed_info['guild'], "resources"), headers=self.passed_info['headers']) 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!")
+                        break
+                    else:
+                        await asyncio.sleep(0.5)
+                        continue
+            await self.config.guild(self.passed_info['guild']).power_action_in_progress.set(False)
 
         @ui.button(label="No", style=discord.ButtonStyle.danger)
         async def no_button(self, button:ui.Button, interaction:discord.Interaction):
@@ -260,26 +277,29 @@ class Pterodactyl(commands.Cog):
             await message.delete(delay=3)
 
     class PowerButtons(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):
+            await self.config.guild(self.passed_info['guild']).power_action_in_progress.set(True)
             headers = self.passed_info['headers']
-            requests.post(await Pterodactyl.get_url(self, self.passed_info['guild'], "power"), headers=headers, json={"signal": self.passed_info['signal']})
+            await self.session.post(await Pterodactyl.get_url(self, self.passed_info['guild'], "power"), headers=headers, json={"signal": self.passed_info['signal']})
             message = await self.passed_info['interaction'].edit_original_response(content=self.passed_info['message'], view=None)
+            await asyncio.sleep(10)
             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:
-                        response_dict = await response.json()
-                        if response_dict['attributes']['current_state'] == self.passed_info['target_signal']:
-                            await message.edit(content=self.passed_info['completed_message'])
-                            break
-                        else:
-                            await asyncio.sleep(1)
-                            continue
+                async with self.session.get(await Pterodactyl.get_url(self, self.passed_info['guild'], "resources"), headers=headers) as response:
+                    response_dict = await response.json()
+                    if response_dict['attributes']['current_state'] == self.passed_info['target_signal']:
+                        await message.edit(content=self.passed_info['completed_message'])
+                        break
+                    else:
+                        await asyncio.sleep(1)
+                        continue
+            await self.config.guild(self.passed_info['guild']).power_action_in_progress.set(False)
 
         @ui.button(label="No", style=discord.ButtonStyle.danger)
         async def no_button(self, button:ui.Button, interaction:discord.Interaction):
@@ -290,7 +310,7 @@ class Pterodactyl(commands.Cog):
     get_group = app_commands.Group(name='get', description="Retrieves information from the Pterodactyl API.")
 
     @get_group.command(name='url', description="Retrieves the URL for the specified endpoint.")
-    @get_group.guild_only()
+    @app_commands.guild_only()
     async def retrieve_url(self, interaction: discord.Interaction, endpoint: str = None):
         """Retrieves the URL for the specified endpoint."""
         try:
@@ -306,13 +326,13 @@ class Pterodactyl(commands.Cog):
     configure = app_commands.Group(name="config", description="Configures the Pterodactyl cog.")
 
     @configure.command(name="api", description="Sets the information used to access the Pterodactyl API.")
-    @configure.guild_only()
+    @app_commands.guild_only()
     async def configure_api(self, interaction: discord.Interaction):
         """Sets the information used to access the Pterdoactyl API."""
         await interaction.response.send_modal(self.APIConfigModal(self.config))
 
     @configure.command(name="update", description="Sets the startup arguments for the update command.")
-    @configure.guild_only()
+    @app_commands.guild_only()
     async def configure_update(self, interaction: discord.Interaction):
         """Sets the startup arguments for the update command."""
         await interaction.response.send_modal(self.StartupConfigModal(self.config))