diff options
| -rw-r--r-- | bot/cogs/rmq.py | 5 | ||||
| -rw-r--r-- | bot/cogs/snekbox.py | 34 | 
2 files changed, 18 insertions, 21 deletions
| diff --git a/bot/cogs/rmq.py b/bot/cogs/rmq.py index ea2d46707..92645457e 100644 --- a/bot/cogs/rmq.py +++ b/bot/cogs/rmq.py @@ -63,10 +63,9 @@ class RMQ:          message = Message(json.dumps(data).encode("utf-8"))          await self.channel.default_exchange.publish(message, queue) -    async def consume(self, queue: str, callback, **kwargs): +    async def consume(self, queue: str, **kwargs):          queue_obj = await self.channel.declare_queue(queue, **kwargs) - -        await queue_obj.consume(callback) +        return await queue_obj.get(timeout=30)      async def handle_message(self, message, data):          log.debug(f"Message: {message}") diff --git a/bot/cogs/snekbox.py b/bot/cogs/snekbox.py index e228608b8..5b061f926 100644 --- a/bot/cogs/snekbox.py +++ b/bot/cogs/snekbox.py @@ -53,29 +53,27 @@ class Snekbox:                  snekid=str(ctx.author.id), message=code              ) -            async def callback(message: Message): -                output = message.body.decode() +            message = await self.rmq.consume(str(ctx.author.id), **RMQ_ARGS) +            output = message.body.decode() -                if "```" in output: -                    output = "Code block escape attempt detected; will not output result" -                else: -                    output = [f"{i:03d} | {line}" for i, line in enumerate(output.split("\n"), start=1)] -                    output = "\n".join(output) +            if "```" in output: +                output = "Code block escape attempt detected; will not output result" +            else: +                output = [f"{i:03d} | {line}" for i, line in enumerate(output.split("\n"), start=1)] +                output = "\n".join(output) -                    if len(output) >= 1900: -                        output = f"{output[:1900]}... (truncated)" +                if len(output) >= 1900: +                    output = f"{output[:1900]}... (truncated)" -                await ctx.send( -                    f"{ctx.author.mention} Your eval job has completed.\n\n```{output}```" -                ) - -                await ctx.send( -                    f"{ctx.author.mention} Your eval job has completed." -                ) +            await ctx.send( +                f"{ctx.author.mention} Your eval job has completed.\n\n```{output}```" +            ) -                del self.jobs[ctx.author.id] +            await ctx.send( +                f"{ctx.author.mention} Your eval job has completed." +            ) -            await self.rmq.consume(str(ctx.author.id), callback, **RMQ_ARGS) +            del self.jobs[ctx.author.id]          except Exception:              del self.jobs[ctx.author.id]              raise | 
