From c2ff2261380f5642c42983a7261b15208737ab38 Mon Sep 17 00:00:00 2001 From: Marc Ahlgrim Date: Mon, 3 Oct 2022 12:23:25 +0200 Subject: [PATCH] VirusTotal Scanner Cog added Signed-off-by: Marc Ahlgrim --- classes/bot.py | 12 +++ classes/config.py | 4 + cogs/scanner.py | 188 ++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 204 insertions(+) create mode 100644 cogs/scanner.py diff --git a/classes/bot.py b/classes/bot.py index b93ef32..26ecd17 100644 --- a/classes/bot.py +++ b/classes/bot.py @@ -188,6 +188,12 @@ class OMFBot(AutoShardedBot): except Exception as e: print(f"ERROR in setup_hook: {e}") + # transmit the VirusTotal Token to the scanner object + scanner = self.get_cog('Scanner') + if scanner is not None: + scanner.registervTotal(self.configData.getVToken()) + print(f'TOKEN: {self.configData.getVToken()}') + # ###################################################### # send_random_message is called when the server is idle @@ -384,6 +390,12 @@ class OMFBot(AutoShardedBot): async def on_message(self, message : discord.Message ): + try: + x=message.author + except Exception as e: + # message has been deleted by the scanner + return + # ignore ephemeral messages if message.flags.ephemeral: return diff --git a/classes/config.py b/classes/config.py index a823522..ecd42cb 100644 --- a/classes/config.py +++ b/classes/config.py @@ -43,6 +43,10 @@ class Config(): secretNode=self.getNode("secret") return secretNode.get('BOT_TOKEN') + def getVToken(self): + secretNode=self.getNode("secret") + return secretNode.get('VTOTAL_TOKEN') + def getClientID(self): secretNode=self.getNode("secret") return secretNode.get('CLIENT_ID') diff --git a/cogs/scanner.py b/cogs/scanner.py new file mode 100644 index 0000000..686b6f9 --- /dev/null +++ b/cogs/scanner.py @@ -0,0 +1,188 @@ +# ##################################### +# +# a virustotal scanner cog +# listens for messages and scans URLs, +# domains and files +# +# ##################################### + +from asyncore import loop +import hashlib +from discord.ext import commands +import discord +import re +import vt +import asyncio + +# ##################################### +# +# some return values and constants +# +# ##################################### + + +stFile=1 +stURL=2 +stDomain=3 + +rvHarmless="harmless" +rvSuspicious="suspicious" +rvMalicious="malicious" +rvUnknown="unknown" + + +# ##################################### +# +# Class definition +# +# ##################################### + +class Scanner(commands.Cog): + + vTotalAPIKEY=None + vtClient = None + + # ############################################ + # init routine + # ############################################ + + def __init__(self, bot): + self.bot = bot + genericDomainString='[-a-zA-Z0-9@:%._\\+~#=]{1,256}\\.[a-zA-Z0-9()]{1,6}\\b(?:[-a-zA-Z0-9()@:%_\\+.~#?&\\/=]*)' + self.domainRegExp=f'{genericDomainString}' + self.urlRegExp =f'(?:http|ftp|sftp)s?://{genericDomainString}' + + # ############################################ + # registervTotal creates the VirusTotal API + # ############################################ + + def registervTotal(self,vToken): + self.vTotalAPIKEY = vToken + + + # ############################################################## + # getVerdict submits domains, URLS and/or files to Virustotal + # for scan + # ############################################################## + + async def getVerdict(self,vtObjectSring): + vtInfo:vt.Object + theVerdict=rvHarmless + try: + vtInfo=await self.vtClient.get_object_async(vtObjectSring) + except Exception as e: + theVerdict=rvUnknown + print(f"scanner vtInfo returned {e}") + + if (not theVerdict == rvUnknown): + vtInfoAnalysis=vtInfo.get("last_analysis_stats") + print(f"scanner vtInfo returned {vtInfoAnalysis}") + m=vtInfoAnalysis["malicious"] + h=vtInfoAnalysis["harmless"] + s=vtInfoAnalysis["suspicious"] + if (s>h): + theVerdict = rvSuspicious + if (m>h): + theVerdict = rvMalicious + return theVerdict + + # ############################################################## + # scan_message submits Attachments, URLs to getVerdict + # ############################################################## + + async def scan_message(self,msg): + + if msg.author.id == self.bot.user.id: + return + + # if we have not yet used the client or if the + # client had been closed in between then we re- + # open it + + if (self.vtClient is None): + try: + self.vtClient = vt.Client(self.vTotalAPIKEY) + except Exception as e: + print(f"Error in scan_message: {e}") + return + + # check for domain names + checkEntries=re.findall(self.domainRegExp,msg.content.replace("\n",""),re.IGNORECASE) + if len(checkEntries) >0: + for c in checkEntries: + print("{} is a Domain".format(c)) + + # check for links + checkEntries=re.findall(self.urlRegExp,msg.content.replace("\n",""),re.IGNORECASE) + if len(checkEntries) >0: + for c in checkEntries: + newMessage="**SECURITY WARNING**\nAll URLs are scanned\n" + xmsg=await msg.reply(newMessage) + print("{} is a URL".format(c)) + url_id = vt.url_id(c) + theVerdict = await self.getVerdict(f"/urls/{url_id}") + newMessage=f'{newMessage}\n>> Scan result: **{theVerdict}**\n' + await xmsg.edit(content=newMessage) + if (theVerdict==rvMalicious): + await xmsg.edit(content=f'**MESSAGE HAS BEEN REMOVED FOR SECURITY REASONS**\n') + await msg.delete() + # take further action + break + print("URL Scan finished") + if (not theVerdict==rvSuspicious): + await xmsg.delete() + + # check embedded files + if len(msg.attachments) > 0: + newMessage="**SECURITY WARNING**\nFiles posted here are systematically scanned\nEven if the scan is negative the file may still be malicious\n" + xmsg=await msg.reply(newMessage) + for theAttachment in msg.attachments: + newMessage=f'{newMessage}\n>> hashing file {theAttachment.filename}' + await xmsg.edit(content=newMessage) + attachmentContent = await theAttachment.read() + sha256String = hashlib.sha256(attachmentContent).hexdigest(); + newMessage=f'{newMessage}\n>> submitting hash to Scan Engine' + await xmsg.edit(content=newMessage) + theVerdict=await self.getVerdict(f'/files/{sha256String}') + newMessage=f'{newMessage}\n>> Scan result: **{theVerdict}**\n' + await xmsg.edit(content=newMessage) + if (theVerdict==rvMalicious): + await xmsg.edit(content=f'**FILE HAS BEEN REMOVED FOR SECURITY REASONS**\n') + await msg.delete() + # take further action + + + # ##################################### + # hook into edited messages + # ##################################### + + + @commands.Cog.listener() + async def on_raw_message_edit(self, rawdata: discord.RawMessageUpdateEvent): + channel = await self.bot.fetch_channel(rawdata.channel_id) + try: + msg = await channel.fetch_message(rawdata.message_id) + except Exception as e: + # message has been deleted by the scanner + return + if msg.author.id == self.bot.user.id: + return + asyncio.create_task(self.scan_message(msg)) + #await self.scan_message(msg) + + # ##################################### + # hook into new messages + # ##################################### + + @commands.Cog.listener() + async def on_message(self, msg): + if msg.author.id == self.bot.user.id: + return + asyncio.create_task(self.scan_message(msg)) + +# ##################################### +# The cog init +# ##################################### + +async def setup(bot): + await bot.add_cog(Scanner(bot))