forward console responses to commands

bot.py

@@ -1,7 +1,9 @@
+import asyncio
 import inspect
 import itertools
 import logging
 import os
+import re
 
 import discord
 from discord.ext import commands
@@ -13,18 +15,23 @@ import heavynode
 DISCORD_TOKEN = os.environ['discord_token']
 DISCORD_SERVER_ID = 530446700058509323
 HEAVYNODE_TOKEN = os.environ['heavynode_token']
+SESSION_COOKIE = os.environ['pterodactyl_session_cookie']
 
 
 logging.basicConfig()
 
-bot = lib.MineBot(command_prefix='!')
+intents = discord.Intents.default()
-hn = heavynode.Client(HEAVYNODE_TOKEN)
+intents.members = True
+bot = lib.MineBot(command_prefix='!', intents=intents)
+
+hn = heavynode.Client(HEAVYNODE_TOKEN, SESSION_COOKIE)
+bot.add_cleanup(hn.shutdown)
 
 
 async def is_admin(ctx):
     user = ctx.message.author
     guild = discord.utils.get(bot.guilds, id=DISCORD_SERVER_ID)
-    member = discord.utils.get(guild.members, id=user.id)
+    member = guild.get_member(user.id)
     if member is not None:
         for role in member.roles:
             if role.name == 'Admin' or role.name == 'Mod':
@@ -32,20 +39,27 @@ async def is_admin(ctx):
     return False
 
 
+async def whitelist_cmd(action, player):
+    expr = re.compile(f'\[Server thread/INFO\]: .*({player}|player).*(whitelist(ed)?|exist)', re.IGNORECASE)
+    resp = await hn.cmd_with_response(f'whitelist {action} {player}', expr)
+    _, msg = resp.split('[Server thread/INFO]: ')
+    return msg
+
+
 @bot.command()
 @commands.check(is_admin)
 async def add(ctx, player):
     """Add a player to the server whitelist. Must use exact Minecraft name."""
-    await hn.send_command(f'whitelist add {player}')
+    msg = await whitelist_cmd('add', player)
-    await ctx.send(f'"{player}" added to whitelist.')
+    await ctx.send(msg)
 
 
 @bot.command()
 @commands.check(is_admin)
 async def remove(ctx, player):
     """Remove a player from the server whitelist. Must use exact Minecraft name."""
-    await hn.send_command(f'whitelist remove {player}')
+    msg = await whitelist_cmd('remove', player)
-    await ctx.send(f'"{player}" removed from whitelist.')
+    await ctx.send(msg + '.')
 
 
 @add.error
@@ -53,8 +67,12 @@ async def remove(ctx, player):
 async def whitelist_error(ctx, error):
     if isinstance(error, commands.CheckFailure):
         await ctx.send('You must be a server admin to use this command.')
-    elif isinstance(error, heavynode.HttpError):
+    elif isinstance(error.original, (heavynode.HttpError, heavynode.GenericError)):
         await ctx.send('Failed to communicate with server.')
+    elif isinstance(error.original, asyncio.exceptions.TimeoutError):
+        await ctx.send('Command sent; response inconclusive. Check server output for more information.')
+    elif isinstance(error.original, ValueError):
+        await ctx.send('Command sent; could not interpret server response.')
     else:
         raise error
 

heavynode.py

