import asyncio import json from typing import Optional import discord import websockets from pydactyl import PterodactylClient from redbot.core import commands from redbot.core.bot import Red from redbot.core.utils.chat_formatting import box from pterodactyl.config import config, register_config from pterodactyl.logger import logger class Pterodactyl(commands.Cog): """Pterodactyl allows you to manage your Pterodactyl Panel from Discord.""" def __init__(self, bot: Red): self.bot = bot self.client: Optional[PterodactylClient] = None self.task: Optional[asyncio.Task] = None self.websocket: Optional[websockets.WebSocketClientProtocol] = None register_config(config) async def cog_load(self) -> None: self.task = self.get_task() async def cog_unload(self) -> None: self.task.cancel() await self.client._session.close() # pylint: disable=protected-access def get_task(self) -> asyncio.Task: from pterodactyl.websocket import establish_websocket_connection 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) -> 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 logger.error("WebSocket task has failed: %s", e, exc_info=e) self.task.cancel() self.task = self.get_task() @commands.Cog.listener() 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: logger.error("WebSocket connection closed: %s", e) self.task.cancel() self.task = self.get_task() 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)]}) logger.debug("Sending chat message to server:\n%s", msg) try: await self.websocket.send(msg) except websockets.exceptions.ConnectionClosed as 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) -> 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) -> 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 config.base_url() return await ctx.send(f"Base URL is currently set to {base_url}") await config.base_url.set(base_url) await ctx.send(f"Base URL set to {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 config.api_key.set(api_key) await ctx.send(f"API key set to `{api_key[:5]}...{api_key[-4:]}`") 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 config.server_id.set(server_id) await ctx.send(f"Server ID set to {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 config.console_channel.set(channel.id) await ctx.send(f"Console channel set to {channel.mention}") @pterodactyl_config.group(name = "chat") async def pterodactyl_config_chat(self, ctx: commands.Context): """Configure chat settings.""" @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 config.chat_channel.set(channel.id) await ctx.send(f"Chat channel set to {channel.mention}") @pterodactyl_config_chat.command(name = "regex") async def pterodactyl_config_chat_regex(self, ctx: commands.Context, *, regex: str = None) -> None: """Set the regex pattern to match chat messages on the server. See [documentation]() for more information.""" #TODO - fix this link if regex is None: regex = await config.chat_regex() return await ctx.send(f"Chat regex is currently set to:\n{box(regex, '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") async def pterodactyl_config_chat_command(self, ctx: commands.Context, *, command: str = None) -> None: """Set the command that will be used to send messages from Discord. Required placeholders: `.$U` (username), `.$M` (message), `.$C` (color) See [documentation]() for more information.""" #TODO - fix this link if command is None: command = await config.chat_command() return await ctx.send(f"Chat command is currently set to:\n{box(command, 'json')}") await config.chat_command.set(command) await ctx.send(f"Chat command set to:\n{box(command, 'json')}") @pterodactyl_config_chat.command(name = "server") async def pterodactyl_config_chat_server(self, ctx: commands.Context, *, regex: str = None) -> None: """Set the regex pattern to match server messages on the server. See [documentation]() for more information.""" #TODO - fix this link if regex is None: regex = await config.server_regex() return await ctx.send(f"Server regex is currently set to:\n{box(regex, 'regex')}") await config.server_regex.set(regex) await ctx.send(f"Server regex set to:\n{box(regex, 'regex')}")