minebot/heavynode/main.py

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, cookie_name, cookie_value):
        self.token = token
        self.cookie_name = cookie_name
        self.cookie_value = cookie_value
        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 = {self.cookie_name: self.cookie_value}
        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()