diff --git a/pterodactyl/ptero.py b/pterodactyl/ptero.py index def4d8d..ed0e383 100644 --- a/pterodactyl/ptero.py +++ b/pterodactyl/ptero.py @@ -1,10 +1,10 @@ import asyncio import aiohttp 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.""" @@ -17,8 +17,16 @@ class Pterodactyl(commands.Cog): api_key=None, server_id=None, startup_jar=None, - startup_arguments=None + startup_arguments=None, + power_action_in_progress=False ) + self.session: aiohttp.ClientSession = 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.""" @@ -44,22 +52,17 @@ 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(name="update", description="Updates the server.") @app_commands.guild_only() async def update(self, interaction: discord.Interaction): """Updates the server using the arguments provided in the server's configuration.""" + if await self.config.guild(interaction.guild).power_action_in_progress() is True: + await interaction.response.send_message(ephemeral=True, content="Power action already in progress!\nTry again later.") + return + session = self.session 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 +71,15 @@ 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() + try: 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() + 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 = await response.json() + async with session.get(await self.get_url(interaction.guild, "startup"), headers=headers) as response: + list_var_response_dict = await response.json() updater_startup_vars = [ { "key": "FLAGS", @@ -94,24 +101,25 @@ class Pterodactyl(commands.Cog): } ] if response_dict['attributes']['current_state'] == "offline": + await self.config.guild(interaction.guild).power_action_in_progress.set(True) 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"), headers=headers, json=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 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"), 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(0.5) + 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"), headers=headers, json=data) await interaction_message.edit(content="Updater finished.\nUpdate process completed!") + await self.config.guild(interaction.guild).power_action_in_progress.set(False) elif response_dict['attributes']['current_state'] == "running" or response_dict['attributes']['current_state'] == "starting": passed_info = { "headers": headers, @@ -120,20 +128,23 @@ class Pterodactyl(commands.Cog): "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.") @power.command(name='start', description="Starts the server.") - @power.guild_only() + @app_commands.guild_only() async def start(self, interaction: discord.Interaction): """Starts the server.""" + if await self.config.guild(interaction.guild).power_action_in_progress() is True: + await interaction.response.send_message(ephemeral=True, content="Power action already in progress!\nTry again later.") + return await interaction.response.defer(ephemeral=True, thinking=True) interaction_message = await interaction.original_response() headers = await self.get_headers(interaction.guild) - response = requests.get(await self.get_url(interaction.guild, "resources"), headers=headers) - requests_json = response.json() - current_status = requests_json['attributes']['current_state'] + async with self.session.get(await self.get_url(interaction.guild, "resources"), headers=headers) as response: + response_json = await response.json() + current_status = response_json['attributes']['current_state'] if current_status == "offline": passed_info = { "headers": headers, @@ -144,21 +155,24 @@ class Pterodactyl(commands.Cog): "message": "Server starting...", "completed_message": "Server started!" } - await interaction_message.edit(content="Are you sure you'd like to start the server?", view=self.PowerButtons(timeout=180, passed_info=passed_info)) + await interaction_message.edit(content="Are you sure you'd like to start the server?", view=self.PowerButtons(timeout=180, passed_info=passed_info, session=self.session)) else: message = await interaction_message.edit(content="The server is already running!") await message.delete(delay=3) @power.command(name='restart', description="Restarts the server.") - @power.guild_only() + @app_commands.guild_only() async def restart(self, interaction: discord.Interaction): """Restarts the server.""" + if await self.config.guild(interaction.guild).power_action_in_progress() is True: + await interaction.response.send_message(ephemeral=True, content="Power action already in progress!\nTry again later.") + return await interaction.response.defer(ephemeral=True, thinking=True) interaction_message = await interaction.original_response() headers = await self.get_headers(interaction.guild) - response = requests.get(await self.get_url(interaction.guild, "resources"), headers=headers) - requests_json = response.json() - current_status = requests_json['attributes']['current_state'] + async with self.session.get(await self.get_url(interaction.guild, "resources"), headers=headers) as response: + response_json = await response.json() + current_status = response_json['attributes']['current_state'] if current_status == "running": passed_info = { "headers": headers, @@ -169,7 +183,7 @@ class Pterodactyl(commands.Cog): "message": "Server restarting...", "completed_message": "Server restarted!" } - await interaction_message.edit(content="Are you sure you'd like to restart the server?", view=self.PowerButtons(timeout=180, passed_info=passed_info)) + await interaction_message.edit(content="Are you sure you'd like to restart the server?", view=self.PowerButtons(timeout=180, passed_info=passed_info, session=self.session)) elif current_status == "offline": message = await interaction_message.edit(content="The server is offline!") await message.delete(delay=3) @@ -178,15 +192,18 @@ class Pterodactyl(commands.Cog): await message.delete(delay=3) @power.command(name='stop', description="Stops the server.") - @power.guild_only() + @app_commands.guild_only() async def stop(self, interaction: discord.Interaction): """Stops the server.""" + if await self.config.guild(interaction.guild).power_action_in_progress() is True: + await interaction.response.send_message(ephemeral=True, content="Power action already in progress!\nTry again later.") + return await interaction.response.defer(ephemeral=True, thinking=True) interaction_message = await interaction.original_response() headers = await self.get_headers(interaction.guild) - response = requests.get(await self.get_url(interaction.guild, "resources"), headers=headers) - requests_json = response.json() - current_status = requests_json['attributes']['current_state'] + async with self.session.get(await self.get_url(interaction.guild, "resources"), headers=headers) as response: + response_json = await response.json() + current_status = response_json['attributes']['current_state'] if current_status == "running" or current_status == "starting": passed_info = { "headers": headers, @@ -197,62 +214,62 @@ class Pterodactyl(commands.Cog): "message": "Server stopping...", "completed_message": "Server stopped!" } - await interaction_message.edit(content="Are you sure you'd like to stop the server?", view=self.PowerButtons(timeout=180, passed_info=passed_info)) + await interaction_message.edit(content="Are you sure you'd like to stop the server?", view=self.PowerButtons(timeout=180, passed_info=passed_info, session=self.session)) elif current_status == "offline": message = await interaction_message.edit(content="The server is already offline!") 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"}) + await self.config.guild(self.passed_info['guild']).power_action_in_progress.set(True) + session = self.session + await 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 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"), 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!") + break + else: + await asyncio.sleep(0.5) + 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"}) + await session.put(url=await Pterodactyl.get_url(self, self.passed_info['guild'], "startup/variable"), headers=self.passed_info['headers'], json=data) + 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 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"), 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!") + 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(await Pterodactyl.get_url(self, self.passed_info['guild'], "startup/variable"), headers=self.passed_info['headers'], json=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"), headers=self.passed_info['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=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"), 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!") + break + else: + await asyncio.sleep(0.5) + continue + await self.config.guild(self.passed_info['guild']).power_action_in_progress.set(False) @ui.button(label="No", style=discord.ButtonStyle.danger) async def no_button(self, button:ui.Button, interaction:discord.Interaction): @@ -260,26 +277,29 @@ class Pterodactyl(commands.Cog): await message.delete(delay=3) class PowerButtons(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): + await self.config.guild(self.passed_info['guild']).power_action_in_progress.set(True) headers = self.passed_info['headers'] - requests.post(await Pterodactyl.get_url(self, self.passed_info['guild'], "power"), headers=headers, json={"signal": self.passed_info['signal']}) + await self.session.post(await Pterodactyl.get_url(self, self.passed_info['guild'], "power"), headers=headers, json={"signal": self.passed_info['signal']}) message = await self.passed_info['interaction'].edit_original_response(content=self.passed_info['message'], view=None) + await asyncio.sleep(10) 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'] == self.passed_info['target_signal']: - await message.edit(content=self.passed_info['completed_message']) - break - else: - await asyncio.sleep(1) - continue + async with self.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'] == self.passed_info['target_signal']: + await message.edit(content=self.passed_info['completed_message']) + break + else: + await asyncio.sleep(1) + continue + await self.config.guild(self.passed_info['guild']).power_action_in_progress.set(False) @ui.button(label="No", style=discord.ButtonStyle.danger) async def no_button(self, button:ui.Button, interaction:discord.Interaction): @@ -290,7 +310,7 @@ class Pterodactyl(commands.Cog): get_group = app_commands.Group(name='get', description="Retrieves information from the Pterodactyl API.") @get_group.command(name='url', description="Retrieves the URL for the specified endpoint.") - @get_group.guild_only() + @app_commands.guild_only() async def retrieve_url(self, interaction: discord.Interaction, endpoint: str = None): """Retrieves the URL for the specified endpoint.""" try: @@ -306,13 +326,13 @@ class Pterodactyl(commands.Cog): configure = app_commands.Group(name="config", description="Configures the Pterodactyl cog.") @configure.command(name="api", description="Sets the information used to access the Pterodactyl API.") - @configure.guild_only() + @app_commands.guild_only() async def configure_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(name="update", description="Sets the startup arguments for the update command.") - @configure.guild_only() + @app_commands.guild_only() async def configure_update(self, interaction: discord.Interaction): """Sets the startup arguments for the update command.""" await interaction.response.send_modal(self.StartupConfigModal(self.config))