Merge pull request #528 from bendiller/antimalware-cog

Antimalware cog

4 files changed, 86 insertions, 0 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
new file mode 100644
index 000000000..ababd6f18
--- /dev/null
+++ b/bot/cogs/antimalware.py
@@ -0,0 +1,56 @@
+import logging
+
+from discord import Message, NotFound
+from discord.ext.commands import Bot, Cog
+
+from bot.constants import AntiMalware as AntiMalwareConfig, Channels
+
+log = logging.getLogger(__name__)
+
+
+class AntiMalware(Cog):
+    """Delete messages which contain attachments with non-whitelisted file extensions."""
+
+    def __init__(self, bot: Bot):
+        self.bot = bot
+
+    @Cog.listener()
+    async def on_message(self, message: Message) -> None:
+        """Identify messages with prohibited attachments."""
+        rejected_attachments = False
+        detected_pyfile = False
+        for attachment in message.attachments:
+            if attachment.filename.lower().endswith('.py'):
+                detected_pyfile = True
+                break  # Other detections irrelevant because we prioritize the .py message.
+            if not attachment.filename.lower().endswith(tuple(AntiMalwareConfig.whitelist)):
+                rejected_attachments = True
+
+        if detected_pyfile or rejected_attachments:
+            # 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 = 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."""
+    bot.add_cog(AntiMalware(bot))
+    log.info("Cog loaded: AntiMalware")
diff --git a/bot/constants.py b/bot/constants.py
index f4f45eb2c..4beae84e9 100644
--- a/bot/constants.py
+++ b/bot/constants.py
@@ -345,6 +345,7 @@ class Channels(metaclass=YAMLGetter):
     help_7: int
     helpers: int
     message_log: int
+    meta: int
     mod_alerts: int
     modlog: int
     off_topic_0: int
@@ -460,6 +461,12 @@ class AntiSpam(metaclass=YAMLGetter):
     rules: Dict[str, Dict[str, int]]
 
 
+class AntiMalware(metaclass=YAMLGetter):
+    section = "anti_malware"
+
+    whitelist: list
+
+
 class BigBrother(metaclass=YAMLGetter):
     section = 'big_brother'
 
diff --git a/config-default.yml b/config-default.yml
index ca405337e..197743296 100644
--- a/config-default.yml
+++ b/config-default.yml
@@ -107,6 +107,7 @@ guild:
         help_7:                           587375768556797982
         helpers:                          385474242440986624
         message_log:       &MESSAGE_LOG   467752170159079424
+        meta:                             429409067623251969
         mod_alerts:                       473092532147060736
         modlog:            &MODLOG        282638479504965634
         off_topic_0:                      291284109232308226
@@ -322,6 +323,27 @@ anti_spam:
             max: 3
 
 
+anti_malware:
+    whitelist:
+        - '.3gp'
+        - '.3g2'
+        - '.avi'
+        - '.bmp'
+        - '.gif'
+        - '.h264'
+        - '.jpg'
+        - '.jpeg'
+        - '.m4v'
+        - '.mkv'
+        - '.mov'
+        - '.mp4'
+        - '.mpeg'
+        - '.mpg'
+        - '.png'
+        - '.tiff'
+        - '.wmv'
+
+
 reddit:
     request_delay: 60
     subreddits: