Create cog disabling API (Cog-Creators#4043)

mikeshardmind · Jackenmen · Drapersniper · web-flow · commit 1d80fe9aecbd · 2020-07-28T20:52:36.000+02:00
* create cog disbale base

* Because defaults...

* lol

* announcer needs to respect this

* defaultdict mishap

* Allow None as guild

- Mostly for interop with with ctx.guild

* a whitespace issue

* Apparently, I broke this too

* Apply suggestions from code review

Co-authored-by: jack1142 &lt;6032823+jack1142@users.noreply.github.com&gt;

* This can probably be more optimized later, but since this is a cached value, it's not a large issue

* Report tunnel closing

* mod too

* whitespace issue

* Fix Artifact of prior method naming

* these 3 places should have the check if i understood it correctly

* Announce the closed tunnels

* tunnel oversight

* Make the player stop at next track

* added where draper said to put it

* Apply suggestions from code review

Co-authored-by: jack1142 &lt;6032823+jack1142@users.noreply.github.com&gt;

Co-authored-by: jack1142 &lt;6032823+jack1142@users.noreply.github.com&gt;
Co-authored-by: Drapersniper &lt;27962761+drapersniper@users.noreply.github.com&gt;
diff --git a/docs/guide_cog_creation.rst b/docs/guide_cog_creation.rst
@@ -229,3 +229,12 @@ Not all of these are strict requirements (some are) but are all generally advisa
 
   - We announce this in advance.
   - If you need help, ask.
+
+14. Check events against ``bot.cog_disabled_in_guild``
+
+  - Not all events need to be checked, only those that interact with a guild.
+  - Some discretion may apply, for example,
+    a cog which logs command invocation errors could choose to ignore this
+    but a cog which takes actions based on messages should not.
+
+15. Respect settings when treating non command messages as commands.
diff --git a/redbot/cogs/admin/announcer.py b/redbot/cogs/admin/announcer.py
@@ -40,6 +40,8 @@ def cancel(self):
         self.active = False
 
     async def _get_announce_channel(self, guild: discord.Guild) -> Optional[discord.TextChannel]:
+        if await self.ctx.bot.cog_disabled_in_guild_raw("Admin", guild.id):
+            return
         channel_id = await self.config.guild(guild).announce_channel()
         return guild.get_channel(channel_id)
 
diff --git a/redbot/cogs/alias/alias.py b/redbot/cogs/alias/alias.py
@@ -326,6 +326,11 @@ async def _list_global_alias(self, ctx: commands.Context):
 
     @commands.Cog.listener()
     async def on_message_without_command(self, message: discord.Message):
+
+        if message.guild is not None:
+            if await self.bot.cog_disabled_in_guild(self, message.guild):
+                return
+
         try:
             prefix = await self.get_prefix(message)
         except ValueError:
diff --git a/redbot/cogs/audio/core/events/cog.py b/redbot/cogs/audio/core/events/cog.py
@@ -25,6 +25,13 @@ async def on_red_audio_track_start(
     ):
         if not (track and guild):
             return
+
+        if await self.bot.cog_disabled_in_guild(self, guild):
+            player = lavalink.get_player(guild.id)
+            await player.stop()
+            await player.disconnect()
+            return
+
         track_identifier = track.track_identifier
         if self.playlist_api is not None:
             daily_cache = self._daily_playlist_cache.setdefault(
diff --git a/redbot/cogs/audio/core/events/dpy.py b/redbot/cogs/audio/core/events/dpy.py
@@ -178,6 +178,8 @@ def cog_unload(self) -> None:
     async def on_voice_state_update(
         self, member: discord.Member, before: discord.VoiceState, after: discord.VoiceState
     ) -> None:
+        if await self.bot.cog_disabled_in_guild(self, member.guild):
+            return
         await self.cog_ready_event.wait()
         if after.channel != before.channel:
             try:
diff --git a/redbot/cogs/audio/core/events/lavalink.py b/redbot/cogs/audio/core/events/lavalink.py
@@ -19,7 +19,13 @@ async def lavalink_event_handler(
         current_track = player.current
         current_channel = player.channel
         guild = self.rgetattr(current_channel, "guild", None)
+        if await self.bot.cog_disabled_in_guild(self, guild):
+            await player.stop()
+            await player.disconnect()
+            return
         guild_id = self.rgetattr(guild, "id", None)
+        if not guild:
+            return
         current_requester = self.rgetattr(current_track, "requester", None)
         current_stream = self.rgetattr(current_track, "is_stream", None)
         current_length = self.rgetattr(current_track, "length", None)
diff --git a/redbot/cogs/audio/core/tasks/player.py b/redbot/cogs/audio/core/tasks/player.py
@@ -20,6 +20,8 @@ async def player_automated_timer(self) -> None:
         while True:
             async for p in AsyncIter(lavalink.all_players()):
                 server = p.channel.guild
+                if await self.bot.cog_disabled_in_guild(self, server):
+                    continue
 
                 if [self.bot.user] == p.channel.members:
                     stop_times.setdefault(server.id, time.time())
diff --git a/redbot/cogs/customcom/customcom.py b/redbot/cogs/customcom/customcom.py
@@ -516,6 +516,9 @@ async def on_message_without_command(self, message):
         if len(message.content) < 2 or is_private or not user_allowed or message.author.bot:
             return
 
+        if await self.bot.cog_disabled_in_guild(self, message.guild):
+            return
+
         ctx = await self.bot.get_context(message)
 
         if ctx.prefix is None:
diff --git a/redbot/cogs/filter/filter.py b/redbot/cogs/filter/filter.py
@@ -369,6 +369,10 @@ async def check_filter(self, message: discord.Message):
     async def on_message(self, message: discord.Message):
         if isinstance(message.channel, discord.abc.PrivateChannel):
             return
+
+        if await self.bot.cog_disabled_in_guild(self, message.guild):
+            return
+
         author = message.author
         valid_user = isinstance(author, discord.Member) and not author.bot
         if not valid_user:
@@ -395,6 +399,11 @@ async def on_member_join(self, member: discord.Member):
         await self.maybe_filter_name(member)
 
     async def maybe_filter_name(self, member: discord.Member):
+
+        guild = member.guild
+        if (not guild) or await self.bot.cog_disabled_in_guild(self, guild):
+            return
+
         if not member.guild.me.guild_permissions.manage_nicknames:
             return  # No permissions to manage nicknames, so can't do anything
         if member.top_role >= member.guild.me.top_role:
diff --git a/redbot/cogs/mod/events.py b/redbot/cogs/mod/events.py
@@ -79,6 +79,10 @@ async def on_message(self, message):
         author = message.author
         if message.guild is None or self.bot.user == author:
             return
+
+        if await self.bot.cog_disabled_in_guild(self, message.guild):
+            return
+
         valid_user = isinstance(author, discord.Member) and not author.bot
         if not valid_user:
             return
@@ -110,6 +114,9 @@ async def on_user_update(self, before: discord.User, after: discord.User):
     @commands.Cog.listener()
     async def on_member_update(self, before: discord.Member, after: discord.Member):
         if before.nick != after.nick and after.nick is not None:
+            guild = after.guild
+            if (not guild) or await self.bot.cog_disabled_in_guild(self, guild):
+                return
             async with self.config.member(before).past_nicks() as nick_list:
                 while None in nick_list:  # clean out null entries from a bug
                     nick_list.remove(None)
diff --git a/redbot/cogs/mod/kickban.py b/redbot/cogs/mod/kickban.py
@@ -142,6 +142,9 @@ async def check_tempban_expirations(self):
                 if not guild.me.guild_permissions.ban_members:
                     continue
 
+                if await self.bot.cog_disabled_in_guild(self, guild):
+                    continue
+
                 async with self.config.guild(guild).current_tempbans() as guild_tempbans:
                     for uid in guild_tempbans.copy():
                         unban_time = datetime.utcfromtimestamp(
diff --git a/redbot/cogs/reports/reports.py b/redbot/cogs/reports/reports.py
@@ -293,10 +293,11 @@ async def report_cleanup(self, ctx: commands.Context):
                     pass
 
     @commands.Cog.listener()
-    async def on_raw_reaction_add(self, payload):
+    async def on_raw_reaction_add(self, payload: discord.RawReactionActionEvent):
         """
         oh dear....
         """
+
         if not str(payload.emoji) == "\N{NEGATIVE SQUARED CROSS MARK}":
             return
 
@@ -314,13 +315,35 @@ async def on_raw_reaction_add(self, payload):
 
     @commands.Cog.listener()
     async def on_message(self, message: discord.Message):
+
+        to_remove = []
+
         for k, v in self.tunnel_store.items():
-            topic = _("Re: ticket# {1} in {0.name}").format(*k)
+
+            guild, ticket_number = k
+            if await self.bot.cog_disabled_in_guild(self, guild):
+                to_remove.append(k)
+                continue
+
+            topic = _("Re: ticket# {ticket_number} in {guild.name}").format(
+                ticket_number=ticket_number, guild=guild
+            )
             # Tunnels won't forward unintended messages, this is safe
             msgs = await v["tun"].communicate(message=message, topic=topic)
             if msgs:
                 self.tunnel_store[k]["msgs"] = msgs
 
+        for key in to_remove:
+            if tun := self.tunnel_store.pop(key, None):
+                guild, ticket = key
+                await tun["tun"].close_because_disabled(
+                    _(
+                        "Correspondence about ticket# {ticket_number} in "
+                        "{guild.name} has been ended due "
+                        "to reports being disabled in that server."
+                    ).format(ticket_number=ticket, guild=guild)
+                )
+
     @commands.guild_only()
     @checks.mod_or_permissions(manage_roles=True)
     @report.command(name="interact")
diff --git a/redbot/cogs/streams/streams.py b/redbot/cogs/streams/streams.py
@@ -702,6 +702,8 @@ async def check_streams(self):
                         continue
                     for message in stream._messages_cache:
                         with contextlib.suppress(Exception):
+                            if await self.bot.cog_disabled_in_guild(self, message.guild):
+                                continue
                             autodelete = await self.config.guild(message.guild).autodelete()
                             if autodelete:
                                 await message.delete()
@@ -714,6 +716,8 @@ async def check_streams(self):
                         channel = self.bot.get_channel(channel_id)
                         if not channel:
                             continue
+                        if await self.bot.cog_disabled_in_guild(self, channel.guild):
+                            continue
                         ignore_reruns = await self.config.guild(channel.guild).ignore_reruns()
                         if ignore_reruns and is_rerun:
                             continue
diff --git a/redbot/core/bot.py b/redbot/core/bot.py
@@ -37,7 +37,12 @@
 from .events import init_events
 from .global_checks import init_global_checks
 
-from .settings_caches import PrefixManager, IgnoreManager, WhitelistBlacklistManager
+from .settings_caches import (
+    PrefixManager,
+    IgnoreManager,
+    WhitelistBlacklistManager,
+    DisabledCogCache,
+)
 
 from .rpc import RPCMixin
 from .utils import common_filters
@@ -132,12 +137,16 @@ def __init__(self, *args, cli_flags=None, bot_dir: Path = Path.cwd(), **kwargs):
         self._config.register_channel(embeds=None, ignored=False)
         self._config.register_user(embeds=None)
 
+        self._config.init_custom("COG_DISABLE_SETTINGS", 2)
+        self._config.register_custom("COG_DISABLE_SETTINGS", disabled=None)
+
         self._config.init_custom(CUSTOM_GROUPS, 2)
         self._config.register_custom(CUSTOM_GROUPS)
 
         self._config.init_custom(SHARED_API_TOKENS, 2)
         self._config.register_custom(SHARED_API_TOKENS)
         self._prefix_cache = PrefixManager(self._config, cli_flags)
+        self._disabled_cog_cache = DisabledCogCache(self._config)
         self._ignored_cache = IgnoreManager(self._config)
         self._whiteblacklist_cache = WhitelistBlacklistManager(self._config)
 
@@ -217,6 +226,41 @@ async def _red_before_invoke_method(self, ctx):
                 return_exceptions=return_exceptions,
             )
 
+    async def cog_disabled_in_guild(
+        self, cog: commands.Cog, guild: Optional[discord.Guild]
+    ) -> bool:
+        """
+        Check if a cog is disabled in a guild
+
+        Parameters
+        ----------
+        cog: commands.Cog
+        guild: Optional[discord.Guild]
+
+        Returns
+        -------
+        bool
+        """
+        if guild is None:
+            return False
+        return await self._disabled_cog_cache.cog_disabled_in_guild(cog.qualified_name, guild.id)
+
+    async def cog_disabled_in_guild_raw(self, cog_name: str, guild_id: int) -> bool:
+        """
+        Check if a cog is disabled in a guild without the cog or guild object
+
+        Parameters
+        ----------
+        cog_name: str
+            This should be the cog's qualified name, not neccessarily the classname
+        guild_id: int
+
+        Returns
+        -------
+        bool
+        """
+        return await self._disabled_cog_cache.cog_disabled_in_guild(cog_name, guild_id)
+
     def remove_before_invoke_hook(self, coro: PreInvokeCoroutine) -> None:
         """
         Functional method to remove a `before_invoke` hook.
diff --git a/redbot/core/commands/requires.py b/redbot/core/commands/requires.py
@@ -511,6 +511,10 @@ async def _verify_bot(self, ctx: "Context") -> None:
             bot_user = ctx.bot.user
         else:
             bot_user = ctx.guild.me
+            cog = ctx.cog
+            if cog and await ctx.bot.cog_disabled_in_guild(cog, ctx.guild):
+                raise discord.ext.commands.DisabledCommand()
+
         bot_perms = ctx.channel.permissions_for(bot_user)
         if not (bot_perms.administrator or bot_perms >= self.bot_perms):
             raise BotMissingPermissions(missing=self._missing_perms(self.bot_perms, bot_perms))
diff --git a/redbot/core/core_commands.py b/redbot/core/core_commands.py
@@ -2174,9 +2174,83 @@ async def localblacklist_clear(self, ctx: commands.Context):
     @checks.guildowner_or_permissions(administrator=True)
     @commands.group(name="command")
     async def command_manager(self, ctx: commands.Context):
-        """Manage the bot's commands."""
+        """Manage the bot's commands and cogs."""
         pass
 
+    @checks.is_owner()
+    @command_manager.command(name="defaultdisablecog")
+    async def command_default_disable_cog(self, ctx: commands.Context, *, cogname: str):
+        """Set the default state for a cog as disabled."""
+        cog = self.bot.get_cog(cogname)
+        if not cog:
+            return await ctx.send(_("Cog with the given name doesn't exist."))
+        if cog == self:
+            return await ctx.send(_("You can't disable this cog by default."))
+        await self.bot._disabled_cog_cache.default_disable(cogname)
+        await ctx.send(_("{cogname} has been set as disabled by default.").format(cogname=cogname))
+
+    @checks.is_owner()
+    @command_manager.command(name="defaultenablecog")
+    async def command_default_enable_cog(self, ctx: commands.Context, *, cogname: str):
+        """Set the default state for a cog as enabled."""
+        cog = self.bot.get_cog(cogname)
+        if not cog:
+            return await ctx.send(_("Cog with the given name doesn't exist."))
+        await self.bot._disabled_cog_cache.default_enable(cogname)
+        await ctx.send(_("{cogname} has been set as enabled by default.").format(cogname=cogname))
+
+    @commands.guild_only()
+    @command_manager.command(name="disablecog")
+    async def command_disable_cog(self, ctx: commands.Context, *, cogname: str):
+        """Disable a cog in this guild."""
+        cog = self.bot.get_cog(cogname)
+        if not cog:
+            return await ctx.send(_("Cog with the given name doesn't exist."))
+        if cog == self:
+            return await ctx.send(_("You can't disable this cog as you would lock yourself out."))
+        if await self.bot._disabled_cog_cache.disable_cog_in_guild(cogname, ctx.guild.id):
+            await ctx.send(_("{cogname} has been disabled in this guild.").format(cogname=cogname))
+        else:
+            await ctx.send(
+                _("{cogname} was already disabled (nothing to do).").format(cogname=cogname)
+            )
+
+    @commands.guild_only()
+    @command_manager.command(name="enablecog")
+    async def command_enable_cog(self, ctx: commands.Context, *, cogname: str):
+        """Enable a cog in this guild."""
+        if await self.bot._disabled_cog_cache.enable_cog_in_guild(cogname, ctx.guild.id):
+            await ctx.send(_("{cogname} has been enabled in this guild.").format(cogname=cogname))
+        else:
+            # putting this here allows enabling a cog that isn't loaded but was disabled.
+            cog = self.bot.get_cog(cogname)
+            if not cog:
+                return await ctx.send(_("Cog with the given name doesn't exist."))
+
+            await ctx.send(
+                _("{cogname} was not disabled (nothing to do).").format(cogname=cogname)
+            )
+
+    @commands.guild_only()
+    @command_manager.command(name="listdisabledcogs")
+    async def command_list_disabled_cogs(self, ctx: commands.Context):
+        """List the cogs which are disabled in this guild."""
+        disabled = [
+            cog.qualified_name
+            for cog in self.bot.cogs.values()
+            if await self.bot._disabled_cog_cache.cog_disabled_in_guild(
+                cog.qualified_name, ctx.guild.id
+            )
+        ]
+        if disabled:
+            output = _("The following cogs are disabled in this guild:\n")
+            output += humanize_list(disabled)
+
+            for page in pagify(output):
+                await ctx.send(page)
+        else:
+            await ctx.send(_("There are no disabled cogs in this guild."))
+
     @command_manager.group(name="listdisabled", invoke_without_command=True)
     async def list_disabled(self, ctx: commands.Context):
         """
diff --git a/redbot/core/settings_caches.py b/redbot/core/settings_caches.py
diff --git a/redbot/core/utils/tunnel.py b/redbot/core/utils/tunnel.py
