diff options
| -rw-r--r-- | bot/__main__.py | 1 | ||||
| -rw-r--r-- | bot/cogs/rmq.py | 14 | ||||
| -rw-r--r-- | bot/cogs/snekbox.py | 67 |
3 files changed, 82 insertions, 0 deletions
diff --git a/bot/__main__.py b/bot/__main__.py index d04be725c..42ab0f224 100644 --- a/bot/__main__.py +++ b/bot/__main__.py @@ -67,6 +67,7 @@ bot.load_extension("bot.cogs.eval") bot.load_extension("bot.cogs.fun") bot.load_extension("bot.cogs.hiphopify") bot.load_extension("bot.cogs.snakes") +bot.load_extension("bot.cogs.snekbox") bot.load_extension("bot.cogs.tags") bot.load_extension("bot.cogs.verification") bot.load_extension("bot.cogs.utils") diff --git a/bot/cogs/rmq.py b/bot/cogs/rmq.py index 1b137334e..ea2d46707 100644 --- a/bot/cogs/rmq.py +++ b/bot/cogs/rmq.py @@ -4,6 +4,7 @@ import logging import pprint import aio_pika +from aio_pika import Message from dateutil import parser as date_parser from discord import Colour, Embed from discord.ext.commands import Bot @@ -54,6 +55,19 @@ class RMQ: with message.process(): await self.handle_message(message, message.body.decode()) + async def send_text(self, queue: str, data: str): + message = Message(data.encode("utf-8")) + await self.channel.default_exchange.publish(message, queue) + + async def send_json(self, queue: str, **data): + message = Message(json.dumps(data).encode("utf-8")) + await self.channel.default_exchange.publish(message, queue) + + async def consume(self, queue: str, callback, **kwargs): + queue_obj = await self.channel.declare_queue(queue, **kwargs) + + await queue_obj.consume(callback) + async def handle_message(self, message, data): log.debug(f"Message: {message}") log.debug(f"Data: {data}") diff --git a/bot/cogs/snekbox.py b/bot/cogs/snekbox.py new file mode 100644 index 000000000..55fbe9a4e --- /dev/null +++ b/bot/cogs/snekbox.py @@ -0,0 +1,67 @@ +import datetime +import logging + +from aio_pika import Message +from discord.ext.commands import Bot, Context, command + +from bot.cogs.rmq import RMQ +from bot.constants import Roles +from bot.decorators import with_role + +log = logging.getLogger(__name__) + +RMQ_ARGS = { + "durable": False, + "arguments": {"x-message-ttl": 5000}, + "auto_delete": True +} + + +class Snekbox: + """ + Safe evaluation using Snekbox + """ + + jobs = None # type: dict + + def __init__(self, bot: Bot): + self.bot = bot + self.jobs = {} + + @property + def rmq(self) -> RMQ: + return self.bot.get_cog("RMQ") + + @command(name="snekbox.eval()", aliases=["snekbox.eval"]) + @with_role(Roles.admin, Roles.owner, Roles.devops, Roles.moderator) + async def do_eval(self, ctx: Context, code: str): + if ctx.author.id in self.jobs: + await ctx.send(f"{ctx.author.mention} You've already got a job running - please wait for it to finish!") + return + + log.info(f"Received code from {ctx.author.name}#{ctx.author.discriminator} for evaluation:\n{code}") + self.jobs[ctx.author.id] = datetime.datetime.now() + + try: + await self.rmq.send_json( + "input", + snekid=str(ctx.author.id), message=code + ) + + async def callback(message: Message): + await ctx.send( + f"{ctx.author.mention}\n\n" + f"```{message.body.decode()}```" + ) + + del self.jobs[ctx.author.id] + + await self.rmq.consume(str(ctx.author.id), callback, **RMQ_ARGS) + except Exception: + del self.jobs[ctx.author.id] + raise + + +def setup(bot): + bot.add_cog(Snekbox(bot)) + log.info("Cog loaded: Snekbox") |