diff options
| -rw-r--r-- | bot/exts/utils/clean.py | 254 |
1 files changed, 183 insertions, 71 deletions
diff --git a/bot/exts/utils/clean.py b/bot/exts/utils/clean.py index cb662e852..25582165a 100644 --- a/bot/exts/utils/clean.py +++ b/bot/exts/utils/clean.py @@ -1,20 +1,26 @@ import logging -import random import re -from typing import Iterable, Optional +import time +from collections import defaultdict +from typing import Any, Callable, DefaultDict, Iterable, List, Optional, Tuple -from discord import Colour, Embed, Message, TextChannel, User, errors +from discord import Colour, Embed, Message, NotFound, TextChannel, User, errors from discord.ext import commands from discord.ext.commands import Cog, Context, group, has_any_role +from discord.ext.commands.errors import BadArgument, MaxConcurrencyReached, MissingRequiredArgument from bot.bot import Bot from bot.constants import ( - Channels, CleanMessages, Colours, Event, Icons, MODERATION_ROLES, NEGATIVE_REPLIES + Channels, CleanMessages, Colours, Event, Icons, MODERATION_ROLES ) from bot.exts.moderation.modlog import ModLog +from bot.utils.channel import is_mod_channel log = logging.getLogger(__name__) +# Type alias for checks +Predicate = Callable[[Message], bool] + class Clean(Cog): """ @@ -36,15 +42,92 @@ class Clean(Cog): """Get currently loaded ModLog cog instance.""" return self.bot.get_cog("ModLog") + @staticmethod + def is_older_than_14d(message: Message) -> bool: + """ + Precisely checks if message is older than 14 days, bulk deletion limit. + + Inspired by how purge works internally. + Comparison on message age could possibly be less accurate which in turn would resort in problems + with message deletion if said messages are very close to the 14d mark. + """ + two_weeks_old_snowflake = int((time.time() - 14 * 24 * 60 * 60) * 1000.0 - 1420070400000) << 22 + return message.id < two_weeks_old_snowflake + + async def _delete_messages_individually(self, messages: List[Message]) -> None: + for message in messages: + # Ensure that deletion was not canceled + if not self.cleaning: + return + try: + await message.delete() + except NotFound: + # Message doesn't exist or was already deleted + continue + + def _get_messages_from_cache(self, amount: int, to_delete: Predicate) -> Tuple[DefaultDict, List[int]]: + """Helper function for getting messages from the cache.""" + message_mappings = defaultdict(list) + message_ids = [] + for message in self.bot.cached_messages: + if not self.cleaning: + # Cleaning was canceled + return (message_mappings, message_ids) + + if to_delete(message): + message_mappings[message.channel].append(message) + message_ids.append(message.id) + + if len(message_ids) == amount: + # We've got the requested amount of messages + return message_mappings, message_ids + + # Amount exceeds amount of messages matching the check + return message_mappings, message_ids + + async def _get_messages_from_channels( + self, + amount: int, + channels: Iterable[TextChannel], + to_delete: Predicate, + until_message: Optional[Message] = None + ) -> tuple[defaultdict[Any, list], list]: + message_mappings = defaultdict(list) + message_ids = [] + + for channel in channels: + + async for message in channel.history(limit=amount): + + if not self.cleaning: + # Cleaning was canceled, return empty containers + return defaultdict(list), [] + + if until_message: + + # We could use ID's here however in case if the message we are looking for gets deleted, + # We won't have a way to figure that out thus checking for datetime should be more reliable + if message.created_at < until_message.created_at: + # Means we have found the message until which we were supposed to be deleting. + break + + if to_delete(message): + message_mappings[message.channel].append(message) + message_ids.append(message.id) + + return message_mappings, message_ids + async def _clean_messages( self, amount: int, ctx: Context, channels: Iterable[TextChannel], bots_only: bool = False, + use_cache: bool = False, user: User = None, regex: Optional[str] = None, until_message: Optional[Message] = None, + after_message: Optional[Message] = None, ) -> None: """A helper function that does the actual message cleaning.""" def predicate_bots_only(message: Message) -> bool: @@ -79,25 +162,30 @@ class Clean(Cog): else: return bool(re.search(regex.lower(), content.lower())) + def predicate_range(message: Message) -> bool: + """Check if message is older than message provided in after_message but younger than until_message.""" + return after_message.created_at <= message.created_at <= until_message.created_at + # Is this an acceptable amount of messages to clean? if amount > CleanMessages.message_limit: - embed = Embed( - color=Colour(Colours.soft_red), - title=random.choice(NEGATIVE_REPLIES), - description=f"You cannot clean more than {CleanMessages.message_limit} messages." - ) - await ctx.send(embed=embed) - return + raise BadArgument(f"You cannot clean more than {CleanMessages.message_limit} messages.") + + if after_message: + # Ensure that until_message is specified. + if not until_message: + raise MissingRequiredArgument("`until_message` must be specified if `after_message` is specified.") + + # Messages are not in same channel + if after_message.channel != until_message.channel: + raise BadArgument("You cannot do range clean across several channel.") + + # Ensure that after_message is younger than until_message + if after_message.created_at >= until_message.created_at: + raise BadArgument("`after` message must be younger than `until` message") # Are we already performing a clean? if self.cleaning: - embed = Embed( - color=Colour(Colours.soft_red), - title=random.choice(NEGATIVE_REPLIES), - description="Please wait for the currently ongoing clean operation to complete." - ) - await ctx.send(embed=embed) - return + raise MaxConcurrencyReached("Please wait for the currently ongoing clean operation to complete.") # Set up the correct predicate if bots_only: @@ -106,77 +194,87 @@ class Clean(Cog): predicate = predicate_specific_user # Delete messages from specific user elif regex: predicate = predicate_regex # Delete messages that match regex + elif after_message: + predicate = predicate_range # Delete messages older than specific message else: - predicate = None # Delete all messages + predicate = lambda m: True # Delete all messages # noqa: E731 # Default to using the invoking context's channel if not channels: channels = [ctx.channel] - # Delete the invocation first - self.mod_log.ignore(Event.message_delete, ctx.message.id) - try: - await ctx.message.delete() - except errors.NotFound: - # Invocation message has already been deleted - log.info("Tried to delete invocation message, but it was already deleted.") + if not is_mod_channel(ctx.channel): + + # Delete the invocation first + self.mod_log.ignore(Event.message_delete, ctx.message.id) + try: + await ctx.message.delete() + except errors.NotFound: + # Invocation message has already been deleted + log.info("Tried to delete invocation message, but it was already deleted.") - messages = [] - message_ids = [] self.cleaning = True - # Find the IDs of the messages to delete. IDs are needed in order to ignore mod log events. - for channel in channels: - async for message in channel.history(limit=amount): + if use_cache: + message_mappings, message_ids = self._get_messages_from_cache(amount=amount, to_delete=predicate) + else: + message_mappings, message_ids = await self._get_messages_from_channels( + amount=amount, + channels=channels, + to_delete=predicate, + until_message=until_message + ) + + if not self.cleaning: + # Means that the cleaning was canceled + return + + # Now let's delete the actual messages with purge. + self.mod_log.ignore(Event.message_delete, *message_ids) + + for channel, messages in message_mappings.items(): + + to_delete = [] + + for current_index, message in enumerate(messages): - # If at any point the cancel command is invoked, we should stop. if not self.cleaning: + # Means that the cleaning was canceled return - # If we are looking for specific message. - if until_message: + if self.is_older_than_14d(message): + # further messages are too old to be deleted in bulk + await self._delete_messages_individually(messages[current_index:]) + if not self.cleaning: + # Means that deletion was canceled while deleting the individual messages + return + break - # we could use ID's here however in case if the message we are looking for gets deleted, - # we won't have a way to figure that out thus checking for datetime should be more reliable - if message.created_at < until_message.created_at: - # means we have found the message until which we were supposed to be deleting. - break + to_delete.append(message) - # Since we will be using `delete_messages` method of a TextChannel and we need message objects to - # use it as well as to send logs we will start appending messages here instead adding them from - # purge. - messages.append(message) + if len(to_delete) == 100: + # we can only delete up to 100 messages in a bulk + await channel.delete_messages(to_delete) + to_delete.clear() - # If the message passes predicate, let's save it. - if predicate is None or predicate(message): - message_ids.append(message.id) + if len(to_delete) > 0: + # deleting any leftover messages if there are any + await channel.delete_messages(to_delete) self.cleaning = False - # Now let's delete the actual messages with purge. - self.mod_log.ignore(Event.message_delete, *message_ids) - for channel in channels: - if until_message: - for i in range(0, len(messages), 100): - # while purge automatically handles the amount of messages - # delete_messages only allows for up to 100 messages at once - # thus we need to paginate the amount to always be <= 100 - await channel.delete_messages(messages[i:i + 100]) - else: - messages += await channel.purge(limit=amount, check=predicate) + log_messages = [] + + for messages in message_mappings.values(): + log_messages.extend(messages) - # Reverse the list to restore chronological order - if messages: - messages = reversed(messages) - log_url = await self.mod_log.upload_log(messages, ctx.author.id) + if log_messages: + # Reverse the list to restore chronological order + log_messages = reversed(log_messages) + log_url = await self.mod_log.upload_log(log_messages, ctx.author.id) else: # Can't build an embed, nothing to clean! - embed = Embed( - color=Colour(Colours.soft_red), - description="No matching messages could be found." - ) - await ctx.send(embed=embed, delete_after=10) - return + raise BadArgument("No matching messages could be found.") # Build the embed and send it target_channels = ", ".join(channel.mention for channel in channels) @@ -211,7 +309,8 @@ class Clean(Cog): channels: commands.Greedy[TextChannel] = None ) -> None: """Delete messages posted by the provided user, stop cleaning after traversing `amount` messages.""" - await self._clean_messages(amount, ctx, user=user, channels=channels) + use_cache = not channels + await self._clean_messages(amount, ctx, user=user, channels=channels, use_cache=use_cache) @clean_group.command(name="all", aliases=["everything"]) @has_any_role(*MODERATION_ROLES) @@ -245,11 +344,12 @@ class Clean(Cog): channels: commands.Greedy[TextChannel] = None ) -> None: """Delete all messages that match a certain regex, stop cleaning after traversing `amount` messages.""" - await self._clean_messages(amount, ctx, regex=regex, channels=channels) + use_cache = not channels + await self._clean_messages(amount, ctx, regex=regex, channels=channels, use_cache=use_cache) - @clean_group.command(name="message", aliases=["messages"]) + @clean_group.command(name="until") @has_any_role(*MODERATION_ROLES) - async def clean_message(self, ctx: Context, message: Message) -> None: + async def clean_until(self, ctx: Context, message: Message) -> None: """Delete all messages until certain message, stop cleaning after hitting the `message`.""" await self._clean_messages( CleanMessages.message_limit, @@ -258,6 +358,18 @@ class Clean(Cog): until_message=message ) + @clean_group.command(name="from-to", aliases=["after-until", "between"]) + @has_any_role(*MODERATION_ROLES) + async def clean_from_to(self, ctx: Context, after_message: Message, until_message: Message) -> None: + """Delete all messages within range of messages.""" + await self._clean_messages( + CleanMessages.message_limit, + ctx, + channels=[until_message.channel], + until_message=until_message, + after_message=after_message, + ) + @clean_group.command(name="stop", aliases=["cancel", "abort"]) @has_any_role(*MODERATION_ROLES) async def clean_cancel(self, ctx: Context) -> None: |