spam: block harmful file extensions

spam/spam.py

@@ -28,6 +28,7 @@ def __init__(self, bot):
             "strings": {},
             "invites": {},
             "channels": {},
+            "files": [],
             "active": False,
             "feed": None
         }
@@ -94,6 +95,33 @@ async def _channel(self, ctx, channel: discord.TextChannel):
         else:
             await ctx.maybe_send_embed("Couldn't find channel.")
 
+    @spam.group(name="file")
+    async def file_group(self, ctx):
+        """ Add a file extension to block """
+        pass
+
+    @file_group.command(name="add")
+    async def file_add(self, ctx, extension: str):
+        """ Add a file extension to block """
+        extension = extension.lower().replace(".", "")
+        async with self.config.guild(ctx.guild).files() as files:
+            if extension not in files:
+                files.append(extension)
+                await ctx.maybe_send_embed(f"Added `{extension}` to blocked file extensions")
+            else:
+                await ctx.maybe_send_embed(f"`{extension}` is already blocked")
+
+    @file_group.command(name="remove")
+    async def file_remove(self, ctx, extension: str):
+        """ Remove a file extension from block list """
+        extension = extension.lower().replace(".", "")
+        async with self.config.guild(ctx.guild).files() as files:
+            if extension in files:
+                files.remove(extension)
+                await ctx.maybe_send_embed(f"Removed `{extension}` from blocked file extensions")
+            else:
+                await ctx.maybe_send_embed(f"`{extension}` is not in the block list")
+
     @spam.command()
     async def remove(self, ctx, name):
         """ Remove a text from spam protection detection """
@@ -158,6 +186,16 @@ async def channel(self, ctx):
         msg = "**Allowed Channels:**\n"+msg
         await menu(ctx, list(pagify(msg)), DEFAULT_CONTROLS)
 
+    @_list.command()
+    async def file(self, ctx):
+        """ Shows a list of blocked file extensions """
+        files = await self.config.guild(ctx.guild).files()
+        if not files:
+            await ctx.maybe_send_embed("List is empty.")
+            return
+        msg = "**Blocked File Extensions:**\n" + "\n".join(f"- `{extension}`" for extension in files)
+        await menu(ctx, list(pagify(msg)), DEFAULT_CONTROLS)
+
     @checks.admin_or_permissions(manage_roles=True)
     @spam.command()
     async def setfeed(self, ctx, channel_id):
@@ -179,6 +217,21 @@ async def on_message(self, ctx):
         if await self.bot.is_mod(ctx.author):
             return
 
+        if ctx.attachments:
+            files = await self.config.guild(ctx.guild).files()
+            for attachment in ctx.attachments:
+                file_extension = attachment.filename.split('.')[-1].lower()
+                if file_extension in files:
+                    await ctx.delete()
+                    feed = await self.config.guild(ctx.guild).feed()
+                    if feed:
+                        embed = discord.Embed(colour=discord.Colour(0xf5a623), description=f"Spam protection deleted a file with blocked extension `.{file_extension}` in <#{ctx.channel.id}>")
+                        embed.add_field(name="**Author:**", value=f"<@{ctx.author.id}>", inline=False)
+                        embed.add_field(name="**Message:**", value=ctx.content, inline=False)
+                        embed.add_field(name="**File:**", value=attachment.filename, inline=False)
+                        await self.bot.get_channel(int(await self.config.guild(ctx.guild).feed())).send(embed=embed)
+                    return
+
         find = INVITE_RE.findall(ctx.clean_content)
         # print(find)
         if find: