SeaCogs/pterodactyl/pterodactyl.py

190 lines
9.4 KiB
Python

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')}")
@pterodactyl_config.group(name = "messages", aliases = ['msg', 'msgs', 'message'])
async def pterodactyl_config_messages(self, ctx: commands.Context):
"""Configure message settings."""
@pterodactyl_config_messages.command(name = "startup")
async def pterodactyl_config_messages_startup(self, ctx: commands.Context, *, message: str = None) -> None:
"""Set the message that will be sent when the server starts."""
if message is None:
message = await config.startup_msg()
return await ctx.send(f"Startup message is currently set to: {message}")
await config.startup_msg.set(message)
await ctx.send(f"Startup message set to: {message}")
@pterodactyl_config_messages.command(name = "shutdown")
async def pterodactyl_config_messages_shutdown(self, ctx: commands.Context, *, message: str = None) -> None:
"""Set the message that will be sent when the server stops."""
if message is None:
message = await config.shutdown_msg()
return await ctx.send(f"Shutdown message is currently set to: {message}")
await config.shutdown_msg.set(message)
await ctx.send(f"Shutdown message set to: {message}")