Expand in_channel check to accept multi-channels, bypass roles.

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.