[Core] Command disable feature (#2099)

* [Core] Command disable feature

Signed-off-by: Toby Harradine <tobyharradine@gmail.com>

* [Core] Allow user to set the "command disabled" message

Signed-off-by: Toby Harradine <tobyharradine@gmail.com>

* Reformat

Signed-off-by: Toby Harradine <tobyharradine@gmail.com>

redbot/core/bot.py

@@ -58,6 +58,8 @@ class RedBase(commands.GroupMixin, commands.bot.BotBase, RPCMixin):
             help__page_char_limit=1000,
             help__max_pages_in_guild=2,
             help__tagline="",
+            disabled_commands=[],
+            disabled_command_msg="That command is disabled.",
         )
 
         self.db.register_guild(
@@ -69,6 +71,7 @@ class RedBase(commands.GroupMixin, commands.bot.BotBase, RPCMixin):
             embeds=None,
             use_bot_color=False,
             fuzzy=False,
+            disabled_commands=[],
         )
 
         self.db.register_user(embeds=None)
@@ -340,6 +343,13 @@ class RedBase(commands.GroupMixin, commands.bot.BotBase, RPCMixin):
                 )
         super().add_cog(cog)
 
+    def add_command(self, command: commands.Command):
+        if not isinstance(command, commands.Command):
+            raise TypeError("Command objects must derive from redbot.core.commands.Command")
+
+        super().add_command(command)
+        self.dispatch("command_add", command)
+
 
 class Red(RedBase, discord.AutoShardedClient):
     """

redbot/core/commands/commands.py

@@ -4,8 +4,10 @@ This module contains extended classes and functions which are intended to
 replace those from the `discord.ext.commands` module.
 """
 import inspect
-from typing import TYPE_CHECKING
+import weakref
+from typing import Awaitable, Callable, TYPE_CHECKING
 
+import discord
 from discord.ext import commands
 
 from .errors import ConversionFailure
@@ -104,6 +106,49 @@ class Command(commands.Command):
             # We should expose anything which might be a bug in the converter
             raise exc
 
+    def disable_in(self, guild: discord.Guild) -> bool:
+        """Disable this command in the given guild.
+
+        Parameters
+        ----------
+        guild : discord.Guild
+            The guild to disable the command in.
+
+        Returns
+        -------
+        bool
+            ``True`` if the command wasn't already disabled.
+
+        """
+        disabler = get_command_disabler(guild)
+        if disabler in self.checks:
+            return False
+        else:
+            self.checks.append(disabler)
+            return True
+
+    def enable_in(self, guild: discord.Guild) -> bool:
+        """Enable this command in the given guild.
+
+        Parameters
+        ----------
+        guild : discord.Guild
+            The guild to enable the command in.
+
+        Returns
+        -------
+        bool
+            ``True`` if the command wasn't already enabled.
+
+        """
+        disabler = get_command_disabler(guild)
+        try:
+            self.checks.remove(disabler)
+        except ValueError:
+            return False
+        else:
+            return True
+
 
 class GroupMixin(commands.GroupMixin):
     """Mixin for `Group` and `Red` classes.
@@ -162,6 +207,12 @@ class Group(GroupMixin, Command, commands.Group):
             if self.autohelp and not self.invoke_without_command:
                 await self._verify_checks(ctx)
                 await ctx.send_help()
+        elif self.invoke_without_command:
+            # So invoke_without_command when a subcommand of this group is invoked
+            # will skip the the invokation of *this* command. However, because of
+            # how our permissions system works, we don't want it to skip the checks
+            # as well.
+            await self._verify_checks(ctx)
 
         await super().invoke(ctx)
 
@@ -184,3 +235,25 @@ def group(name=None, **attrs):
     Same interface as `discord.ext.commands.group`.
     """
     return command(name, cls=Group, **attrs)
