Implement a bypassable cooldown decorator

2 files changed, 50 insertions, 3 deletions

diff --git a/bot/cogs/information.py b/bot/cogs/information.py
index b6a3c4a40..3a7ba0444 100644
--- a/bot/cogs/information.py
+++ b/bot/cogs/information.py
@@ -8,11 +8,11 @@ from typing import Any, Mapping, Optional
 import discord
 from discord import CategoryChannel, Colour, Embed, Member, Role, TextChannel, VoiceChannel, utils
 from discord.ext import commands
-from discord.ext.commands import Bot, Cog, Context, command, group
+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, in_channel, with_role
-from bot.utils.checks import with_role_check
+from bot.utils.checks import cooldown_with_role_bypass, with_role_check
 from bot.utils.time import time_since
 
 log = logging.getLogger(__name__)
@@ -268,6 +268,7 @@ class Information(Cog):
         # 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:
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