@@ -1,60 +0,0 @@
-import asyncio
-import urllib.parse
-
-import aiohttp
-
-
-class HttpError(Exception):
-    def __init__(self, msg, response):
-        self.response = response
-        super().__init__(msg)
-
-
-class Client:
-    def __init__(self, token):
-        self.token = token
-        self.baseurl = 'https://control.heavynode.com/api'
-        # global state is icky, but it sure is convenient
-        loop = asyncio.get_event_loop()
-        loop.create_task(self.fetch_server())
-        # since we don't start the loop here, this will just
-        # hang out in a pending state until someone else does
-
-    async def make_request(self, method, path, *args, **kwargs):
-        h = {
-            'Authorization': f'Bearer {self.token}',
-            'Accept': 'Application/vnd.pterodactyl.v1+json',
-        }
-        if 'headers' in kwargs:
-            kwargs['headers'].update(h)
-        else:
-            kwargs['headers'] = h
-
-        if path[0] != '/': 
-            path = '/' + path
-        url = self.baseurl + path
-
-        # use context managers so connection is properly closed after request
-        # we don't make many requests so this is reasonable
-        async with aiohttp.ClientSession() as session:
-            async with session.request(method, url, *args, **kwargs) as r:
-                if r.status >= 400:
-                    raise HttpError(f'Request failed with status code {r.status}', r)
-                elif r.status == 204:
-                    return None # no content
-                elif r.headers['Content-Type'].lower() in {'application/json', 'application/vnd.pterodactyl.v1+json'}:
-                    return await r.json()
-                else:
-                    return await r.text
-
-    async def send_command(self, cmd):
-        """Send console command to minecraft server."""
-        server_id = self.server['identifier']
-        payload = {'command': cmd}
-        return await self.make_request('POST', f'/client/servers/{server_id}/command', json=payload)
-
-    async def fetch_server(self):
-        """Get the server to which we have access.
-        Assume there's only one."""
-        r = await self.make_request('GET', '/client')
-        self.server = r['data'][0]['attributes']

heavynode/__init__.py

@@ -0,0 +1 @@
+from .main import Client, HttpError, GenericError

heavynode/events.py

@@ -0,0 +1,38 @@
+import asyncio
+import contextlib
+
+
+class Beacon:
+    def __init__(self):
+        self.watchers = set()
+
+    @contextlib.contextmanager
+    def msg_q(self):
+        q = asyncio.Queue()
+        self.watchers.add(q)
+        try:
+            yield q
+        finally:
+            self.watchers.remove(q)
+
+    def flash(self, data):
+        for w in self.watchers:
+            w.put_nowait(data)
+
+    async def watch(self, count=float('inf')):
+        with self.msg_q() as q:
+            i = 0
+            while i < count:
+                yield await q.get()
+                i += 1
+
+    async def next_flash(self):
+        async for event in self.watch(1):
+            return event
+
+    async def watch_for(self, predicate, timeout=None):
+        async def waiter():
+            async for event in self.watch():
+                if predicate(event):
+                    return event
+        return await asyncio.wait_for(waiter(), timeout=timeout)

heavynode/main.py

