cleanup(pterodactyl): split out a bunch of stuff into separate files
Some checks failed
Actions / Lint Code (Ruff & Pylint) (pull_request) Failing after 19s
Actions / Build Documentation (MkDocs) (pull_request) Successful in 23s

This commit is contained in:
Seaswimmer 2024-02-29 23:26:24 -05:00
parent 417f297193
commit e1c98aa78f
Signed by: cswimr
GPG key ID: B8953EC01E5C4063
4 changed files with 189 additions and 170 deletions

18
pterodactyl/config.py Normal file
View file

@ -0,0 +1,18 @@
from redbot.core import Config
config: Config = Config.get_conf(None, identifier=457581387213637448123567, cog_name="Pterodactyl")
def register_config(config_obj: Config) -> None:
config_obj.register_global(
base_url=None,
api_key=None,
server_id=None,
console_channel=None,
startup_jar=None,
startup_arguments=None,
power_action_in_progress=False,
chat_regex=r"\[(\d{2}:\d{2}:\d{2})\sINFO\]:\s<(\w+)>\s(.*)",
chat_command='tellraw @a ["",{"text":".$U ","color":".$C"},{"text":" (DISCORD): ","color":"blue"},{"text":".$M","color":"white"}]',
api_endpoint="minecraft",
chat_channel=None
)

3
pterodactyl/logger.py Normal file
View file

@ -0,0 +1,3 @@
import logging
logger = logging.getLogger('red.sea.pterodactyl')

View file

