working on updater functionality

pterodactyl/ptero.py

@@ -1,4 +1,7 @@
+import asyncio
+import aiohttp
 import discord
+import requests
 from redbot.core import commands, app_commands, Config
 
 class Pterodactyl(commands.Cog):
@@ -10,11 +13,13 @@ class Pterodactyl(commands.Cog):
         self.config.register_guild(
             base_url=None,
             api_key=None,
-            server_id=None
+            server_id=None,
+            startup_jar=None,
+            startup_arguments=None
         )
 
     async def get_url(self, guild, endpoint = None):
-        """Returns the base url for the Servers API, or the url for a specific endpoint if one is provided."""
+        """Returns the base url for the Servers API, or the url for a specific API endpoint if one is provided."""
         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:
@@ -26,6 +31,118 @@ 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()
+    async def update(self, interaction: discord.Interaction):
+        if await self.config.guild(interaction.guild).api_key() is None:
+            raise LookupError("API Key not set.")
+        elif await self.config.guild(interaction.guild).startup_jar() is None:
+            raise LookupError("Startup jar not set.")
+        elif await self.config.guild(interaction.guild).startup_arguments() is None:
+            raise LookupError("Startup arguments not set.")
+        else:
+            api_key = await self.config.guild(interaction.guild).api_key()
+            startup_jar = await self.config.guild(interaction.guild).startup_jar()
+            startup_commands = await self.config.guild(interaction.guild).startup_arguments()
+        headers = {
+            "Authorization": f"Bearer {api_key}",
+            "Content-Type": "application/json",
+            "Accept": "Application/vnd.pterodactyl.v1+json"
+        }
+        response = requests.get(Pterodactyl.get_url(interaction.guild, "resources"), headers=headers)
+        response_dict = response.json()
+        list_var = requests.get(Pterodactyl.get_url(interaction.guild, "startup"), headers=headers)
+        list_var_response_dict = list_var.json()
+        old_startup_args = [list_var_response_dict['data'][0]['attributes']['server_value'], list_var_response_dict['meta']['startup_command']]
+        put_data = [
+            {
+                "key": "FLAGS",
+                "value": startup_commands
+            },
+            {
+                "key": "SERVER_JARFILE",
+                "value": startup_jar
+            }
+        ]
+        if response_dict['attributes']['current_state'] == "offline":
+            for data in put_data:
+                await self.put(Pterodactyl.get_url(interaction.guild, "startup/variable"), headers, data)
+            new_put_data = [
+                {
+                    "key": "FLAGS",
+                    "value": old_startup_args[1]
+                },
+                {
+                    "key": "SERVER_JARFILE",
+                    "value": old_startup_args[0]
+                }
+            ]
+            requests.post(Pterodactyl.get_url(interaction.guild, "power"), headers=headers, json={"signal": "start"})
+            await interaction.response.send_message("Packwiz installer started...", ephemeral=True)
+            while True:
+                async with aiohttp.ClientSession() as session:
+                    async with session.get(Pterodactyl.get_url(interaction.guild, "resources"), headers=headers) as response:
+                        response_dict = await response.json()
+                        if response_dict['attributes']['current_state'] == "offline":
+                            await interaction.original_response.edit("Packwiz installer finished!")
+                            break
+                        else:
+                            await asyncio.sleep(1)
+                            continue
+            for data in new_put_data:
+                await self.put(Pterodactyl.get_url(interaction.guild, "startup/variable"), headers, data)
+        elif response_dict['attributes']['current_state'] == "running":
+            passed_info = [old_startup_args[0], old_startup_args[1], headers, put_data, new_put_data]
+            await interaction.response.send_message("Server is already running! Are you sure you'd like to stop the server for updates?", ephemeral=True, view=self.UpdateButtons(timeout=180, passed_info=passed_info))
+
+    class UpdateButtons(discord.ui.View):
+        def __init__(self, timeout, passed_info):
+            super().__init__()
+            self.passed_info = passed_info
+
+        @discord.ui.button(label="Yes", style=discord.ButtonStyle.green, emoji="✅")
+        async def yes_button(self, button:discord.ui.Button, interaction:discord.Interaction):
+            await interaction.original_response.edit("Server stopping!")
+            requests.post(await Pterodactyl.get_url(interaction.guild, "power"), headers=self.passed_info[2], json={"signal": "stop"})
+            for data in self.passed_info[3]:
+                await Pterodactyl.put(await Pterodactyl.get_url(interaction.guild, "startup/variable"), self.passed_info[2], data)
+            requests.post(await Pterodactyl.get_url(interaction.guild, "power"), headers=self.passed_info[2], json={"signal": "start"})
+            await interaction.original_response.edit("Packwiz installer started...")
+            while True:
+                async with aiohttp.ClientSession() as session:
+                    async with session.get(await Pterodactyl.get_url(interaction.guild, "resources"), headers=self.passed_info[2]) as response:
+                        response_dict = await response.json()
+                        if response_dict['attributes']['current_state'] == "offline":
+                            await interaction.original_response.edit("Packwiz installer finished!")
+                            break
+                        else:
+                            await asyncio.sleep(1)
+                            continue
+            for data in self.passed_info[4]:
+                await Pterodactyl.put(Pterodactyl.get_url(interaction.guild, "startup/variable"), self.passed_info[2], data)
+            asyncio.sleep(1)
+            requests.post(await Pterodactyl.get_url(interaction.guild, "power"), self.passed_info[2], json={"signal": "start"})
+            await interaction.original_response.edit("Server starting...")
+            while True:
+                async with aiohttp.ClientSession() as session:
+                    async with session.get(await Pterodactyl.get_url(interaction.guild, "resources"), headers=self.passed_info[2]) as response:
+                        response_dict = await response.json()
+                        if response_dict['attributes']['current_state'] == "running":
+                            await interaction.original_response.edit("Server started!\nUpdate process completed!!")
+                            break
+                        else:
+                            await asyncio.sleep(1)
+                            continue
+
+        @discord.ui.button(label="No", style=discord.ButtonStyle.red, emoji="❌")
+        async def no_button(self, button:discord.ui.Button, interaction:discord.Interaction):
+            await interaction.response.edit_message(content=f"Command cancelled.")
+
     @app_commands.command()
     async def test(self, interaction: discord.Interaction, endpoint: str = None):
         """This does stuff!"""
