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()