@ -1,16 +1,15 @@
import asyncio
import json
import logging
import re
from typing import Optional, Union
import aiohttp
import discord
import websockets
from pydactyl import PterodactylClient, exceptions
from redbot.core import Config, commands
from redbot.core import commands
from redbot.core.bot import Red
from redbot.core.utils.chat_formatting import box, pagify
from redbot.core.utils.chat_formatting import box
from pterodactyl.config import config
from pterodactyl.logger import logger
from pterodactyl.websocket import establish_websocket_connection
class Pterodactyl(commands.Cog):
@ -18,229 +17,104 @@ class Pterodactyl(commands.Cog):
def __init__(self, bot: Red):
self.bot = bot
self.config = Config.get_conf(self, identifier=457581387213637448123567, force_registration=True)
self.config.register_global(
base_url=None,
api_key=None,
server_id=None,
console_channel=None,
startup_jar=None,
startup_arguments=None,
power_action_in_progress=False,
chat_regex=r"\[(\d{2}:\d{2}:\d{2})\sINFO\]:\s<(\w+)>\s(.*)",
chat_command='tellraw @a ["",{"text":".$U ","color":".$C"},{"text":" (DISCORD): ","color":"blue"},{"text":".$M","color":"white"}]',
api_endpoint="minecraft",
chat_channel=None
)
self.logger = logging.getLogger('red.sea.pterodactyl')
self.client = None
self.task = None
self.websocket = None
async def establish_websocket_connection(self) -> None:
self.logger.debug("Establishing WebSocket connection")
base_url = await self.config.base_url()
api_key = await self.config.api_key()
server_id = await self.config.server_id()
async def cog_load(self) -> None:
self.task = self.get_task()
try:
client = PterodactylClient(base_url, api_key, debug=True).client
self.client = client
websocket_credentials = client.servers.get_websocket(server_id)
self.logger.debug("""Websocket connection details retrieved:
Socket: %s
Token: %s...""",
websocket_credentials['data']['socket'],
websocket_credentials['data']['token'][:20]
)
#NOTE - The token is truncated to prevent it from being logged in its entirety, for security reasons
except exceptions.ClientConfigError as e:
return self.logger.error('Failed to initialize Pterodactyl client: %s', e)
except exceptions.PterodactylApiError as e:
return self.logger.error('Failed to retrieve Pterodactyl websocket: %s', e)
async def cog_unload(self) -> None:
self.task.cancel()
await self.client._session.close() # pylint: disable=protected-access
async with websockets.connect(websocket_credentials['data']['socket'], origin=base_url, ping_timeout=60) as websocket:
self.logger.info("WebSocket connection established")
auth_message = json.dumps({"event": "auth", "args": [websocket_credentials['data']['token']]})
await websocket.send(auth_message)
self.logger.debug("Authentication message sent")
self.websocket = websocket
current_status = ''
while True:
message = await websocket.recv()
if json.loads(message)['event'] in ('token expiring', 'token expired'):
self.logger.debug("Received token expiring/expired event. Refreshing token.")
websocket_credentials = client.servers.get_websocket(server_id)
auth_message = json.dumps({"event": "auth", "args": [websocket_credentials['data']['token']]})
await websocket.send(auth_message)
self.logger.debug("Authentication message sent")
if json.loads(message)['event'] == 'auth success':
self.logger.info("WebSocket authentication successful")
if json.loads(message)['event'] == 'console output' and await self.config.console_channel() is not None:
if current_status in ('running', 'offline', ''):
content = self.remove_ansi_escape_codes(json.loads(message)['args'][0])
channel = self.bot.get_channel(await self.config.console_channel())
if channel is not None:
if content.startswith('['):
pagified_content = pagify(content, delims=[" ", "\n"])
for page in pagified_content:
await channel.send(content=page)
chat_message = await self.check_if_chat_message(content)
if chat_message:
info = await self.get_info(chat_message['username'])
if info is not None:
await self.send_chat_discord(chat_message['username'], chat_message['message'], info['data']['player']['avatar'])
else:
await self.send_chat_discord(chat_message['username'], chat_message['message'], 'https://seafsh.cc/u/j3AzqQ.png')
if json.loads(message)['event'] == 'status':
current_status = json.loads(message)['args'][0]
if await self.config.console_channel() is not None:
console = self.bot.get_channel(await self.config.console_channel())
if console is not None:
await console.send(f"Server status changed! `{json.loads(message)['args'][0]}`")
def remove_ansi_escape_codes(self, text: str) -> str:
ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
#NOTE - https://chat.openai.com/share/d92f9acf-d776-4fd6-a53f-b14ac15dd540
return ansi_escape.sub('', text)
async def check_if_chat_message(self, text: str) -> Union[bool, dict]:
self.logger.debug("Checking if message is a chat message")
regex = await self.config.chat_regex()
match: Optional[re.Match[str]] = re.match(regex, text)
if match:
groups = {"time": match.group(1), "username": match.group(2), "message": match.group(3)}
self.logger.debug("Message is a chat message\n%s", json.dumps(groups))
return groups
self.logger.debug("Message is not a chat message")
return False
async def get_info(self, username: str) -> Optional[dict]:
self.logger.debug("Retrieving player info for %s", username)
endpoint = await self.config.api_endpoint()
async with aiohttp.ClientSession() as session:
async with session.get(f"https://playerdb.co/api/player/{endpoint}/{username}") as response:
if response.status == 200:
self.logger.debug("Player info retrieved for %s\n%s", username, json.dumps(await response.json()))
return await response.json()
self.logger.error("Failed to retrieve player info for %s: %s", username, response.status)
return None
async def send_chat_discord(self, username: str, message: str, avatar_url: str) -> None:
self.logger.debug("Sending chat message to Discord")
channel = self.bot.get_channel(await self.config.chat_channel())
if channel is not None:
webhooks = await channel.webhooks()
webhook = discord.utils.get(webhooks, name="Pterodactyl Chat")
if webhook is None:
webhook = await channel.create_webhook(name="Pterodactyl Chat")
await webhook.send(content=message, username=username, avatar_url=avatar_url)
self.logger.debug("Chat message sent to Discord")
else:
self.logger.debug("Chat channel not set. Skipping sending chat message to Discord")
async def get_chat_command(self, username: str, message: str, color: discord.Color):
command = await self.config.chat_command()
command = command.replace(".$U", username).replace(".$M", message).replace(".$C", str(color))
return command
def get_task(self):
task = self.bot.loop.create_task(self.establish_websocket_connection(), name="Pterodactyl Websocket Connection")
def get_task(self) -> asyncio.Task:
task = self.bot.loop.create_task(establish_websocket_connection(self), name="Pterodactyl Websocket Connection")
task.add_done_callback(self.error_callback)
return task
def error_callback(self, fut): #NOTE - Thanks flame442 and zephyrkul for helping me figure this out
def error_callback(self, fut) -> None: #NOTE - Thanks flame442 and zephyrkul for helping me figure this out
try:
fut.result()
except asyncio.CancelledError:
pass
except Exception as e: # pylint: disable=broad-exception-caught
self.logger.error("WebSocket task has failed: %s", e, exc_info=e)
logger.error("WebSocket task has failed: %s", e, exc_info=e)
self.task.cancel()
self.task = self.get_task()
async def cog_load(self):
self.task = self.get_task()
async def cog_unload(self):
self.task.cancel()
await self.client._session.close() # pylint: disable=protected-access
@commands.Cog.listener()
async def on_message_without_command(self, message: discord.Message):
if message.channel.id == await self.config.console_channel() and message.author.bot is False:
self.logger.debug("Received console command from %s: %s", message.author.id, message.content)
async def on_message_without_command(self, message: discord.Message) -> None:
if message.channel.id == await config.console_channel() and message.author.bot is False:
logger.debug("Received console command from %s: %s", message.author.id, message.content)
await message.channel.send(f"Received console command from {message.author.id}: {message.content[:1900]}")
try:
await self.websocket.send(json.dumps({"event": "send command", "args": [message.content]}))
except websockets.exceptions.ConnectionClosed as e:
self.logger.error("WebSocket connection closed: %s", e)
logger.error("WebSocket connection closed: %s", e)
self.task.cancel()
self.task = self.get_task()
if message.channel.id == await self.config.chat_channel() and message.author.bot is False:
self.logger.debug("Received chat message from %s: %s", message.author.id, message.content)
channel = self.bot.get_channel(await self.config.console_channel())
if message.channel.id == await config.chat_channel() and message.author.bot is False:
logger.debug("Received chat message from %s: %s", message.author.id, message.content)
channel = self.bot.get_channel(await config.console_channel())
if channel:
await channel.send(f"Received chat message from {message.author.id}: {message.content[:1900]}")
msg = json.dumps({"event": "send command", "args": [await self.get_chat_command(message.author.display_name, message.content, message.author.color)]})
self.logger.debug("Sending chat message to server:\n%s", msg)
logger.debug("Sending chat message to server:\n%s", msg)
try:
await self.websocket.send(msg)
except websockets.exceptions.ConnectionClosed as e:
self.logger.error("WebSocket connection closed: %s", e)
logger.error("WebSocket connection closed: %s", e)
self.task.cancel()
self.task = self.get_task()
async def get_chat_command(self, username: str, message: str, color: discord.Color) -> str:
command: str = await config.chat_command()
command = command.replace(".$U", username).replace(".$M", message).replace(".$C", str(color))
return command
@commands.group(autohelp = True, name = "pterodactyl", aliases = ["ptero"])
async def pterodactyl(self, ctx: commands.Context):
async def pterodactyl(self, ctx: commands.Context) -> None:
"""Pterodactyl allows you to manage your Pterodactyl Panel from Discord."""
@pterodactyl.group(autohelp = True, name = "config", aliases = ["settings", "set"])
async def pterodactyl_config(self, ctx: commands.Context):
async def pterodactyl_config(self, ctx: commands.Context) -> None:
"""Configure Pterodactyl settings."""
@pterodactyl_config.command(name = "url")
async def pterodactyl_config_base_url(self, ctx: commands.Context, *, base_url: str = None) -> None:
"""Set the base URL of your Pterodactyl Panel. Please include the protocol (http/https)."""
if base_url is None:
base_url = await self.config.base_url()
base_url = await config.base_url()
return await ctx.send(f"Base URL is currently set to {base_url}")
await self.config.base_url.set(base_url)
await config.base_url.set(base_url)
await ctx.send(f"Base URL set to {base_url}")
self.logger.info("Configuration value set: base_url = %s\nRestarting task...", base_url)
logger.info("Configuration value set: base_url = %s\nRestarting task...", base_url)
self.task.cancel()
self.task = self.get_task()
@pterodactyl_config.command(name = "apikey")
async def pterodactyl_config_api_key(self, ctx: commands.Context, *, api_key: str) -> None:
"""Set the API key for your Pterodactyl Panel."""
await self.config.api_key.set(api_key)
await config.api_key.set(api_key)
await ctx.send(f"API key set to `{api_key[:5]}...{api_key[-4:]}`")
self.logger.info("Configuration value set: api_key = %s\nRestarting task...", api_key)
logger.info("Configuration value set: api_key = %s\nRestarting task...", api_key)
self.task.cancel()
self.task = self.get_task()
@pterodactyl_config.command(name = "serverid")
async def pterodactyl_config_server_id(self, ctx: commands.Context, *, server_id: str) -> None:
"""Set the server ID for your Pterodactyl Panel."""
await self.config.server_id.set(server_id)
await config.server_id.set(server_id)
await ctx.send(f"Server ID set to {server_id}")
self.logger.info("Configuration value set: server_id = %s\nRestarting task...", server_id)
logger.info("Configuration value set: server_id = %s\nRestarting task...", server_id)
self.task.cancel()
self.task = self.get_task()
@pterodactyl_config.command(name = "consolechannel")
async def pterodactyl_config_console_channel(self, ctx: commands.Context, channel: discord.TextChannel) -> None:
"""Set the channel to send console output to."""
await self.config.console_channel.set(channel.id)
await config.console_channel.set(channel.id)
await ctx.send(f"Console channel set to {channel.mention}")
@pterodactyl_config.group(name = "chat")
@ -250,7 +124,7 @@ class Pterodactyl(commands.Cog):
@pterodactyl_config_chat.command(name = "channel")
async def pterodactyl_config_chat_channel(self, ctx: commands.Context, channel: discord.TextChannel) -> None:
"""Set the channel to send chat output to."""
await self.config.chat_channel.set(channel.id)
await config.chat_channel.set(channel.id)
await ctx.send(f"Chat channel set to {channel.mention}")
@pterodactyl_config_chat.command(name = "regex")
@ -260,9 +134,9 @@ class Pterodactyl(commands.Cog):
See [documentation]() for more information."""
#TODO - fix this link
if regex is None:
regex = await self.config.chat_regex()
regex = await config.chat_regex()
return await ctx.send(f"Chat regex is currently set to:\n{box(regex, 'regex')}")
await self.config.chat_regex.set(regex)
await config.chat_regex.set(regex)
await ctx.send(f"Chat regex set to:\n{box(regex, 'regex')}")
@pterodactyl_config_chat.command(name = "command")
@ -273,7 +147,7 @@ class Pterodactyl(commands.Cog):
See [documentation]() for more information."""
#TODO - fix this link
if command is None:
command = await self.config.chat_command()
command = await config.chat_command()
return await ctx.send(f"Chat command is currently set to:\n{box(command, 'json')}")
await self.config.chat_command.set(command)
await config.chat_command.set(command)
await ctx.send(f"Chat command set to:\n{box(command, 'json')}")