@@ -39,12 +156,19 @@ class Pterodactyl(commands.Cog):
             return
         await interaction.response.send_message(url, ephemeral=True)
 
-    @app_commands.command()
+    configure = app_commands.Group(name="config", description="Configures the Pterodactyl cog.")
-    async def config(self, interaction: discord.Interaction):
-        """Configures the Pterodactyl cog."""
-        await interaction.response.send_modal(self.ConfigModal(self.config))
 
-    class ConfigModal(discord.ui.Modal, title="Pterodactyl Manager Configuration"):
+    @configure.command()
+    async def 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()
+    async def startup(self, interaction: discord.Interaction):
+        """Sets the startup arguments for the update command."""
+        await interaction.response.send_modal(self.StartupConfigModal(self.config))
+
+    class APIConfigModal(discord.ui.Modal, title="Pterodactyl Manager Configuration"):
         def __init__(self, config):
             super().__init__()
             self.config = config
@@ -84,3 +208,34 @@ class Pterodactyl(commands.Cog):
             if message == "":
                 message = f"No changes were made.\nCurrent configuration:\n- Base URL: `{await self.config.guild(interaction.guild).base_url()}`\n- API Key: `{await self.config.guild(interaction.guild).api_key()}`\n- Server ID: `{await self.config.guild(interaction.guild).server_id()}`"
             await interaction.response.send_message(message, ephemeral=True)
+
+    class StartupConfigModal(discord.ui.Modal, title="Pterodactyl Manager Configuration"):
+        def __init__(self, config):
+            super().__init__()
+            self.config = config
+        startup_jar = discord.ui.TextInput(
+            label="Startup .jar",
+            placeholder="Input the name of your updater .jar here.",
+            style=discord.TextStyle.short,
+            required=False,
+            max_length=300
+        )
+        startup_arguments = discord.ui.TextInput(
+            label="Startup Arguments",
+            placeholder="Input your startup arguments here.\nExample:\n-g -s server https://repository.link/here",
+            style=discord.TextStyle.paragraph,
+            required=False,
+            max_length=1000
+        )
+
+        async def on_submit(self, interaction: discord.Interaction):
+            message = ""
+            if self.startup_jar.value != "":
+                await self.config.guild(interaction.guild).startup_jar.set(self.startup_jar.value)
+                message += f"- Startup jar set to `{self.startup_jar.value}`.\n"
+            if self.startup_arguments.value != "":
+                await self.config.guild(interaction.guild).startup_arguments.set(self.startup_arguments.value)
+                message += f"- Startup arguments set to:\n```{self.startup_arguments.value}```.\n"
+            if message == "":
+                message = f"No changes were made.\nCurrent configuration:\n- Startup jar: `{await self.config.guild(interaction.guild).startup_jar()}`\n- Startup arguments:\n```{await self.config.guild(interaction.guild).startup_arguments()}```"
+            await interaction.response.send_message(message, ephemeral=True)