diff options
Diffstat (limited to '')
| -rw-r--r-- | bot/cogs/token_remover.py | 17 | ||||
| -rw-r--r-- | bot/cogs/webhook_remover.py | 14 | ||||
| -rw-r--r-- | tests/bot/cogs/test_token_remover.py | 32 | 
3 files changed, 47 insertions, 16 deletions
diff --git a/bot/cogs/token_remover.py b/bot/cogs/token_remover.py index d55e079e9..ef979f222 100644 --- a/bot/cogs/token_remover.py +++ b/bot/cogs/token_remover.py @@ -4,7 +4,7 @@ import logging  import re  import typing as t -from discord import Colour, Message +from discord import Colour, Message, NotFound  from discord.ext.commands import Cog  from bot import utils @@ -63,6 +63,10 @@ class TokenRemover(Cog):          See: https://discordapp.com/developers/docs/reference#snowflakes          """ +        # Ignore DMs; can't delete messages in there anyway. +        if not msg.guild or msg.author.bot: +            return +          found_token = self.find_token_in_message(msg)          if found_token:              await self.take_action(msg, found_token) @@ -79,7 +83,13 @@ class TokenRemover(Cog):      async def take_action(self, msg: Message, found_token: Token) -> None:          """Remove the `msg` containing the `found_token` and send a mod log message."""          self.mod_log.ignore(Event.message_delete, msg.id) -        await msg.delete() + +        try: +            await msg.delete() +        except NotFound: +            log.debug(f"Failed to remove token in message {msg.id}: message already deleted.") +            return +          await msg.channel.send(DELETION_MESSAGE_TEMPLATE.format(mention=msg.author.mention))          log_message = self.format_log_message(msg, found_token) @@ -112,9 +122,6 @@ class TokenRemover(Cog):      @classmethod      def find_token_in_message(cls, msg: Message) -> t.Optional[Token]:          """Return a seemingly valid token found in `msg` or `None` if no token is found.""" -        if msg.author.bot: -            return -          # Use finditer rather than search to guard against method calls prematurely returning the          # token check (e.g. `message.channel.send` also matches our token pattern)          for match in TOKEN_RE.finditer(msg.content): diff --git a/bot/cogs/webhook_remover.py b/bot/cogs/webhook_remover.py index 1b5c3f821..543869215 100644 --- a/bot/cogs/webhook_remover.py +++ b/bot/cogs/webhook_remover.py @@ -1,7 +1,7 @@  import logging  import re -from discord import Colour, Message +from discord import Colour, Message, NotFound  from discord.ext.commands import Cog  from bot.bot import Bot @@ -35,7 +35,13 @@ class WebhookRemover(Cog):          """Delete `msg` and send a warning that it contained the Discord webhook `redacted_url`."""          # Don't log this, due internal delete, not by user. Will make different entry.          self.mod_log.ignore(Event.message_delete, msg.id) -        await msg.delete() + +        try: +            await msg.delete() +        except NotFound: +            log.debug(f"Failed to remove webhook in message {msg.id}: message already deleted.") +            return +          await msg.channel.send(ALERT_MESSAGE_TEMPLATE.format(user=msg.author.mention))          message = ( @@ -59,6 +65,10 @@ class WebhookRemover(Cog):      @Cog.listener()      async def on_message(self, msg: Message) -> None:          """Check if a Discord webhook URL is in `message`.""" +        # Ignore DMs; can't delete messages in there anyway. +        if not msg.guild or msg.author.bot: +            return +          matches = WEBHOOK_URL_RE.search(msg.content)          if matches:              await self.delete_and_respond(msg, matches[1] + "xxx") diff --git a/tests/bot/cogs/test_token_remover.py b/tests/bot/cogs/test_token_remover.py index a10124d2d..3349caa73 100644 --- a/tests/bot/cogs/test_token_remover.py +++ b/tests/bot/cogs/test_token_remover.py @@ -3,7 +3,7 @@ from re import Match  from unittest import mock  from unittest.mock import MagicMock -from discord import Colour +from discord import Colour, NotFound  from bot import constants  from bot.cogs import token_remover @@ -121,15 +121,16 @@ class TokenRemoverTests(unittest.IsolatedAsyncioTestCase):          find_token_in_message.assert_called_once_with(self.msg)          take_action.assert_not_awaited() -    @autospec("bot.cogs.token_remover", "TOKEN_RE") -    def test_find_token_ignores_bot_messages(self, token_re): -        """The token finder should ignore messages authored by bots.""" -        self.msg.author.bot = True +    @autospec(TokenRemover, "find_token_in_message") +    async def test_on_message_ignores_dms_bots(self, find_token_in_message): +        """Shouldn't parse a message if it is a DM or authored by a bot.""" +        cog = TokenRemover(self.bot) +        dm_msg = MockMessage(guild=None) +        bot_msg = MockMessage(author=MagicMock(bot=True)) -        return_value = TokenRemover.find_token_in_message(self.msg) - -        self.assertIsNone(return_value) -        token_re.finditer.assert_not_called() +        for msg in (dm_msg, bot_msg): +            await cog.on_message(msg) +            find_token_in_message.assert_not_called()      @autospec("bot.cogs.token_remover", "TOKEN_RE")      def test_find_token_no_matches(self, token_re): @@ -281,6 +282,19 @@ class TokenRemoverTests(unittest.IsolatedAsyncioTestCase):              channel_id=constants.Channels.mod_alerts          ) +    @mock.patch.object(TokenRemover, "mod_log", new_callable=mock.PropertyMock) +    async def test_take_action_delete_failure(self, mod_log_property): +        """Shouldn't send any messages if the token message can't be deleted.""" +        cog = TokenRemover(self.bot) +        mod_log_property.return_value = mock.create_autospec(ModLog, spec_set=True, instance=True) +        self.msg.delete.side_effect = NotFound(MagicMock(), MagicMock()) + +        token = mock.create_autospec(Token, spec_set=True, instance=True) +        await cog.take_action(self.msg, token) + +        self.msg.delete.assert_called_once_with() +        self.msg.channel.send.assert_not_awaited() +  class TokenRemoverExtensionTests(unittest.TestCase):      """Tests for the token_remover extension."""  |