diff options
| -rw-r--r-- | bot/exts/utils/clean.py | 224 |
1 files changed, 182 insertions, 42 deletions
diff --git a/bot/exts/utils/clean.py b/bot/exts/utils/clean.py index 8acaf9131..f98e5c255 100644 --- a/bot/exts/utils/clean.py +++ b/bot/exts/utils/clean.py @@ -1,9 +1,11 @@ import logging import random import re -from typing import Iterable, Optional +import time +from collections import defaultdict +from typing import Callable, DefaultDict, Iterable, List, Optional -from discord import Colour, Embed, Message, TextChannel, User +from discord import Colour, Embed, Message, NotFound, TextChannel, User from discord.ext import commands from discord.ext.commands import Cog, Context, group, has_any_role @@ -15,6 +17,9 @@ from bot.exts.moderation.modlog import ModLog log = logging.getLogger(__name__) +# Type alias for checks +CheckHint = Callable[[Message], bool] + class Clean(Cog): """ @@ -36,15 +41,80 @@ class Clean(Cog): """Get currently loaded ModLog cog instance.""" return self.bot.get_cog("ModLog") + 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, predicate: CheckHint) -> List[DefaultDict, List[int]]: + """Helper function for getting messages from the cache.""" + message_mappings = defaultdict(lambda: []) + message_ids = [] + for message in self.bot.cached_messages: + if not self.cleaning: + # Cleaning was canceled + return (message_mappings, message_ids) + + if predicate(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], + predicate: CheckHint, + until_message: Optional[Message] = None + ) -> DefaultDict: + message_mappings = defaultdict(lambda: []) + message_ids = [] + + for channel in channels: + + async for message in channel.history(amount=amount): + + if not self.cleaning: + # Cleaning was canceled + return (message_mappings, message_ids) + + if predicate(message): + message_mappings[message.channel].append(message) + message_ids.append(message.id) + + 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 + + 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,6 +149,10 @@ 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 message.created_at > after_message.created_at and message.created_at < until_message.created_at + # Is this an acceptable amount of messages to clean? if amount > CleanMessages.message_limit: embed = Embed( @@ -89,6 +163,38 @@ class Clean(Cog): await ctx.send(embed=embed) return + if after_message: + + # Ensure that until_message is specified. + if not until_message: + embed = Embed( + color=Colour(Colours.soft_red), + title=random.choice(NEGATIVE_REPLIES), + description="`until_message` must be specified if `after_message` is specified." + ) + await ctx.send(embed=embed) + return + + # Check if the messages are not in same channel + if after_message.channel != until_message.channel: + embed = Embed( + color=Colour(Colours.soft_red), + title=random.choice(NEGATIVE_REPLIES), + description="You cannot do range clean across different channel." + ) + await ctx.send(embed=embed) + return + + # Ensure that after_message is younger than until_message + if after_message.created_at >= until_message.created_at: + embed = Embed( + color=Colour(Colours.soft_red), + title=random.choice(NEGATIVE_REPLIES), + description="`after` message must be younger than `until` message" + ) + await ctx.send(embed=embed) + return + # Are we already performing a clean? if self.cleaning: embed = Embed( @@ -106,8 +212,10 @@ 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: @@ -117,54 +225,68 @@ class Clean(Cog): self.mod_log.ignore(Event.message_delete, ctx.message.id) await ctx.message.delete() - 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, predicate=predicate) + else: + message_mappings, message_ids = await self._get_messages_from_channels( + amount=amount, + channels=channels, + predicate=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) + + # Creates ID like int object that would represent an object that is exactly 14 days old + minimum_time = int((time.time() - 14 * 24 * 60 * 60) * 1000.0 - 1420070400000) << 22 + + 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 message.id < minimum_time: + # 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( @@ -207,7 +329,10 @@ 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 = True + if channels: + use_cache = False + 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) @@ -241,11 +366,14 @@ 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 = True + if channels: + use_cache = False + 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, @@ -254,6 +382,18 @@ class Clean(Cog): until_message=message ) + @clean_group.command(name="from-to", aliases=["after-until", "range"]) + @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: |