@@ -0,0 +1,161 @@
+import asyncio
+import contextlib
+import collections
+import re
+import urllib.parse
+
+import aiohttp
+import socketio
+
+from . import events
+
+
+class HttpError(Exception):
+    def __init__(self, msg, response):
+        self.response = response
+        super().__init__(msg)
+
+class GenericError(Exception):
+    pass
+
+
+class Socket:
+    def __init__(self, secret, server_id, buffer_size=300):
+        self.secret = secret
+        self.server_id = server_id
+        self.buffer_size = buffer_size
+        # event hooks for observing
+        self.events = collections.defaultdict(events.Beacon)
+        self.buffers = collections.defaultdict(list)
+
+    async def handle_event(self, name, data):
+        buf = self.buffers[name]
+        buf.append(data)
+        if len(buf) > self.buffer_size:
+            del buf[0]
+        self.events[name].flash(data)
+
+    async def connect_socket(self):
+        socket = socketio.AsyncClient()
+        ns = '/v1/ws/' + self.server_id
+        
+        @socket.event
+        async def connect():
+            print('Websocket connected.')
+
+        @socket.event
+        async def disconnect():
+            print('Websocket disconnected.')
+
+        @socket.event(namespace=ns)
+        async def proc(data):
+            await self.handle_event('proc', data['data'])
+
+        @socket.event(namespace=ns)
+        async def console(data):
+            await self.handle_event('console', data['line'])
+
+        self.ws = socket
+        await socket.connect('https://mcp02.lax.us.heavynode.net:8443/socketio?token=' + self.secret, namespaces=[ns])
+
+    async def listen(self, event, expression, timeout=60):
+        if type(expression) == str:
+            regexp = re.compile(expression)
+        elif type(expression) == re.Pattern:
+            regexp = expression
+        else:
+            raise ValueError('Expression must be string or compiled regex.')
+
+        def test(line):
+            return regexp.search(line)
+        return await self.events[event].watch_for(test, timeout=timeout)
+
+    async def disconnect(self):
+        try:
+            await asyncio.wait_for(self.ws.disconnect(), timeout=5)
+        except asyncio.TimeoutError:
+            print('WARNING: Could not disconnect websocket cleanly.')
+
+
+class Client:
+    def __init__(self, token, session_cookie):
+        self.token = token
+        self.session_cookie = session_cookie
+        self.baseurl = 'https://control.heavynode.com/api'
+        self.stats = []
+        # global state is icky, but it sure is convenient
+        loop = asyncio.get_event_loop()
+        loop.create_task(self.setup_async())
+        # since we don't start the loop here, this will just
+        # hang out in a pending state until someone else does
+
+    async def make_request(self, method, path, *args, **kwargs):
+        h = {
+            'Authorization': f'Bearer {self.token}',
+            'Accept': 'Application/vnd.pterodactyl.v1+json',
+        }
+        if 'headers' in kwargs:
+            kwargs['headers'].update(h)
+        else:
+            kwargs['headers'] = h
+
+        if path[0] != '/': 
+            path = '/' + path
+        url = self.baseurl + path
+
+        # use context managers so connection is properly closed after request
+        # we don't make many requests so this is reasonable
+        async with aiohttp.ClientSession() as session:
+            try:
+                async with session.request(method, url, *args, **kwargs) as r:
+                    if r.status >= 400:
+                        raise HttpError(f'Request failed with status code {r.status}', r)
+                    elif r.status == 204:
+                        return None # no content
+                    elif r.headers['Content-Type'].lower() in {'application/json', 'application/vnd.pterodactyl.v1+json'}:
+                        return await r.json()
+                    else:
+                        return await r.text
+            except Exception as e:
+                raise GenericError('Something went wrong further down the stack. Original error shown above.') from e
+
+    async def send_command(self, cmd):
+        """Send console command to minecraft server."""
+        server_id = self.server['identifier']
+        payload = {'command': cmd}
+        return await self.make_request('POST', f'/client/servers/{server_id}/command', json=payload)
+
+    async def cmd_with_response(self, cmd, regex):
+        if type(regex) == str:
+            regex = re.compile(regex)
+        elif type(regex) != re.Pattern:
+            raise ValueError('Expression must be string or compiled regex.')
+
+        listener = asyncio.create_task(self.socket.listen('console', regex, timeout=5))
+        await self.send_command(cmd)
+        return await asyncio.wait_for(listener, timeout=None)
+
+
+    async def setup_async(self):
+        """Get the server to which we have access.
+        Assume there's only one."""
+        r = await self.make_request('GET', '/client')
+        self.server = r['data'][0]['attributes']
+        print('Added server:', self.server['identifier'])
+        
+        secret = await self.fetch_daemon_secret()
+        self.socket = Socket(secret, self.server['uuid'])
+        await self.socket.connect_socket()
+
+    async def fetch_daemon_secret(self):
+        cookie = {'pterodactyl_session': self.session_cookie}
+        async with aiohttp.ClientSession(cookies=cookie) as session:
+            r = await session.get('https://control.heavynode.com/server/' + self.server['identifier'])
+            m = re.search('"daemonSecret"\s?:\s?"([^"]*)"', await r.text())
+        if m != None:
+            return m.groups()[0]
+        else:
+            raise GenericError('Heavynode client setup failed: Could not acquire daemon secret.')
+
+    async def shutdown(self):
+        await self.socket.disconnect()