diff options
| -rw-r--r-- | bot/__main__.py | 1 | ||||
| -rw-r--r-- | bot/cogs/antimalware.py | 17 |
2 files changed, 12 insertions, 6 deletions
diff --git a/bot/__main__.py b/bot/__main__.py index 19a7e5ec6..f352cd60e 100644 --- a/bot/__main__.py +++ b/bot/__main__.py @@ -39,6 +39,7 @@ bot.load_extension("bot.cogs.logging") bot.load_extension("bot.cogs.security") # Commands, etc +bot.load_extension("bot.cogs.antimalware") bot.load_extension("bot.cogs.antispam") bot.load_extension("bot.cogs.bot") bot.load_extension("bot.cogs.clean") diff --git a/bot/cogs/antimalware.py b/bot/cogs/antimalware.py index e4688295e..94566c156 100644 --- a/bot/cogs/antimalware.py +++ b/bot/cogs/antimalware.py @@ -1,6 +1,6 @@ import logging -from discord import Message, utils +from discord import Message, NotFound from discord.ext.commands import Bot, Cog from bot.constants import AntiMalware as AntiMalwareConfig, Channels @@ -10,6 +10,7 @@ log = logging.getLogger(__name__) class AntiMalware(Cog): """Cog providing anti-malware behavior.""" + def __init__(self, bot: Bot): self.bot = bot self.whitelist = tuple(AntiMalwareConfig.whitelist) @@ -17,27 +18,31 @@ class AntiMalware(Cog): @Cog.listener() async def on_message(self, message: Message) -> None: """Identify messages with prohibited attachments.""" - log.trace("Entered AntiMalware.on_message()") rejected_attachments = [a for a in message.attachments if not a.filename.lower().endswith(self.whitelist)] detected_pyfile = len([a for a in message.attachments if a.filename.lower().endswith('.py')]) > 0 if len(rejected_attachments) > 0: - log.trace("Identified rejected attachment(s)") - # Send a message indicating the problem to the user (with special treatment for .py) + # Send a message to the user indicating the problem (with special treatment for .py) author = message.author if detected_pyfile: msg = f"{author.mention}, it looks like you tried to attach a Python file - please " \ f"use a code-pasting service such as https://paste.pythondiscord.com/ instead." else: - meta_channel = utils.get(message.guild.channels, id=Channels.meta) + meta_channel = self.bot.get_channel(Channels.meta) msg = f"{author.mention}, it looks like you tried to attach a file type we don't " \ f"allow. Feel free to ask in {meta_channel.mention} if you think this is a mistake." await message.channel.send(msg) + # Delete the offending message: + try: + await message.delete() + except NotFound: + log.info(f"Tried to delete message `{message.id}`, but message could not be found.") + def setup(bot: Bot) -> None: - """AntiMalware cog load.""" + """Antimalware cog load.""" bot.add_cog(AntiMalware(bot)) log.info("Cog loaded: AntiMalware") |