GalacticCogs/pterodactyl/ptero.py

import asyncio
import aiohttp
import discord
import requests
from redbot.core import commands, app_commands, Config

class Pterodactyl(commands.Cog):
    """Pterodactyl allows you to manage your Pterodactyl Panel from Discord."""

    def __init__(self, bot):
        self.bot = bot
        self.config = Config.get_conf(self, identifier=457581387213637448123567)
        self.config.register_guild(
            base_url=None,
            api_key=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 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:
            raise LookupError("Base URL not set.")
        base_url = await self.config.guild(guild).base_url()
        server_id = await self.config.guild(guild).server_id()
        url = f"https://{base_url}/api/client/servers/{server_id}/"
        if endpoint:
            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!"""
        try:
            if endpoint:
                url = await self.get_url(interaction.guild, endpoint)
            else:
                url = await self.get_url(interaction.guild)
        except LookupError as e:
            await interaction.response.send_message(f"Something went wrong.\nError: `{e}`", ephemeral=True)
            return
        await interaction.response.send_message(url, ephemeral=True)

    configure = app_commands.Group(name="config", description="Configures the Pterodactyl cog.")

    @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
        base_url = discord.ui.TextInput(
            label="Base URL",
            placeholder="Input your Pterodactyl Panel's Base URL here, without HTTPS or HTTP.",
            style=discord.TextStyle.paragraph,
            required=False,
            max_length=300
        )
        api_key = discord.ui.TextInput(
            label="API Key",
            placeholder="Input your Pterodactyl Client API Key here.",
            style=discord.TextStyle.short,
            required=False,
            max_length=300
        )
        server_id = discord.ui.TextInput(
            label="Server ID",
            placeholder="Input your Pterodactyl server's Server ID here.",
            style=discord.TextStyle.short,
            required=False,
            max_length=300
        )

        async def on_submit(self, interaction: discord.Interaction):
            message = ""
            if self.base_url.value != "":
                await self.config.guild(interaction.guild).base_url.set(self.base_url.value)
                message += f"Base URL set to `{self.base_url.value}`.\n"
            if self.api_key.value != "":
                await self.config.guild(interaction.guild).api_key.set(self.api_key.value)
                message += f"API Key set to `{self.api_key.value}`.\n"
            if self.server_id.value != "":
                await self.config.guild(interaction.guild).server_id.set(self.server_id.value)
                message += f"Server ID set to `{self.server_id.value}`.\n"
            if message == "":
                send = 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()}`"
            else:
                send = f"Configuration changed:\n{message}"
            await interaction.response.send_message(send, 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 == "":
                send = 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()}```"
            else:
                send = f"Configuration changed:\n{message}"
            await interaction.response.send_message(send, ephemeral=True)