diff options
| -rw-r--r-- | bot/cogs/antimalware.py | 41 | 
1 files changed, 24 insertions, 17 deletions
| diff --git a/bot/cogs/antimalware.py b/bot/cogs/antimalware.py index 9e9e81364..373619895 100644 --- a/bot/cogs/antimalware.py +++ b/bot/cogs/antimalware.py @@ -1,4 +1,5 @@  import logging +from os.path import splitext  from discord import Embed, Message, NotFound  from discord.ext.commands import Cog @@ -28,24 +29,30 @@ class AntiMalware(Cog):              return          embed = Embed() -        for attachment in message.attachments: -            filename = attachment.filename.lower() -            if filename.endswith('.py'): -                embed.description = ( -                    f"It looks like you tried to attach a Python file - please " -                    f"use a code-pasting service such as {URLs.site_schema}{URLs.site_paste}" -                ) -                break  # Other detections irrelevant because we prioritize the .py message. -            if not filename.endswith(tuple(AntiMalwareConfig.whitelist)): -                whitelisted_types = ', '.join(AntiMalwareConfig.whitelist) -                meta_channel = self.bot.get_channel(Channels.meta) -                embed.description = ( -                    f"It looks like you tried to attach a file type that we " -                    f"do not allow. We currently allow the following file " -                    f"types: **{whitelisted_types}**. \n\n Feel free to ask " -                    f"in {meta_channel.mention} if you think this is a mistake." -                ) +        file_extensions = {splitext(message.filename.lower())[1] for message in message.attachments} +        extensions_blocked = file_extensions - set(AntiMalwareConfig.whitelist) +        if ".py" in extensions_blocked: +            # Short-circuit on *.py files to provide a pastebin link +            embed.description = ( +                "It looks like you tried to attach a Python file - " +                f"please use a code-pasting service such as {URLs.site_schema}{URLs.site_paste}" +            ) +        elif extensions_blocked: +            blocked_extensions_str = ', '.join(extensions_blocked) +            whitelisted_types = ', '.join(AntiMalwareConfig.whitelist) +            meta_channel = self.bot.get_channel(Channels.meta) +            embed.description = ( +                f"It looks like you tried to attach file type(s) that we do not allow ({blocked_extensions_str}). " +                f"We currently allow the following file types: **{whitelisted_types}**.\n\n" +                f"Feel free to ask in {meta_channel.mention} if you think this is a mistake." +            ) +          if embed.description: +            log.info( +                f"User '{message.author}' ({message.author.id}) uploaded blacklisted file(s): {blocked_extensions_str}", +                extra={"attachment_list": [attachment.filename for attachment in message.attachments]} +            ) +              await message.channel.send(f"Hey {message.author.mention}!", embed=embed)              # Delete the offending message: | 