+
+
+__command_disablers = weakref.WeakValueDictionary()
+
+
+def get_command_disabler(guild: discord.Guild) -> Callable[["Context"], Awaitable[bool]]:
+    """Get the command disabler for a guild.
+
+    A command disabler is a simple check predicate which returns
+    ``False`` if the context is within the given guild.
+    """
+    try:
+        return __command_disablers[guild]
+    except KeyError:
+
+        async def disabler(ctx: "Context") -> bool:
+            if ctx.guild == guild:
+                raise commands.DisabledCommand()
+            return True
+
+        __command_disablers[guild] = disabler
+        return disabler

redbot/core/core_commands.py

@@ -1,4 +1,5 @@
 import asyncio
+import contextlib
 import datetime
 import importlib
 import itertools
@@ -1561,6 +1562,138 @@ class Core(CoreLogic):
         await ctx.bot.db.guild(ctx.guild).blacklist.set([])
         await ctx.send(_("blacklist has been cleared."))
 
+    @checks.guildowner_or_permissions(administrator=True)
+    @commands.group(name="command")
+    async def command_manager(self, ctx: commands.Context):
+        """Manage the bot's commands."""
+        pass
+
+    @command_manager.group(name="disable", invoke_without_command=True)
+    async def command_disable(self, ctx: commands.Context, *, command: str):
+        """Disable a command.
+
+        If you're the bot owner, this will disable commands
+        globally by default.
+        """
+        # Select the scope based on the author's privileges
+        if await ctx.bot.is_owner(ctx.author):
+            await ctx.invoke(self.command_disable_global, command=command)
+        else:
+            await ctx.invoke(self.command_disable_guild, command=command)
+
+    @checks.is_owner()
+    @command_disable.command(name="global")
+    async def command_disable_global(self, ctx: commands.Context, *, command: str):
+        """Disable a command globally."""
+        command_obj: commands.Command = ctx.bot.get_command(command)
+        if command_obj is None:
+            await ctx.send(
+                _("I couldn't find that command. Please note that it is case sensitive.")
+            )
+            return
+
+        async with ctx.bot.db.disabled_commands() as disabled_commands:
+            if command not in disabled_commands:
+                disabled_commands.append(command_obj.qualified_name)
+
+        if not command_obj.enabled:
+            await ctx.send(_("That command is already disabled globally."))
+            return
+        command_obj.enabled = False
+
+        await ctx.tick()
+
+    @commands.guild_only()
+    @command_disable.command(name="server", aliases=["guild"])
+    async def command_disable_guild(self, ctx: commands.Context, *, command: str):
+        """Disable a command in this server only."""
+        command_obj: commands.Command = ctx.bot.get_command(command)
+        if command_obj is None:
+            await ctx.send(
+                _("I couldn't find that command. Please note that it is case sensitive.")
+            )
+            return
+
+        async with ctx.bot.db.guild(ctx.guild).disabled_commands() as disabled_commands:
+            if command not in disabled_commands:
+                disabled_commands.append(command_obj.qualified_name)
+
+        done = command_obj.disable_in(ctx.guild)
+
+        if not done:
+            await ctx.send(_("That command is already disabled in this server."))
+        else:
+            await ctx.tick()
+
+    @command_manager.group(name="enable", invoke_without_command=True)
+    async def command_enable(self, ctx: commands.Context, *, command: str):
+        """Enable a command.
+
+        If you're a bot owner, this will try to enable a globally
+        disabled command by default.
+        """
+        if await ctx.bot.is_owner(ctx.author):
+            await ctx.invoke(self.command_enable_global, command=command)
+        else:
+            await ctx.invoke(self.command_enable_guild, command=command)
+
+    @commands.is_owner()
+    @command_enable.command(name="global")
+    async def command_enable_global(self, ctx: commands.Context, *, command: str):
+        """Enable a command globally."""
+        command_obj: commands.Command = ctx.bot.get_command(command)
+        if command_obj is None:
+            await ctx.send(
+                _("I couldn't find that command. Please note that it is case sensitive.")
+            )
+            return
+
+        async with ctx.bot.db.disabled_commands() as disabled_commands:
+            with contextlib.suppress(ValueError):
+                disabled_commands.remove(command_obj.qualified_name)
+
+        if command_obj.enabled:
+            await ctx.send(_("That command is already enabled globally."))
+            return
+
+        command_obj.enabled = True
+        await ctx.tick()
+
+    @commands.guild_only()
+    @command_enable.command(name="server", aliases=["guild"])
+    async def command_enable_guild(self, ctx: commands.Context, *, command: str):
+        """Enable a command in this server."""
+        command_obj: commands.Command = ctx.bot.get_command(command)
+        if command_obj is None:
+            await ctx.send(
+                _("I couldn't find that command. Please note that it is case sensitive.")
+            )
+            return
+
+        async with ctx.bot.db.guild(ctx.guild).disabled_commands() as disabled_commands:
+            with contextlib.suppress(ValueError):
+                disabled_commands.remove(command_obj.qualified_name)
+
+        done = command_obj.enable_in(ctx.guild)
+
+        if not done:
+            await ctx.send(_("That command is already enabled in this server."))
+        else:
+            await ctx.tick()
+
+    @checks.is_owner()
+    @command_manager.command(name="disabledmsg")
+    async def command_disabledmsg(self, ctx: commands.Context, *, message: str = ""):
+        """Set the bot's response to disabled commands.
+
+        Leave blank to send nothing.
+
+        To include the command name in the message, include the
+        `{command}` placeholder.
+        """
+        await ctx.bot.db.disabled_command_msg.set(message)
+        await ctx.tick()
+
     # RPC handlers
     async def rpc_load(self, request):
         cog_name = request.params[0]

