rewrite override_in_channel to accept an optional new channel whitelist

add checks for the new whitelist to in_channel_check

1 files changed, 29 insertions, 10 deletions

diff --git a/bot/decorators.py b/bot/decorators.py
index dbaad4a2..d7ce29cb 100644
--- a/bot/decorators.py
+++ b/bot/decorators.py
@@ -64,12 +64,13 @@ def without_role(*role_ids: int) -> bool:
 
 
 def in_channel_check(*channels: int, bypass_roles: typing.Container[int] = None) -> typing.Callable[[Context], bool]:
-    """Checks that the message is in a whitelisted channel or optionally has a bypass role."""
+    """Checks that the message is in a whitelisted channel or optionally has a bypass role.
+
+       If `in_channel_override` is present, check if it contains channels and use them in place of global whitelist"""
     def predicate(ctx: Context) -> bool:
         if not ctx.guild:
             log.debug(f"{ctx.author} tried to use the '{ctx.command.name}' command from a DM.")
             return True
-
         if ctx.channel.id in channels:
             log.debug(
                 f"{ctx.author} tried to call the '{ctx.command.name}' command "
@@ -78,11 +79,24 @@ def in_channel_check(*channels: int, bypass_roles: typing.Container[int] = None)
             return True
 
         if hasattr(ctx.command.callback, "in_channel_override"):
-            log.debug(
-                f"{ctx.author} called the '{ctx.command.name}' command "
-                f"and the command was whitelisted to bypass the in_channel check."
-            )
-            return True
+            override = ctx.command.callback.in_channel_override
+            if override is None:
+                log.debug(
+                    f"{ctx.author} called the '{ctx.command.name}' command "
+                    f"and the command was whitelisted to bypass the in_channel check."
+                )
+                return True
+            else:
+                if ctx.channel.id in override:
+                    log.debug(
+                        f"{ctx.author} tried to call the '{ctx.command.name}' command "
+                        f"and the command was used in an overridden whitelisted channel."
+                    )
+                    return True
+                channels_str = ', '.join(f"<#{c_id}>" for c_id in override)
+                raise InChannelCheckFailure(
+                    f"Sorry, but you may only use this command within {channels_str}."
+                )
 
         if bypass_roles and any(r.id in bypass_roles for r in ctx.author.roles):
             log.debug(
@@ -107,14 +121,19 @@ def in_channel_check(*channels: int, bypass_roles: typing.Container[int] = None)
 in_channel = commands.check(in_channel_check)
 
 
-def override_in_channel(func: typing.Callable) -> typing.Callable:
+def override_in_channel(channels: typing.Tuple[int] = None) -> typing.Callable:
     """
     Set command callback attribute for detection in `in_channel_check`.
 
+    Override global whitelist if channels are specified.
+
     This decorator has to go before (below) below the `command` decorator.
     """
-    func.in_channel_override = True
-    return func
+    def inner(func: typing.Callable):
+        func.in_channel_override = channels
+        return func
+
+    return inner
 
 
 def locked() -> typing.Union[typing.Callable, None]: