Add a method to check our restrictions

cogs/utils/checks.py

@@ -4,6 +4,7 @@ import rethinkdb as r
 from discord.ext import commands
 import discord
 from . import config
+from . import utilities
 
 loop = asyncio.get_event_loop()
 
@@ -77,6 +78,85 @@ def should_ignore(bot, message):
     return str(message.author.id) in ignored['members'] or str(message.channel.id) in ignored['channels']
 
 
+def check_restricted():
+    async def predicate(ctx):
+        # Return true if this is a private channel, we'll handle that in the registering of the command
+        if type(ctx.message.channel) is discord.DMChannel:
+            return True
+
+        # First get all the restrictions
+        restrictions = ctx.bot.db.load('server_settings', key=ctx.message.guild.id, pluck='restrictions') or {}
+        # Now lets check the "from" restrictions
+        for from_restriction in restrictions.get('from', []):
+            # Get the source and destination
+            # Source should ALWAYS be a command in this case
+            source = from_restriction.get('source')
+            destination = from_restriction.get('destination')
+            # Convert destination to the object we want
+            destination = await utilities.convert(ctx, destination)
+            # If we couldn't find the destination, just continue with other restrictions
+            # Also if this restriction we're checking isn't for this command
+            if destination is None or source != ctx.command.qualified_name:
+                continue
+
+            # This means that the type of restriction we have is `command from channel`
+            # Which means we do not want commands to be ran in this channel
+            if destination == ctx.message.channel:
+                return False
+            # This type is `command from Role` meaning anyone with this role can't run this command
+            elif destination in ctx.message.author.roles:
+                return False
+            # This is `command from Member` meaning this user specifically cannot run this command
+            elif destination == ctx.message.author:
+                return False
+
+        # If we are here, then there are no blacklists stopping this from running
+
+        # Now for the to restrictions this is a little different, we need to make a whitelist and
+        # see if our current channel is in this whitelist, as well as any whitelisted roles are in the author's roles
+        # Only if there is no whitelist, do we want to blanket return True
+        to_restrictions = restrictions.get('to', [])
+        if len(to_restrictions) == 0:
+            return True
+
+        # Otherwise there is a whitelist, and we need to start it
+        whitelisted_channels = []
+        whitelisted_roles = []
+
+        for to_restriction in to_restrictions:
+            # Get the source and destination
+            # Source should ALWAYS be a command in this case
+            source = to_restriction.get('source')
+            destination = to_restriction.get('destination')
+            # Convert destination to the object we want
+            destination = await utilities.convert(ctx, destination)
+            # If we couldn't find the destination, just continue with other restrictions
+            # Also if this restriction we're checking isn't for this command
+            if destination is None or source != ctx.command.qualified_name:
+                continue
+
+            # Append to our two whitelists depending on what type this is
+            if isinstance(destination, discord.TextChannel):
+                whitelisted_channels.append(destination)
+            elif isinstance(destination, discord.Role):
+                whitelisted_roles.append(destination)
+
+        if whitelisted_channels:
+            if ctx.channel not in whitelisted_channels:
+                return False
+        if whitelisted_roles:
+            if not any(x in ctx.message.author.roles for x in whitelisted_roles):
+                return False
+
+        # If we have passed all of these, then we are allowed to run this command
+        # This looks like a whole lot, but all of these lists will be very tiny in almost all cases
+        # And only delving deep into the specific lists that may be large, will we finally see "large" lists
+        # Which means this still will not be slow in other cases
+        return True
+
+    return commands.check(predicate)
+
+
 def custom_perms(**perms):
     def predicate(ctx):
         # Return true if this is a private channel, we'll handle that in the registering of the command