diff options
| -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.""" |