diff options
| -rw-r--r-- | bot/cogs/snekbox.py | 26 | ||||
| -rw-r--r-- | bot/cogs/utils.py | 16 | ||||
| -rw-r--r-- | bot/decorators.py | 44 |
3 files changed, 53 insertions, 33 deletions
diff --git a/bot/cogs/snekbox.py b/bot/cogs/snekbox.py index 1b51da843..cb0454249 100644 --- a/bot/cogs/snekbox.py +++ b/bot/cogs/snekbox.py @@ -6,12 +6,12 @@ import textwrap from discord import Colour, Embed from discord.ext.commands import ( - Bot, CommandError, Context, MissingPermissions, - NoPrivateMessage, check, command, guild_only + Bot, CommandError, Context, NoPrivateMessage, command, guild_only ) from bot.cogs.rmq import RMQ from bot.constants import Channels, ERROR_REPLIES, NEGATIVE_REPLIES, Roles, URLs +from bot.decorators import InChannelCheckFailure, in_channel from bot.utils.messages import wait_for_deletion @@ -51,22 +51,8 @@ RAW_CODE_REGEX = re.compile( r"\s*$", # any trailing whitespace until the end of the string re.DOTALL # "." also matches newlines ) -BYPASS_ROLES = (Roles.owner, Roles.admin, Roles.moderator, Roles.helpers) -WHITELISTED_CHANNELS = (Channels.bot,) -WHITELISTED_CHANNELS_STRING = ', '.join(f"<#{channel_id}>" for channel_id in WHITELISTED_CHANNELS) - - -async def channel_is_whitelisted_or_author_can_bypass(ctx: Context): - """ - Checks that the author is either helper or above - or the channel is a whitelisted channel. - """ - if ctx.channel.id in WHITELISTED_CHANNELS: - return True - if any(r.id in BYPASS_ROLES for r in ctx.author.roles): - return True - raise MissingPermissions("You are not allowed to do that here.") +BYPASS_ROLES = (Roles.owner, Roles.admin, Roles.moderator, Roles.helpers) class Snekbox: @@ -84,7 +70,7 @@ class Snekbox: @command(name='eval', aliases=('e',)) @guild_only() - @check(channel_is_whitelisted_or_author_can_bypass) + @in_channel(Channels.bot, bypass_roles=BYPASS_ROLES) async def eval_command(self, ctx: Context, *, code: str = None): """ Run some code. get the result back. We've done our best to make this safe, but do let us know if you @@ -205,9 +191,9 @@ class Snekbox: embed.description = "You're not allowed to use this command in private messages." await ctx.send(embed=embed) - elif isinstance(error, MissingPermissions): + elif isinstance(error, InChannelCheckFailure): embed.title = random.choice(NEGATIVE_REPLIES) - embed.description = f"Sorry, but you may only use this command within {WHITELISTED_CHANNELS_STRING}." + embed.description = str(error) await ctx.send(embed=embed) else: diff --git a/bot/cogs/utils.py b/bot/cogs/utils.py index 94f3938e4..65c729414 100644 --- a/bot/cogs/utils.py +++ b/bot/cogs/utils.py @@ -1,4 +1,5 @@ import logging +import random import re import unicodedata from email.parser import HeaderParser @@ -7,11 +8,13 @@ from io import StringIO from discord import Colour, Embed from discord.ext.commands import AutoShardedBot, Context, command -from bot.constants import Roles -from bot.decorators import with_role +from bot.constants import Channels, NEGATIVE_REPLIES, Roles +from bot.decorators import InChannelCheckFailure, in_channel log = logging.getLogger(__name__) +BYPASS_ROLES = (Roles.owner, Roles.admin, Roles.moderator, Roles.helpers) + class Utils: """ @@ -25,7 +28,6 @@ class Utils: self.base_github_pep_url = "https://raw.githubusercontent.com/python/peps/master/pep-" @command(name='pep', aliases=('get_pep', 'p')) - @with_role(Roles.verified) async def pep_command(self, ctx: Context, pep_number: str): """ Fetches information about a PEP and sends it to the channel. @@ -89,6 +91,7 @@ class Utils: await ctx.message.channel.send(embed=pep_embed) @command() + @in_channel(Channels.bot, bypass_roles=BYPASS_ROLES) async def charinfo(self, ctx, *, characters: str): """ Shows you information on up to 25 unicode characters. @@ -132,6 +135,13 @@ class Utils: await ctx.send(embed=embed) + async def __error(self, ctx, error): + embed = Embed(colour=Colour.red()) + if isinstance(error, InChannelCheckFailure): + embed.title = random.choice(NEGATIVE_REPLIES) + embed.description = str(error) + await ctx.send(embed=embed) + def setup(bot): bot.add_cog(Utils(bot)) diff --git a/bot/decorators.py b/bot/decorators.py index fe974cbd3..87877ecbf 100644 --- a/bot/decorators.py +++ b/bot/decorators.py @@ -1,18 +1,51 @@ import logging import random +import typing from asyncio import Lock from functools import wraps from weakref import WeakValueDictionary from discord import Colour, Embed from discord.ext import commands -from discord.ext.commands import Context +from discord.ext.commands import CheckFailure, Context from bot.constants import ERROR_REPLIES log = logging.getLogger(__name__) +class InChannelCheckFailure(CheckFailure): + pass + + +def in_channel(*channels: int, bypass_roles: typing.Container[int] = None): + """ + Checks that the message is in a whitelisted channel or optionally has a bypass role. + """ + def predicate(ctx: Context): + if ctx.channel.id in channels: + log.debug(f"{ctx.author} tried to call the '{ctx.command.name}' command. " + f"The command was used in a whitelisted channel.") + return True + + if bypass_roles: + if any(r.id in bypass_roles for r in ctx.author.roles): + log.debug(f"{ctx.author} tried to call the '{ctx.command.name}' command. " + f"The command was not used in a whitelisted channel, " + f"but the author had a role to bypass the in_channel check.") + return True + + log.debug(f"{ctx.author} tried to call the '{ctx.command.name}' command. " + f"The in_channel check failed.") + + channels_str = ', '.join(f"<#{c_id}>" for c_id in channels) + raise InChannelCheckFailure( + f"Sorry, but you may only use this command within {channels_str}." + ) + + return commands.check(predicate) + + def with_role(*role_ids: int): async def predicate(ctx: Context): if not ctx.guild: # Return False in a DM @@ -46,15 +79,6 @@ def without_role(*role_ids: int): return commands.check(predicate) -def in_channel(channel_id): - async def predicate(ctx: Context): - check = ctx.channel.id == channel_id - log.debug(f"{ctx.author} tried to call the '{ctx.command.name}' command. " - f"The result of the in_channel check was {check}.") - return check - return commands.check(predicate) - - def locked(): """ Allows the user to only run one instance of the decorated command at a time. |