Merge pull request #503 from avayert/master

Add raw command

2 files changed, 130 insertions, 4 deletions

diff --git a/bot/cogs/information.py b/bot/cogs/information.py
index 1afb37103..3a7ba0444 100644
--- a/bot/cogs/information.py
+++ b/bot/cogs/information.py
@@ -1,14 +1,18 @@
 import colorsys
 import logging
+import pprint
 import textwrap
 import typing
+from typing import Any, Mapping, Optional
 
+import discord
 from discord import CategoryChannel, Colour, Embed, Member, Role, TextChannel, VoiceChannel, utils
-from discord.ext.commands import Bot, Cog, Context, command
+from discord.ext import commands
+from discord.ext.commands import Bot, BucketType, Cog, Context, command, group
 
 from bot.constants import Channels, Emojis, MODERATION_ROLES, STAFF_ROLES
-from bot.decorators import InChannelCheckFailure, with_role
-from bot.utils.checks import with_role_check
+from bot.decorators import InChannelCheckFailure, in_channel, with_role
+from bot.utils.checks import cooldown_with_role_bypass, with_role_check
 from bot.utils.time import time_since
 
 log = logging.getLogger(__name__)
@@ -229,6 +233,82 @@ class Information(Cog):
 
         await ctx.send(embed=embed)
 
+    def format_fields(self, mapping: Mapping[str, Any], field_width: Optional[int] = None) -> str:
+        """Format a mapping to be readable to a human."""
+        # sorting is technically superfluous but nice if you want to look for a specific field
+        fields = sorted(mapping.items(), key=lambda item: item[0])
+
+        if field_width is None:
+            field_width = len(max(mapping.keys(), key=len))
+
+        out = ''
+
+        for key, val in fields:
+            if isinstance(val, dict):
+                # if we have dicts inside dicts we want to apply the same treatment to the inner dictionaries
+                inner_width = int(field_width * 1.6)
+                val = '\n' + self.format_fields(val, field_width=inner_width)
+
+            elif isinstance(val, str):
+                # split up text since it might be long
+                text = textwrap.fill(val, width=100, replace_whitespace=False)
+
+                # indent it, I guess you could do this with `wrap` and `join` but this is nicer
+                val = textwrap.indent(text, ' ' * (field_width + len(': ')))
+
+                # the first line is already indented so we `str.lstrip` it
+                val = val.lstrip()
+
+            if key == 'color':
+                # makes the base 10 representation of a hex number readable to humans
+                val = hex(val)
+
+            out += '{0:>{width}}: {1}\n'.format(key, val, width=field_width)
+
+        # remove trailing whitespace
+        return out.rstrip()
+
+    @cooldown_with_role_bypass(2, 60 * 3, BucketType.member, bypass_roles=STAFF_ROLES)
+    @group(invoke_without_command=True)
+    @in_channel(Channels.bot, bypass_roles=STAFF_ROLES)
+    async def raw(self, ctx: Context, *, message: discord.Message, json: bool = False) -> None:
+        """Shows information about the raw API response."""
+        # I *guess* it could be deleted right as the command is invoked but I felt like it wasn't worth handling
+        # doing this extra request is also much easier than trying to convert everything back into a dictionary again
+        raw_data = await ctx.bot.http.get_message(message.channel.id, message.id)
+
+        paginator = commands.Paginator()
+
+        def add_content(title: str, content: str) -> None:
+            paginator.add_line(f'== {title} ==\n')
+            # replace backticks as it breaks out of code blocks. Spaces seemed to be the most reasonable solution.
+            # we hope it's not close to 2000
+            paginator.add_line(content.replace('```', '`` `'))
+            paginator.close_page()
+
+        if message.content:
+            add_content('Raw message', message.content)
+
+        transformer = pprint.pformat if json else self.format_fields
+        for field_name in ('embeds', 'attachments'):
+            data = raw_data[field_name]
+
+            if not data:
+                continue
+
+            total = len(data)
+            for current, item in enumerate(data, start=1):
+                title = f'Raw {field_name} ({current}/{total})'
+                add_content(title, transformer(item))
+
+        for page in paginator.pages:
+            await ctx.send(page)
+
+    @raw.command()
+    async def json(self, ctx: Context, message: discord.Message) -> None:
+        """Shows information about the raw API response in a copy-pasteable Python format."""
+        await ctx.invoke(self.raw, message=message, json=True)
+
 
 def setup(bot: Bot) -> None:
     """Information cog load."""
diff --git a/bot/utils/checks.py b/bot/utils/checks.py
index 19f64ff9f..ad892e512 100644
--- a/bot/utils/checks.py
+++ b/bot/utils/checks.py
@@ -1,6 +1,8 @@
+import datetime
 import logging
+from typing import Callable, Iterable
 
-from discord.ext.commands import Context
+from discord.ext.commands import BucketType, Cog, Command, CommandOnCooldown, Context, Cooldown, CooldownMapping
 
 log = logging.getLogger(__name__)
 
@@ -42,3 +44,47 @@ def in_channel_check(ctx: Context, channel_id: int) -> bool:
     log.trace(f"{ctx.author} tried to call the '{ctx.command.name}' command. "
               f"The result of the in_channel check was {check}.")
     return check
+
+
+def cooldown_with_role_bypass(rate: int, per: float, type: BucketType = BucketType.default, *,
+                              bypass_roles: Iterable[int]) -> Callable:
+    """
+    Applies a cooldown to a command, but allows members with certain roles to be ignored.
+
+    NOTE: this replaces the `Command.before_invoke` callback, which *might* introduce problems in the future.
+    """
+    # make it a set so lookup is hash based
+    bypass = set(bypass_roles)
+
+    # this handles the actual cooldown logic
+    buckets = CooldownMapping(Cooldown(rate, per, type))
+
+    # will be called after the command has been parse but before it has been invoked, ensures that
+    # the cooldown won't be updated if the user screws up their input to the command
+    async def predicate(cog: Cog, ctx: Context) -> None:
+        nonlocal bypass, buckets
+
+        if any(role.id in bypass for role in ctx.author.roles):
+            return
+
+        # cooldown logic, taken from discord.py internals
+        current = ctx.message.created_at.replace(tzinfo=datetime.timezone.utc).timestamp()
+        bucket = buckets.get_bucket(ctx.message)
+        retry_after = bucket.update_rate_limit(current)
+        if retry_after:
+            raise CommandOnCooldown(bucket, retry_after)
+
+    def wrapper(command: Command) -> Command:
+        # NOTE: this could be changed if a subclass of Command were to be used. I didn't see the need for it
+        # so I just made it raise an error when the decorator is applied before the actual command object exists.
+        #
+        # if the `before_invoke` detail is ever a problem then I can quickly just swap over.
+        if not isinstance(command, Command):
+            raise TypeError('Decorator `cooldown_with_role_bypass` must be applied after the command decorator. '
+                            'This means it has to be above the command decorator in the code.')
+
+        command._before_invoke = predicate
+
+        return command
+
+    return wrapper