124
pterodactyl/websocket.py Normal file
View file

@ -0,0 +1,124 @@
import json
import re
from typing import Optional, Union
import aiohttp
import discord
import websockets
from pydactyl import PterodactylClient, exceptions
from redbot.core.utils.chat_formatting import pagify
from pterodactyl.config import config
from pterodactyl.logger import logger
from pterodactyl.pterodactyl import Pterodactyl
async def establish_websocket_connection(coginstance: Pterodactyl) -> None:
logger.debug("Establishing WebSocket connection")
base_url = await config.base_url()
api_key = await config.api_key()
server_id = await config.server_id()
try:
client = PterodactylClient(base_url, api_key, debug=True).client
coginstance.client = client
websocket_credentials = client.servers.get_websocket(server_id)
logger.debug("""Websocket connection details retrieved:
Socket: %s
Token: %s...""",
websocket_credentials['data']['socket'],
websocket_credentials['data']['token'][:20]
)
#NOTE - The token is truncated to prevent it from being logged in its entirety, for security reasons
except exceptions.ClientConfigError as e:
return logger.error('Failed to initialize Pterodactyl client: %s', e)
except exceptions.PterodactylApiError as e:
return logger.error('Failed to retrieve Pterodactyl websocket: %s', e)
async with websockets.connect(websocket_credentials['data']['socket'], origin=base_url, ping_timeout=60) as websocket:
logger.info("WebSocket connection established")
auth_message = json.dumps({"event": "auth", "args": [websocket_credentials['data']['token']]})
await websocket.send(auth_message)
logger.debug("Authentication message sent")
coginstance.websocket = websocket
current_status = ''
while True:
message = await websocket.recv()
if json.loads(message)['event'] in ('token expiring', 'token expired'):
logger.debug("Received token expiring/expired event. Refreshing token.")
websocket_credentials = client.servers.get_websocket(server_id)
auth_message = json.dumps({"event": "auth", "args": [websocket_credentials['data']['token']]})
await websocket.send(auth_message)
logger.debug("Authentication message sent")
if json.loads(message)['event'] == 'auth success':
logger.info("WebSocket authentication successful")
if json.loads(message)['event'] == 'console output' and await config.console_channel() is not None:
if current_status in ('running', 'offline', ''):
content = remove_ansi_escape_codes(json.loads(message)['args'][0])
channel = coginstance.bot.get_channel(await config.console_channel())
if channel is not None:
if content.startswith('['):
pagified_content = pagify(content, delims=[" ", "\n"])
for page in pagified_content:
await channel.send(content=page)
chat_message = await check_if_chat_message(content)
if chat_message:
info = await get_info(chat_message['username'])
if info is not None:
await send_chat_discord(coginstance, chat_message['username'], chat_message['message'], info['data']['player']['avatar'])
else:
await send_chat_discord(coginstance, chat_message['username'], chat_message['message'], 'https://seafsh.cc/u/j3AzqQ.png')
if json.loads(message)['event'] == 'status':
current_status = json.loads(message)['args'][0]
if await config.console_channel() is not None:
console = coginstance.bot.get_channel(await config.console_channel())
if console is not None:
await console.send(f"Server status changed! `{json.loads(message)['args'][0]}`")
def remove_ansi_escape_codes(text: str) -> str:
ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
#NOTE - https://chat.openai.com/share/d92f9acf-d776-4fd6-a53f-b14ac15dd540
return ansi_escape.sub('', text)
async def check_if_chat_message(text: str) -> Union[bool, dict]:
logger.debug("Checking if message is a chat message")
regex = await config.chat_regex()
match: Optional[re.Match[str]] = re.match(regex, text)
if match:
groups = {"time": match.group(1), "username": match.group(2), "message": match.group(3)}
logger.debug("Message is a chat message\n%s", json.dumps(groups))
return groups
logger.debug("Message is not a chat message")
return False
async def get_info(username: str) -> Optional[dict]:
logger.debug("Retrieving player info for %s", username)
endpoint = await config.api_endpoint()
async with aiohttp.ClientSession() as session:
async with session.get(f"https://playerdb.co/api/player/{endpoint}/{username}") as response:
if response.status == 200:
logger.debug("Player info retrieved for %s\n%s", username, json.dumps(await response.json()))
return await response.json()
logger.error("Failed to retrieve player info for %s: %s", username, response.status)
return None
async def send_chat_discord(coginstance: Pterodactyl, username: str, message: str, avatar_url: str) -> None:
logger.debug("Sending chat message to Discord")
channel = coginstance.bot.get_channel(await config.chat_channel())
if channel is not None:
webhooks = await channel.webhooks()
webhook = discord.utils.get(webhooks, name="Pterodactyl Chat")
if webhook is None:
webhook = await channel.create_webhook(name="Pterodactyl Chat")
await webhook.send(content=message, username=username, avatar_url=avatar_url)
logger.debug("Chat message sent to Discord")
else:
logger.debug("Chat channel not set. Skipping sending chat message to Discord")