GalacticFactory/GalacticCogs
GalacticFactory/GalacticCogs
Archived
1
1
Fork 0
Code Issues Pull requests Activity

Merge pull request 'Refactor Pterodactyl cog to use aiohttp' (#2) from aiohttp-refactor into main

Browse source 
Reviewed-on: https://git.seaswimmer.cc/GalacticFactory/GalacticCogs/pulls/2
This commit is contained in:
Seaswimmer 2023-07-24 09:29:54 -04:00
commit 4e09380351
1 changed files with 113 additions and 93 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

206
pterodactyl/ptero.py
View file

@ -1,10 +1,10 @@
import asyncio import asyncio
import aiohttp import aiohttp
import discord import discord
import requests
from discord import ui from discord import ui
from discord.ext import commands 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): class Pterodactyl(commands.Cog):
"""Pterodactyl allows you to manage your Pterodactyl Panel from Discord.""" """Pterodactyl allows you to manage your Pterodactyl Panel from Discord."""
@ -17,8 +17,16 @@ class Pterodactyl(commands.Cog):
api_key=None, api_key=None,
server_id=None, server_id=None,
startup_jar=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): async def get_headers(self, guild: discord.Guild):
"""Returns the headers used to access the Pterodactyl API.""" """Returns the headers used to access the Pterodactyl API."""
@ -44,22 +52,17 @@ class Pterodactyl(commands.Cog):
url += '/' + endpoint url += '/' + endpoint
return url 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.command(name="update", description="Updates the server.")
@app_commands.guild_only() @app_commands.guild_only()
async def update(self, interaction: discord.Interaction): async def update(self, interaction: discord.Interaction):
"""Updates the server using the arguments provided in the server's configuration.""" """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) await interaction.response.defer(ephemeral=True, thinking=True)
interaction_message = await interaction.original_response() interaction_message = await interaction.original_response()
if await self.config.guild(interaction.guild).api_key() is None: if await self.config.guild(interaction.guild).startup_jar() 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:
await interaction_message.edit(f"Something went wrong.\nError: `Startup jar not set.`", ephemeral=True) await interaction_message.edit(f"Something went wrong.\nError: `Startup jar not set.`", ephemeral=True)
raise LookupError("Startup jar not set.") raise LookupError("Startup jar not set.")
elif await self.config.guild(interaction.guild).startup_arguments() is None: elif await self.config.guild(interaction.guild).startup_arguments() is None:
@ -68,11 +71,15 @@ class Pterodactyl(commands.Cog):
else: else:
startup_jar = await self.config.guild(interaction.guild).startup_jar() startup_jar = await self.config.guild(interaction.guild).startup_jar()
startup_commands = await self.config.guild(interaction.guild).startup_arguments() startup_commands = await self.config.guild(interaction.guild).startup_arguments()
try:
headers = await self.get_headers(interaction.guild) headers = await self.get_headers(interaction.guild)
response = requests.get(await self.get_url(interaction.guild, "resources"), headers=headers) except LookupError as e:
response_dict = response.json() await interaction_message.edit(f"Something went wrong.\nError: `{e}`", ephemeral=True)
list_var = requests.get(await self.get_url(interaction.guild, "startup"), headers=headers) return
list_var_response_dict = list_var.json() 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 = [ updater_startup_vars = [
{ {
"key": "FLAGS", "key": "FLAGS",
@ -94,24 +101,25 @@ class Pterodactyl(commands.Cog):
} }
] ]
if response_dict['attributes']['current_state'] == "offline": 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: for data in updater_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)
requests.post(await self.get_url(interaction.guild, "power"), headers=headers, json={"signal": "start"}) await session.post(await self.get_url(interaction.guild, "power"), headers=headers, json={"signal": "start"})
await interaction_message.edit(content="Updater started...") await interaction_message.edit(content="Updater started...")
await asyncio.sleep(1) await asyncio.sleep(1)
while True: while True:
async with aiohttp.ClientSession() as session: async with session.get(await self.get_url(interaction.guild, "resources"), headers=headers) as response:
async with session.get(await self.get_url(interaction.guild, "resources"), headers=headers) as response: response_dict = await response.json()
response_dict = await response.json() if response_dict['attributes']['current_state'] == "offline":
if response_dict['attributes']['current_state'] == "offline": await interaction_message.edit(content="Updater finished.")
await interaction_message.edit(content="Updater finished.") break
break else:
else: await asyncio.sleep(0.5)
await asyncio.sleep(1) continue
continue
for data in old_startup_vars: 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 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": elif response_dict['attributes']['current_state'] == "running" or response_dict['attributes']['current_state'] == "starting":
passed_info = { passed_info = {
"headers": headers, "headers": headers,
@ -120,20 +128,23 @@ class Pterodactyl(commands.Cog):
"interaction": interaction, "interaction": interaction,
"guild": interaction.guild, "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 = app_commands.Group(name='power', description="Controls the server's power state.")
@power.command(name='start', description="Starts the server.") @power.command(name='start', description="Starts the server.")
@power.guild_only() @app_commands.guild_only()
async def start(self, interaction: discord.Interaction): async def start(self, interaction: discord.Interaction):
"""Starts the server.""" """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) await interaction.response.defer(ephemeral=True, thinking=True)
interaction_message = await interaction.original_response() interaction_message = await interaction.original_response()
headers = await self.get_headers(interaction.guild) headers = await self.get_headers(interaction.guild)
response = requests.get(await self.get_url(interaction.guild, "resources"), headers=headers) async with self.session.get(await self.get_url(interaction.guild, "resources"), headers=headers) as response:
requests_json = response.json() response_json = await response.json()
current_status = requests_json['attributes']['current_state'] current_status = response_json['attributes']['current_state']
if current_status == "offline": if current_status == "offline":
passed_info = { passed_info = {
"headers": headers, "headers": headers,
@ -144,21 +155,24 @@ class Pterodactyl(commands.Cog):
"message": "Server starting...", "message": "Server starting...",
"completed_message": "Server started!" "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: else:
message = await interaction_message.edit(content="The server is already running!") message = await interaction_message.edit(content="The server is already running!")
await message.delete(delay=3) await message.delete(delay=3)
@power.command(name='restart', description="Restarts the server.") @power.command(name='restart', description="Restarts the server.")
@power.guild_only() @app_commands.guild_only()
async def restart(self, interaction: discord.Interaction): async def restart(self, interaction: discord.Interaction):
"""Restarts the server.""" """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) await interaction.response.defer(ephemeral=True, thinking=True)
interaction_message = await interaction.original_response() interaction_message = await interaction.original_response()
headers = await self.get_headers(interaction.guild) headers = await self.get_headers(interaction.guild)
response = requests.get(await self.get_url(interaction.guild, "resources"), headers=headers) async with self.session.get(await self.get_url(interaction.guild, "resources"), headers=headers) as response:
requests_json = response.json() response_json = await response.json()
current_status = requests_json['attributes']['current_state'] current_status = response_json['attributes']['current_state']
if current_status == "running": if current_status == "running":
passed_info = { passed_info = {
"headers": headers, "headers": headers,
@ -169,7 +183,7 @@ class Pterodactyl(commands.Cog):
"message": "Server restarting...", "message": "Server restarting...",
"completed_message": "Server restarted!" "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": elif current_status == "offline":
message = await interaction_message.edit(content="The server is offline!") message = await interaction_message.edit(content="The server is offline!")
await message.delete(delay=3) await message.delete(delay=3)
@ -178,15 +192,18 @@ class Pterodactyl(commands.Cog):
await message.delete(delay=3) await message.delete(delay=3)
@power.command(name='stop', description="Stops the server.") @power.command(name='stop', description="Stops the server.")
@power.guild_only() @app_commands.guild_only()
async def stop(self, interaction: discord.Interaction): async def stop(self, interaction: discord.Interaction):
"""Stops the server.""" """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) await interaction.response.defer(ephemeral=True, thinking=True)
interaction_message = await interaction.original_response() interaction_message = await interaction.original_response()
headers = await self.get_headers(interaction.guild) headers = await self.get_headers(interaction.guild)
response = requests.get(await self.get_url(interaction.guild, "resources"), headers=headers) async with self.session.get(await self.get_url(interaction.guild, "resources"), headers=headers) as response:
requests_json = response.json() response_json = await response.json()
current_status = requests_json['attributes']['current_state'] current_status = response_json['attributes']['current_state']
if current_status == "running" or current_status == "starting": if current_status == "running" or current_status == "starting":
passed_info = { passed_info = {
"headers": headers, "headers": headers,
@ -197,62 +214,62 @@ class Pterodactyl(commands.Cog):
"message": "Server stopping...", "message": "Server stopping...",
"completed_message": "Server stopped!" "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": elif current_status == "offline":
message = await interaction_message.edit(content="The server is already offline!") message = await interaction_message.edit(content="The server is already offline!")
await message.delete(delay=3) await message.delete(delay=3)
class UpdateButtons(ui.View): class UpdateButtons(ui.View):
def __init__(self, timeout, passed_info): def __init__(self, timeout, passed_info, session: aiohttp.ClientSession):
super().__init__() super().__init__()
self.passed_info = passed_info self.passed_info = passed_info
self.session = session
self.config = Config.get_conf(None, cog_name='Pterodactyl', identifier=457581387213637448123567) self.config = Config.get_conf(None, cog_name='Pterodactyl', identifier=457581387213637448123567)
@ui.button(label="Yes", style=discord.ButtonStyle.success) @ui.button(label="Yes", style=discord.ButtonStyle.success)
async def yes_button(self, button:ui.Button, interaction:discord.Interaction): async def yes_button(self, button:ui.Button, interaction:discord.Interaction):
headers = self.passed_info['headers'] await self.config.guild(self.passed_info['guild']).power_action_in_progress.set(True)
requests.post(await Pterodactyl.get_url(self, self.passed_info['guild'], "power"), headers=headers, json={"signal": "stop"}) 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) await self.passed_info['interaction'].edit_original_response(content="Server stopping...", view=None)
while True: 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()
response_dict = await response.json() if response_dict['attributes']['current_state'] == "offline":
if response_dict['attributes']['current_state'] == "offline": await self.passed_info['interaction'].edit_original_response(content="\nServer stopped!")
await self.passed_info['interaction'].edit_original_response(content="\nServer stopped!") break
break else:
else: await asyncio.sleep(0.5)
await asyncio.sleep(2) continue
continue
for data in self.passed_info['updater_startup_vars']: 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) await session.put(url=await Pterodactyl.get_url(self, self.passed_info['guild'], "startup/variable"), headers=self.passed_info['headers'], json=data)
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="Updater started...") await self.passed_info['interaction'].edit_original_response(content="Updater started...")
await asyncio.sleep(2) await asyncio.sleep(2)
while True: 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()
response_dict = await response.json() if response_dict['attributes']['current_state'] == "offline":
if response_dict['attributes']['current_state'] == "offline": await self.passed_info['interaction'].edit_original_response(content="Updater finished!")
await self.passed_info['interaction'].edit_original_response(content="Updater finished!") break
break else:
else: await asyncio.sleep(1)
await asyncio.sleep(1) continue
continue
for data in self.passed_info['old_startup_vars']: 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) 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...") await self.passed_info['interaction'].edit_original_response(content="Server starting...")
while True: 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()
response_dict = await response.json() if response_dict['attributes']['current_state'] == "running":
if response_dict['attributes']['current_state'] == "running": await self.passed_info['interaction'].edit_original_response(content="Server started!\nUpdate process completed!")
await self.passed_info['interaction'].edit_original_response(content="Server started!\nUpdate process completed!") break
break else:
else: await asyncio.sleep(0.5)
await asyncio.sleep(1) continue
continue await self.config.guild(self.passed_info['guild']).power_action_in_progress.set(False)
@ui.button(label="No", style=discord.ButtonStyle.danger) @ui.button(label="No", style=discord.ButtonStyle.danger)
async def no_button(self, button:ui.Button, interaction:discord.Interaction): async def no_button(self, button:ui.Button, interaction:discord.Interaction):
@ -260,26 +277,29 @@ class Pterodactyl(commands.Cog):
await message.delete(delay=3) await message.delete(delay=3)
class PowerButtons(ui.View): class PowerButtons(ui.View):
def __init__(self, timeout, passed_info): def __init__(self, timeout, passed_info, session: aiohttp.ClientSession):
super().__init__() super().__init__()
self.passed_info = passed_info self.passed_info = passed_info
self.session = session
self.config = Config.get_conf(None, cog_name='Pterodactyl', identifier=457581387213637448123567) self.config = Config.get_conf(None, cog_name='Pterodactyl', identifier=457581387213637448123567)
@ui.button(label="Yes", style=discord.ButtonStyle.success) @ui.button(label="Yes", style=discord.ButtonStyle.success)
async def yes_button(self, button:ui.Button, interaction:discord.Interaction): 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'] 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) message = await self.passed_info['interaction'].edit_original_response(content=self.passed_info['message'], view=None)
await asyncio.sleep(10)
while True: while True:
async with aiohttp.ClientSession() as session: async with self.session.get(await Pterodactyl.get_url(self, self.passed_info['guild'], "resources"), headers=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()
response_dict = await response.json() if response_dict['attributes']['current_state'] == self.passed_info['target_signal']:
if response_dict['attributes']['current_state'] == self.passed_info['target_signal']: await message.edit(content=self.passed_info['completed_message'])
await message.edit(content=self.passed_info['completed_message']) break
break else:
else: await asyncio.sleep(1)
await asyncio.sleep(1) continue
continue await self.config.guild(self.passed_info['guild']).power_action_in_progress.set(False)
@ui.button(label="No", style=discord.ButtonStyle.danger) @ui.button(label="No", style=discord.ButtonStyle.danger)
async def no_button(self, button:ui.Button, interaction:discord.Interaction): 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 = 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.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): async def retrieve_url(self, interaction: discord.Interaction, endpoint: str = None):
"""Retrieves the URL for the specified endpoint.""" """Retrieves the URL for the specified endpoint."""
try: try:
@ -306,13 +326,13 @@ class Pterodactyl(commands.Cog):
configure = app_commands.Group(name="config", description="Configures the Pterodactyl 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.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): async def configure_api(self, interaction: discord.Interaction):
"""Sets the information used to access the Pterdoactyl API.""" """Sets the information used to access the Pterdoactyl API."""
await interaction.response.send_modal(self.APIConfigModal(self.config)) await interaction.response.send_modal(self.APIConfigModal(self.config))
@configure.command(name="update", description="Sets the startup arguments for the update command.") @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): async def configure_update(self, interaction: discord.Interaction):
"""Sets the startup arguments for the update command.""" """Sets the startup arguments for the update command."""
await interaction.response.send_modal(self.StartupConfigModal(self.config)) await interaction.response.send_modal(self.StartupConfigModal(self.config))