diff --git a/pterodactyl/ptero.py b/pterodactyl/ptero.py index 5e3d798..a2cda7f 100644 --- a/pterodactyl/ptero.py +++ b/pterodactyl/ptero.py @@ -20,6 +20,17 @@ 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: + 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 + 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.""" if await self.config.guild(guild).server_id() is None: @@ -57,11 +68,7 @@ class Pterodactyl(commands.Cog): 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/json" - } + 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) @@ -123,11 +130,12 @@ class Pterodactyl(commands.Cog): @ui.button(label="Yes", style=discord.ButtonStyle.success) async def yes_button(self, button:ui.Button, interaction:discord.Interaction): - requests.post(await Pterodactyl.get_url(self, self.passed_info['guild'], "power"), headers=self.passed_info['headers'], json={"signal": "stop"}) + headers = self.passed_info['headers'] + requests.post(await Pterodactyl.get_url(self, self.passed_info['guild'], "power"), headers=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=self.passed_info['headers']) as response: + 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!") @@ -136,13 +144,13 @@ class Pterodactyl(commands.Cog): 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=self.passed_info['headers'], data=data) - requests.post(url=await Pterodactyl.get_url(self, self.passed_info['guild'], "power"), headers=self.passed_info['headers'], json={"signal": "start"}) + 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 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=self.passed_info['headers']) as response: + 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!") @@ -151,13 +159,13 @@ class Pterodactyl(commands.Cog): 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"), self.passed_info['headers'], data) + await Pterodactyl.put(self, await Pterodactyl.get_url(self, self.passed_info['guild'], "startup/variable"), headers, data) await asyncio.sleep(2) - requests.post(url=await Pterodactyl.get_url(self, self.passed_info['guild'], "power"), headers=self.passed_info['headers'], json={"signal": "start"}) + requests.post(url=await Pterodactyl.get_url(self, self.passed_info['guild'], "power"), headers=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=self.passed_info['headers']) as response: + 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!")