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

Refactor Pterodactyl cog to use aiohttp #2

Merged
cswimr merged 26 commits from aiohttp-refactor into main 2023-07-24 09:29:55 -04:00
Conversation 0 Commits 26 Files changed 1 +34 -27
1 changed files with 34 additions and 27 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files
Showing only changes of commit 85fa243732 - Show all commits

61
pterodactyl/ptero.py
View file

@ -4,7 +4,8 @@ import discord
import requests 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."""
@ -19,21 +20,24 @@ class Pterodactyl(commands.Cog):
startup_jar=None, startup_jar=None,
startup_arguments=None startup_arguments=None
) )
self.session: aiohttp.ClientSession = None
async def get_session(self, guild): async def cog_load(self):
if await self.config.guild(guild).server_id() is None: self.session = aiohttp.ClientSession()
raise LookupError("Server ID not set.")
elif await self.config.guild(guild).base_url() is None: async def cog_unload(self):
raise LookupError("Base URL not set.") await self.session.close()
elif await self.config.guild(guild).api_key() is 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.") raise LookupError("API Key not set.")
headers = { headers = {
"Authorization": f"Bearer {await self.config.guild(guild).api_key()}", "Authorization": f"Bearer {await self.config.guild(guild).api_key()}",
"Content-Type": "application/json", "Content-Type": "application/json",
"Accept": "application/json" "Accept": "application/json"
} }
async with aiohttp.ClientSession(headers=headers) as session: return headers
return session
async def get_url(self, guild, endpoint = None): 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.""" """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() @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."""
session = await self.get_session(interaction.guild) 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).startup_jar() is None: if await self.config.guild(interaction.guild).startup_jar() is None:
@ -70,9 +74,14 @@ 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()
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() 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() list_var_response_dict = response.json()
updater_startup_vars = [ updater_startup_vars = [
{ {
@ -96,12 +105,12 @@ class Pterodactyl(commands.Cog):
] ]
if response_dict['attributes']['current_state'] == "offline": if response_dict['attributes']['current_state'] == "offline":
for data in updater_startup_vars: for data in updater_startup_vars:
await session.put(await self.get_url(interaction.guild, "startup/variable"), data) 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"), 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 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() 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.")
@ -110,11 +119,11 @@ class Pterodactyl(commands.Cog):
await asyncio.sleep(1) await asyncio.sleep(1)
continue continue
for data in old_startup_vars: for data in old_startup_vars:
await session.put(await self.get_url(interaction.guild, "startup/variable"), data) await session.put(await self.get_url(interaction.guild, "startup/variable"), headers=headers, data=data)
session.close()
await interaction_message.edit(content="Updater finished.\nUpdate process completed!") await interaction_message.edit(content="Updater finished.\nUpdate process completed!")
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,
"updater_startup_vars": updater_startup_vars, "updater_startup_vars": updater_startup_vars,
"old_startup_vars": old_startup_vars, "old_startup_vars": old_startup_vars,
"interaction": interaction, "interaction": interaction,
@ -212,10 +221,10 @@ class Pterodactyl(commands.Cog):
@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):
session = self.session 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) await self.passed_info['interaction'].edit_original_response(content="Server stopping...", view=None)
while True: 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() 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!")
@ -224,13 +233,13 @@ class Pterodactyl(commands.Cog):
await asyncio.sleep(2) await asyncio.sleep(2)
continue continue
for data in self.passed_info['updater_startup_vars']: 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 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 self.passed_info['interaction'].edit_original_response(content="Updater started...")
await asyncio.sleep(2) await asyncio.sleep(2)
while True: 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() 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!")
@ -239,12 +248,12 @@ class Pterodactyl(commands.Cog):
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 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 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...") await self.passed_info['interaction'].edit_original_response(content="Server starting...")
while True: 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() 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!")
@ -252,11 +261,9 @@ class Pterodactyl(commands.Cog):
else: else:
await asyncio.sleep(1) await asyncio.sleep(1)
continue continue
session.close()
@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):
self.session.close()
message = await self.passed_info['interaction'].edit_original_response(content=f"Command cancelled.", view=None) message = await self.passed_info['interaction'].edit_original_response(content=f"Command cancelled.", view=None)
await message.delete(delay=3) await message.delete(delay=3)