diff options
| -rw-r--r-- | bot/cogs/information.py | 5 | ||||
| -rw-r--r-- | bot/utils/checks.py | 48 |
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 |