From 9ee860c3f00932a2fa48808329113360d1ff5c0a Mon Sep 17 00:00:00 2001 From: Toby Harradine Date: Thu, 6 Sep 2018 12:52:19 +1000 Subject: [PATCH] [Core] Command disable feature (#2099) * [Core] Command disable feature Signed-off-by: Toby Harradine * [Core] Allow user to set the "command disabled" message Signed-off-by: Toby Harradine * Reformat Signed-off-by: Toby Harradine --- redbot/core/bot.py | 10 +++ redbot/core/commands/commands.py | 75 ++++++++++++++++- redbot/core/core_commands.py | 133 +++++++++++++++++++++++++++++++ redbot/core/events.py | 42 +++++++++- 4 files changed, 257 insertions(+), 3 deletions(-) diff --git a/redbot/core/bot.py b/redbot/core/bot.py index 04ba6aa19..335bea784 100644 --- a/redbot/core/bot.py +++ b/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): """ diff --git a/redbot/core/commands/commands.py b/redbot/core/commands/commands.py index 016cd1ec2..e3361eb00 100644 --- a/redbot/core/commands/commands.py +++ b/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 diff --git a/redbot/core/core_commands.py b/redbot/core/core_commands.py index 38bcd0df1..7f96fdf8f 100644 --- a/redbot/core/core_commands.py +++ b/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] diff --git a/redbot/core/events.py b/redbot/core/events.py index 6d89c683e..8d8480ba0 100644 --- a/redbot/core/events.py +++ b/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.