diff --git a/bot.py b/bot.py index db07c40..23d68ea 100644 --- a/bot.py +++ b/bot.py @@ -4,32 +4,18 @@ import os from discord.ext import commands +import lib import heavynode -class CustomBot(commands.Bot): - def __init__(self, *args, **kwargs): - self.cleanup = [] - super().__init__(*args, **kwargs) - - def add_cleanup(self, callback): - self.cleanup.append(callback) - - async def close(self): - for callback in self.cleanup: - r = callback() - # coroutines etc. return an awaitable object - if inspect.isawaitable(r): - await r - await super().close() +DISCORD_TOKEN = os.environ['discord_token'] +HEAVYNODE_TOKEN = os.environ['heavynode_token'] logging.basicConfig() -server_id = 'd030737c' -hn = heavynode.Client(os.environ['heavynode_token']) - -bot = CustomBot(command_prefix='!') +bot = lib.MineBot(command_prefix='!') +hn = heavynode.Client(HEAVYNODE_TOKEN) bot.add_cleanup(hn.shutdown) @@ -37,7 +23,7 @@ bot.add_cleanup(hn.shutdown) @commands.has_any_role('Admin', 'Mod') async def add(ctx, player): """Add a player to the server whitelist. Must use exact Minecraft name.""" - await hn.send_command(server_id, f'whitelist add {player}') + await hn.send_command(f'whitelist add {player}') await ctx.send(f'"{player}" added to whitelist.') @@ -45,7 +31,7 @@ async def add(ctx, player): @commands.has_any_role('Admin', 'Mod') async def remove(ctx, player): """Remove a player from the server whitelist. Must use exact Minecraft name.""" - await hn.send_command(server_id, f'whitelist remove {player}') + await hn.send_command(f'whitelist remove {player}') await ctx.send(f'"{player}" removed from whitelist.') @@ -54,8 +40,10 @@ 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): + await ctx.send('Failed to communicate with server.') else: raise error -bot.run(os.environ['discord_token'], reconnect=True) +bot.run(DISCORD_TOKEN, reconnect=True) diff --git a/heavynode.py b/heavynode.py index f48f858..466cb65 100644 --- a/heavynode.py +++ b/heavynode.py @@ -1,3 +1,4 @@ +import asyncio import urllib.parse import aiohttp @@ -14,6 +15,11 @@ class Client: self.token = token self.session = aiohttp.ClientSession() 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}'} @@ -29,11 +35,20 @@ class Client: r = await self.session.request(method, url, *args, **kwargs) if r.status >= 400: raise HttpError(f'Request failed with status code {r.status}', r) - return r + return await r.json() - async def send_command(self, serverid, cmd): + 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/{serverid}/command', json=payload) + 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'] async def shutdown(self): + """Gracefully terminate HTTP session for shutdown.""" await self.session.close() diff --git a/lib.py b/lib.py new file mode 100644 index 0000000..d68a4ed --- /dev/null +++ b/lib.py @@ -0,0 +1,20 @@ +import inspect + +from discord.ext import commands + + +class MineBot(commands.Bot): + def __init__(self, *args, **kwargs): + self.cleanup = [] + super().__init__(*args, **kwargs) + + def add_cleanup(self, callback): + self.cleanup.append(callback) + + async def close(self): + for callback in self.cleanup: + r = callback() + # coroutines etc. return an awaitable object + if inspect.isawaitable(r): + await r + await super().close() \ No newline at end of file