diff --git a/pterodactyl/ptero.py b/pterodactyl/ptero.py index 540adaf..201378e 100644 --- a/pterodactyl/ptero.py +++ b/pterodactyl/ptero.py @@ -4,7 +4,8 @@ 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.""" @@ -19,21 +20,24 @@ class Pterodactyl(commands.Cog): startup_jar=None, startup_arguments=None ) + self.session: aiohttp.ClientSession = None - async def get_session(self, guild): - 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: + 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.""" + if 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" } - async with aiohttp.ClientSession(headers=headers) as session: - return session + return headers 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.""" @@ -58,7 +62,7 @@ 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) + session = self.session await interaction.response.defer(ephemeral=True, thinking=True) interaction_message = await interaction.original_response() if await self.config.guild(interaction.guild).startup_jar() is None: @@ -70,9 +74,14 @@ 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() - async with session.get(await self.get_url(interaction.guild, "resources")) as response: + try: + headers = await self.get_headers(interaction.guild) + 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 = response.json() - async with session.get(await self.get_url(interaction.guild, "startup")) as response: + async with session.get(await self.get_url(interaction.guild, "startup", headers=headers)) as response: list_var_response_dict = response.json() updater_startup_vars = [ { @@ -96,12 +105,12 @@ class Pterodactyl(commands.Cog): ] if response_dict['attributes']['current_state'] == "offline": for data in updater_startup_vars: - 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 session.put(await self.get_url(interaction.guild, "startup/variable"), headers=headers, data=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 session.get(await self.get_url(interaction.guild, "resources")) as response: + 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.") @@ -110,11 +119,11 @@ class Pterodactyl(commands.Cog): await asyncio.sleep(1) continue for data in old_startup_vars: - await session.put(await self.get_url(interaction.guild, "startup/variable"), data) - session.close() + await session.put(await self.get_url(interaction.guild, "startup/variable"), headers=headers, data=data) 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, @@ -212,10 +221,10 @@ class Pterodactyl(commands.Cog): @ui.button(label="Yes", style=discord.ButtonStyle.success) async def yes_button(self, button:ui.Button, interaction:discord.Interaction): session = self.session - session.post(await Pterodactyl.get_url(self, self.passed_info['guild'], "power"), json={"signal": "stop"}) + 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 session.get(await Pterodactyl.get_url(self, self.passed_info['guild'], "resources")) as response: + 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!") @@ -224,13 +233,13 @@ class Pterodactyl(commands.Cog): await asyncio.sleep(2) continue for data in self.passed_info['updater_startup_vars']: - async with session.put(self, url=await Pterodactyl.get_url(self, self.passed_info['guild'], "startup/variable"), data=data) as response: + async with session.put(self, url=await Pterodactyl.get_url(self, self.passed_info['guild'], "startup/variable"), headers=self.passed_info['headers'], data=data) as response: pass - await session.post(url=await Pterodactyl.get_url(self, self.passed_info['guild'], "power"), 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="Updater started...") await asyncio.sleep(2) while True: - async with session.get(await Pterodactyl.get_url(self, self.passed_info['guild'], "resources")) as response: + 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!") @@ -239,12 +248,12 @@ class Pterodactyl(commands.Cog): await asyncio.sleep(1) continue for data in self.passed_info['old_startup_vars']: - await session.put(self, await Pterodactyl.get_url(self, self.passed_info['guild'], "startup/variable"), data) + await session.put(self, await Pterodactyl.get_url(self, self.passed_info['guild'], "startup/variable"), headers=self.passed_info['headers'], data=data) await asyncio.sleep(2) - await session.post(url=await Pterodactyl.get_url(self, self.passed_info['guild'], "power"), 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 session.get(await Pterodactyl.get_url(self, self.passed_info['guild'], "resources")) as response: + 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!") @@ -252,11 +261,9 @@ 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)