diff options
Diffstat (limited to '')
| -rw-r--r-- | bot/cogs/antimalware.py | 4 | ||||
| -rw-r--r-- | tests/bot/cogs/test_antimalware.py | 22 | 
2 files changed, 26 insertions, 0 deletions
diff --git a/bot/cogs/antimalware.py b/bot/cogs/antimalware.py index c76bd2c60..7894ec48f 100644 --- a/bot/cogs/antimalware.py +++ b/bot/cogs/antimalware.py @@ -55,6 +55,10 @@ class AntiMalware(Cog):          if not message.attachments or not message.guild:              return +        # Ignore webhook and bot messages +        if message.webhook_id or message.author.bot: +            return +          # Check if user is staff, if is, return          # Since we only care that roles exist to iterate over, check for the attr rather than a User/Member instance          if hasattr(message.author, "roles") and any(role.id in STAFF_ROLES for role in message.author.roles): diff --git a/tests/bot/cogs/test_antimalware.py b/tests/bot/cogs/test_antimalware.py index ecb7abf00..f50c0492d 100644 --- a/tests/bot/cogs/test_antimalware.py +++ b/tests/bot/cogs/test_antimalware.py @@ -23,6 +23,8 @@ class AntiMalwareCogTests(unittest.IsolatedAsyncioTestCase):          }          self.cog = antimalware.AntiMalware(self.bot)          self.message = MockMessage() +        self.message.webhook_id = None +        self.message.author.bot = None          self.whitelist = [".first", ".second", ".third"]      async def test_message_with_allowed_attachment(self): @@ -48,6 +50,26 @@ class AntiMalwareCogTests(unittest.IsolatedAsyncioTestCase):          self.message.delete.assert_not_called() +    async def test_webhook_message_with_illegal_extension(self): +        """A webhook message containing an illegal extension should be ignored.""" +        attachment = MockAttachment(filename="python.disallowed") +        self.message.webhook_id = 697140105563078727 +        self.message.attachments = [attachment] + +        await self.cog.on_message(self.message) + +        self.message.delete.assert_not_called() + +    async def test_bot_message_with_illegal_extension(self): +        """A bot message containing an illegal extension should be ignored.""" +        attachment = MockAttachment(filename="python.disallowed") +        self.message.author.bot = 409107086526644234 +        self.message.attachments = [attachment] + +        await self.cog.on_message(self.message) + +        self.message.delete.assert_not_called() +      async def test_message_with_illegal_extension_gets_deleted(self):          """A message containing an illegal extension should send an embed."""          attachment = MockAttachment(filename="python.disallowed")  |