diff --git a/cogs/__init__.py b/cogs/__init__.py index beca9474..62bc8e5e 100644 --- a/cogs/__init__.py +++ b/cogs/__init__.py @@ -14,6 +14,7 @@ CommitteeHandoverCommandCog, ) from .archive import ArchiveCommandCog +from .auto_slow_mode import AutomaticSlowModeCommandCog, AutomaticSlowModeTaskCog from .check_su_platform_authorisation import ( CheckSUPlatformAuthorisationCommandCog, CheckSUPlatformAuthorisationTaskCog, @@ -56,6 +57,8 @@ "AnnualRolesResetCommandCog", "AnnualYearChannelsIncrementCommandCog", "ArchiveCommandCog", + "AutomaticSlowModeCommandCog", + "AutomaticSlowModeTaskCog", "CheckSUPlatformAuthorisationCommandCog", "CheckSUPlatformAuthorisationTaskCog", "ClearRemindersBacklogTaskCog", @@ -98,6 +101,8 @@ def setup(bot: "TeXBot") -> None: AnnualRolesResetCommandCog, AnnualYearChannelsIncrementCommandCog, ArchiveCommandCog, + AutomaticSlowModeCommandCog, + AutomaticSlowModeTaskCog, ClearRemindersBacklogTaskCog, CommandErrorCog, CommitteeActionsTrackingSlashCommandsCog, diff --git a/cogs/auto_slow_mode.py b/cogs/auto_slow_mode.py new file mode 100644 index 00000000..2e083d8a --- /dev/null +++ b/cogs/auto_slow_mode.py @@ -0,0 +1,136 @@ +"""Module to handle automatic slow mode for Discord channels.""" + +import logging +from typing import TYPE_CHECKING, override + +import discord +from discord.ext import tasks + +from config import settings +from utils import CommandChecks, TeXBotBaseCog +from utils.error_capture_decorators import capture_guild_does_not_exist_error + +if TYPE_CHECKING: + from collections.abc import Sequence + from logging import Logger + from typing import Final + + from utils import TeXBot, TeXBotApplicationContext + + +__all__: "Sequence[str]" = ("AutomaticSlowModeCommandCog", "AutomaticSlowModeTaskCog") + + +logger: "Final[Logger]" = logging.getLogger("TeX-Bot") + + +class AutomaticSlowModeBaseCog(TeXBotBaseCog): + """Base class for automatic slow mode functionality.""" + + async def calculate_message_rate(self, channel: discord.TextChannel) -> int: + """ + Calculate the message rate for a given channel. + + Returns the number of messages per minute, rounded to the nearest integer. + This is based on the previous 5 minutes of messages. + """ + from datetime import UTC, datetime, timedelta + + # TODO: Make the time period user configurable. # noqa: FIX002 + + count = len( + [ + message + async for message in channel.history( + after=datetime.now(UTC) - timedelta(minutes=5), + oldest_first=False, + limit=None, + ) + ] + ) + + logger.debug( + "Channel: %s | Message count in last 5 minutes: %d | Rate: %d", + channel.name, + count, + round(count / 5), + ) + + return round(count / 5) + + +class AutomaticSlowModeTaskCog(AutomaticSlowModeBaseCog): + """Task to handle automatic slow mode for Discord channels.""" + + @override + def __init__(self, bot: "TeXBot") -> None: + """Start all task managers when this cog is initialised.""" + if settings["AUTO_SLOW_MODE"]: + _ = self.auto_slow_mode_task.start() + + super().__init__(bot) + + @override + def cog_unload(self) -> None: + """ + Unload-hook that ends all running tasks whenever the cog is unloaded. + + This may be run dynamically or when the bot closes. + """ + self.auto_slow_mode_task.cancel() + + @tasks.loop(seconds=60) + @capture_guild_does_not_exist_error + async def auto_slow_mode_task(self) -> None: + """Task to automatically adjust slow mode in channels.""" + import time + + time.process_time() + for channel in self.bot.get_all_channels(): + if isinstance(channel, discord.TextChannel): + message_rate: int = await self.calculate_message_rate(channel) + if message_rate > 5: + await channel.edit( + slowmode_delay=10, reason="TeX-Bot auto slow mode enabled." + ) + await channel.edit(slowmode_delay=0, reason="TeX-Bot auto slow mode disabled.") + logger.debug("Time taken to calculate message rate: %s seconds", time.process_time()) + + @auto_slow_mode_task.before_loop + async def before_tasks(self) -> None: + """Pre-execution hook, preventing any tasks from executing before the bot is ready.""" + await self.bot.wait_until_ready() + + +class AutomaticSlowModeCommandCog(AutomaticSlowModeBaseCog): + """Cog to handle automatic slow mode for Discord channels.""" + + @discord.slash_command( # type: ignore[misc, no-untyped-call] + name="toggle-auto-slow-mode", + description="Enable or disable automatic slow mode.", + ) + @CommandChecks.check_interaction_user_has_committee_role + @CommandChecks.check_interaction_user_in_main_guild + async def toggle_auto_slow_mode( # type: ignore[misc] + self, + ctx: "TeXBotApplicationContext", + ) -> None: + """Enable or disable automatic slow mode for a channel.""" + # NOTE: This should be replaced when the settings improved in PR #221 + if settings["AUTO_SLOW_MODE"]: + settings["AUTO_SLOW_MODE"] = False + await ctx.respond( + "Automatic slow mode is now disabled." + "If you would like to keep it disabled, please remember to " + "update the deployment variables as well. " + "If you do not, it will be re-enabled when the bot restarts.", + ) + return + + settings["AUTO_SLOW_MODE"] = True + await ctx.respond( + "Automatic slow mode is now enabled." + "If you would like to keep it enabled, please remember to " + "update the deployment variables as well. " + "If you do not, it will be disabled again when the bot restarts.", + ) diff --git a/config.py b/config.py index 498e06c4..fce7624a 100644 --- a/config.py +++ b/config.py @@ -935,6 +935,20 @@ def _setup_auto_add_committee_to_threads(cls) -> None: raw_auto_add_committee_to_threads in TRUE_VALUES ) + @classmethod + def _setup_auto_slow_mode(cls) -> None: + raw_auto_slow_mode: str = str( + os.getenv("AUTO_SLOW_MODE", "True"), + ).lower() + + if raw_auto_slow_mode not in TRUE_VALUES | FALSE_VALUES: + INVALID_AUTO_SLOW_MODE_MESSAGE: Final[str] = ( + "AUTO_SLOW_MODE must be a boolean value." + ) + raise ImproperlyConfiguredError(INVALID_AUTO_SLOW_MODE_MESSAGE) + + cls._settings["AUTO_SLOW_MODE"] = raw_auto_slow_mode in TRUE_VALUES + @classmethod def _setup_env_variables(cls) -> None: """ @@ -980,6 +994,7 @@ def _setup_env_variables(cls) -> None: cls._setup_moderation_document_url() cls._setup_strike_performed_manually_warning_location() cls._setup_auto_add_committee_to_threads() + cls._setup_auto_slow_mode() except ImproperlyConfiguredError as improper_config_error: webhook_config_logger.error(improper_config_error.message) # noqa: TRY400 raise improper_config_error from improper_config_error