diff --git a/pterodactyl/ptero.py b/pterodactyl/ptero.py index def4d8d..3c72c90 100644 --- a/pterodactyl/ptero.py +++ b/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) - 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, "resources")) as response: + response_dict = response.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,32 +99,31 @@ 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: - 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")) 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 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,59 +206,60 @@ 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: - 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")) 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 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: - 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")) 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(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: - 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")) 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 + 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)