SeaCogs/pterodactyl/pterodactyl.py

import json
import logging

import websockets
from pydactyl import PterodactylClient, exceptions
from redbot.core import Config, commands
from redbot.core.bot import Red


class Pterodactyl(commands.Cog):
    """Pterodactyl allows you to manage your Pterodactyl Panel from Discord."""

    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,
            startup_jar=None,
            startup_arguments=None,
            power_action_in_progress=False
        )
        self.logger = logging.getLogger('red.sea.pterodactyl')
        self.client = None
        self.websocket = None

    async def establish_websocket_connection(self):
        base_url = await self.config.base_url()
        api_key = await self.config.api_key()
        server_id = await self.config.server_id()
        try:
            client = PterodactylClient(base_url, api_key).client
            websocket_credentials = client.servers.get_websocket(server_id)
            self.logger.debug("Websocket connection details retrieved: " + websocket_credentials)
        except exceptions.ClientConfigError as e:
            self.logger.error(f'Failed to initialize Pterodactyl client: {e}')
            return
        except exceptions.PterodactylApiError as e:
            self.logger.error(f'Failed to retrieve Pterodactyl websocket: {e}')
            return

        async with websockets.connect(websocket_credentials['data']['socket']) as websocket:
            self.logger.debug("WebSocket connection established")

            # Send authentication token
            auth_message = json.dumps({"event": "auth", "args": [websocket_credentials['data']['token']]})
            await websocket.send(auth_message)
            self.logger.debug("Authentication message sent")

            self.client = client
            self.websocket = websocket

            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.debug("Authentication successful")
                self.logger.debug("Received message: %s", message)

    async def cog_load(self):
        await self.establish_websocket_connection()

    async def cog_unload(self):
       await self.client._session.close()