redbot/core/events.py

@@ -193,7 +193,9 @@ def init_events(bot, cli_flags):
         elif isinstance(error, commands.BadArgument):
             await ctx.send_help()
         elif isinstance(error, commands.DisabledCommand):
-            await ctx.send("That command is disabled.")
+            disabled_message = await bot.db.disabled_command_msg()
+            if disabled_message:
+                await ctx.send(disabled_message.replace("{command}", ctx.invoked_with))
         elif isinstance(error, commands.CommandInvokeError):
             # Need to test if the following still works
             """
@@ -241,7 +243,7 @@ def init_events(bot, cli_flags):
             await ctx.send("That command is not available in DMs.")
         elif isinstance(error, commands.CommandOnCooldown):
             await ctx.send(
-                "This command is on cooldown. " "Try again in {:.2f}s" "".format(error.retry_after)
+                "This command is on cooldown. Try again in {:.2f}s".format(error.retry_after)
             )
         else:
             log.exception(type(error).__name__, exc_info=error)
@@ -280,6 +282,42 @@ def init_events(bot, cli_flags):
     async def on_command(command):
         bot.counter["processed_commands"] += 1
 
+    @bot.event
+    async def on_command_add(command: commands.Command):
+        disabled_commands = await bot.db.disabled_commands()
+        if command.qualified_name in disabled_commands:
+            command.enabled = False
+        for guild in bot.guilds:
+            disabled_commands = await bot.db.guild(guild).disabled_commands()
+            if command.qualified_name in disabled_commands:
+                command.disable_in(guild)
+
+    async def _guild_added(guild: discord.Guild):
+        disabled_commands = await bot.db.guild(guild).disabled_commands()
+        for command_name in disabled_commands:
+            command_obj = bot.get_command(command_name)
+            if command_obj is not None:
+                command_obj.disable_in(guild)
+
+    @bot.event
+    async def on_guild_join(guild: discord.Guild):
+        await _guild_added(guild)
+
+    @bot.event
+    async def on_guild_available(guild: discord.Guild):
+        # We need to check guild-disabled commands here since some cogs
+        # are loaded prior to `on_ready`.
+        await _guild_added(guild)
+
+    @bot.event
+    async def on_guild_leave(guild: discord.Guild):
+        # Clean up any unneeded checks
+        disabled_commands = await bot.db.guild(guild).disabled_commands()
+        for command_name in disabled_commands:
+            command_obj = bot.get_command(command_name)
+            if command_obj is not None:
+                command_obj.enable_in(guild)
+
 
 def _get_startup_screen_specs():
     """Get specs for displaying the startup screen on stdout.