From 0fc5a123a50a5433a60fb30c000acf146fc07650 Mon Sep 17 00:00:00 2001 From: bendiller Date: Sat, 12 Oct 2019 15:13:55 -0600 Subject: Create barebones Antimalware cog and config - detects bad file extensions --- bot/cogs/antimalware.py | 43 +++++++++++++++++++++++++++++++++++++++++++ bot/constants.py | 6 ++++++ config-default.yml | 6 ++++++ 3 files changed, 55 insertions(+) create mode 100644 bot/cogs/antimalware.py diff --git a/bot/cogs/antimalware.py b/bot/cogs/antimalware.py new file mode 100644 index 000000000..e4688295e --- /dev/null +++ b/bot/cogs/antimalware.py @@ -0,0 +1,43 @@ +import logging + +from discord import Message, utils +from discord.ext.commands import Bot, Cog + +from bot.constants import AntiMalware as AntiMalwareConfig, Channels + +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) + + @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) + 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) + 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) + + +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 1deeaa3b8..81f316d57 100644 --- a/bot/constants.py +++ b/bot/constants.py @@ -460,6 +460,12 @@ class AntiSpam(metaclass=YAMLGetter): rules: Dict[str, Dict[str, int]] +class AntiMalware(metaclass=YAMLGetter): + section = "anti_malware" + + whitelist: tuple + + class BigBrother(metaclass=YAMLGetter): section = 'big_brother' diff --git a/config-default.yml b/config-default.yml index 0dac9bf9f..30d505d6d 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,11 @@ anti_spam: max: 3 +anti_malware: + whitelist: ['.bmp', '.gif', '.jpg', '.jpeg', '.png', '.tiff', # Images + '.3gp', '.3g2', '.avi', '.h264', '.m4v', '.mkv', '.mov', '.mp4', '.mpeg', '.mpg', '.wmv' ] # Videos + + reddit: request_delay: 60 subreddits: -- cgit v1.2.3 From b7afb7979cb2a63869ef5cfa03b325a0be269ddb Mon Sep 17 00:00:00 2001 From: bendiller Date: Sat, 12 Oct 2019 17:01:53 -0600 Subject: Implement message deletion --- bot/__main__.py | 1 + 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") -- cgit v1.2.3 From feb08b23a4992dab67161b2dd956bb92ef1f04ea Mon Sep 17 00:00:00 2001 From: bendiller Date: Sat, 12 Oct 2019 17:34:29 -0600 Subject: Fix Constants.AntiMalware.whitelist type --- bot/constants.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bot/constants.py b/bot/constants.py index 81f316d57..aecd6be59 100644 --- a/bot/constants.py +++ b/bot/constants.py @@ -463,7 +463,7 @@ class AntiSpam(metaclass=YAMLGetter): class AntiMalware(metaclass=YAMLGetter): section = "anti_malware" - whitelist: tuple + whitelist: list class BigBrother(metaclass=YAMLGetter): -- cgit v1.2.3 From 003f8fe85cbb90d51225580bc35b91ebf21fd0a3 Mon Sep 17 00:00:00 2001 From: bendiller Date: Mon, 14 Oct 2019 17:59:37 -0600 Subject: Improve code readability and docstring --- bot/cogs/antimalware.py | 25 ++++++++++++++----------- bot/constants.py | 1 + config-default.yml | 20 ++++++++++++++++++-- 3 files changed, 33 insertions(+), 13 deletions(-) diff --git a/bot/cogs/antimalware.py b/bot/cogs/antimalware.py index 94566c156..156239a63 100644 --- a/bot/cogs/antimalware.py +++ b/bot/cogs/antimalware.py @@ -9,29 +9,32 @@ log = logging.getLogger(__name__) class AntiMalware(Cog): - """Cog providing anti-malware behavior.""" + """Delete messages which contain attachments with non-whitelisted file extensions.""" def __init__(self, bot: Bot): self.bot = bot - self.whitelist = tuple(AntiMalwareConfig.whitelist) @Cog.listener() async def on_message(self, message: Message) -> None: """Identify messages with prohibited attachments.""" - 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: + rejected_attachments = list() + detected_pyfile = list() + for attachment in message.attachments: + if not attachment.filename.lower().endswith(tuple(AntiMalwareConfig.whitelist)): + rejected_attachments.append(attachment) + if attachment.filename.lower().endswith('.py'): + detected_pyfile.append(attachment) + + if 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." + 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." + 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) diff --git a/bot/constants.py b/bot/constants.py index aecd6be59..13f25e4f8 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 diff --git a/config-default.yml b/config-default.yml index 30d505d6d..071478206 100644 --- a/config-default.yml +++ b/config-default.yml @@ -324,8 +324,24 @@ anti_spam: anti_malware: - whitelist: ['.bmp', '.gif', '.jpg', '.jpeg', '.png', '.tiff', # Images - '.3gp', '.3g2', '.avi', '.h264', '.m4v', '.mkv', '.mov', '.mp4', '.mpeg', '.mpg', '.wmv' ] # Videos + whitelist: + - '.3gp' + - '.3g2' + - '.avi' + - '.bmp' + - '.gif' + - '.h264' + - '.jpg' + - '.jpeg' + - '.m4v' + - '.mkv' + - '.mov' + - '.mp4' + - '.mpeg' + - '.mpg' + - '.png' + - '.tiff' + - '.wmv' reddit: -- cgit v1.2.3 From d9b4f391f914b680bdf4d5b93882abe11b8dc9d4 Mon Sep 17 00:00:00 2001 From: bendiller Date: Wed, 16 Oct 2019 17:42:48 -0600 Subject: Improve code readability and provide early exit from loop --- bot/cogs/antimalware.py | 21 +++++++++++++-------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/bot/cogs/antimalware.py b/bot/cogs/antimalware.py index 156239a63..2ef61e8ad 100644 --- a/bot/cogs/antimalware.py +++ b/bot/cogs/antimalware.py @@ -17,24 +17,29 @@ class AntiMalware(Cog): @Cog.listener() async def on_message(self, message: Message) -> None: """Identify messages with prohibited attachments.""" - rejected_attachments = list() - detected_pyfile = list() + rejected_attachments = False + detected_pyfile = False for attachment in message.attachments: if not attachment.filename.lower().endswith(tuple(AntiMalwareConfig.whitelist)): - rejected_attachments.append(attachment) + rejected_attachments = True if attachment.filename.lower().endswith('.py'): - detected_pyfile.append(attachment) + detected_pyfile = True + break # Other detections irrelevant because we prioritize the .py message. if 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.") + 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.") + 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) -- cgit v1.2.3 From 0d164adf6ac13611f1fc66825998d20d26f240ac Mon Sep 17 00:00:00 2001 From: bendiller Date: Thu, 17 Oct 2019 14:08:13 -0600 Subject: Address reviewer request --- bot/cogs/antimalware.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/bot/cogs/antimalware.py b/bot/cogs/antimalware.py index 2ef61e8ad..b8c12accb 100644 --- a/bot/cogs/antimalware.py +++ b/bot/cogs/antimalware.py @@ -20,13 +20,14 @@ class AntiMalware(Cog): rejected_attachments = False detected_pyfile = False for attachment in message.attachments: - if not attachment.filename.lower().endswith(tuple(AntiMalwareConfig.whitelist)): - rejected_attachments = True 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 + break - if rejected_attachments: + 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: -- cgit v1.2.3 From 8722e954af926e7ec6b480df9edde1249c8ab795 Mon Sep 17 00:00:00 2001 From: bendiller Date: Thu, 17 Oct 2019 14:23:23 -0600 Subject: Bugfix - ensure .py attachment is prioritized over other non-whitelisted --- bot/cogs/antimalware.py | 1 - 1 file changed, 1 deletion(-) diff --git a/bot/cogs/antimalware.py b/bot/cogs/antimalware.py index b8c12accb..ababd6f18 100644 --- a/bot/cogs/antimalware.py +++ b/bot/cogs/antimalware.py @@ -25,7 +25,6 @@ class AntiMalware(Cog): break # Other detections irrelevant because we prioritize the .py message. if not attachment.filename.lower().endswith(tuple(AntiMalwareConfig.whitelist)): rejected_attachments = True - break if detected_pyfile or rejected_attachments: # Send a message to the user indicating the problem (with special treatment for .py) -- cgit v1.2.3