antivirus.py

import os, sys, hashlib, getpass, datetime
import requests

VERSION = "1.2.6"
# URL to download signatures from
SIGS_URLS = ["https://raw.githubusercontent.com/iam-py-test/malware-bazaar-AV/main/full_sha256.txt"]
# other hashes to look for
ADDED_HASHES = ["3a2b0defc9e4c9282291a93c655c0dd1bd5c28307869aeb54ee6f3d0f1ccf150"] # test file

# get the folder to scan
floc = input("Enter a directory to scan: ")
# all the threats detected
detected = {}
# the number of threats which could not be removed
removefailed = []
# the total number of files scanned
filescanned = 0
# how many files could not be scanned (i.e. system files) - this may need to be removed as to avoid scaring the user
totalfail = 0
# if removal has been run
removalrun = False
# debug log
debuglog = f"""MalwareBazaar AV (unofficial) - {VERSION}
Date: {datetime.datetime.now().isoformat()}
RUN INFO:
SIGS_URLS: {",".join(SIGS_URLS)}
ADDED_HASHES: {",".join(ADDED_HASHES)}

Log: 
"""
# hashes
hashes = []
hashes += ADDED_HASHES


try:
	# fetch the hashes. To prevent overloading abuse.ch's servers, the list is mirrored on GH by me
	for url in SIGS_URLS:
		hashes += requests.get(url).text.split("\r\n")
	# scan the folder specified by the user
	for root,dirs,files in os.walk(floc):
		for file in files:
			filescanned += 1
			try:
				fpath = os.path.abspath(os.path.join(root,file))
				chash = hashlib.sha256(open(fpath,"rb").read()).hexdigest()
				if chash in hashes:
					detected[fpath] = chash
			except Exception as err:
				# should add debug mode soon
				totalfail += 1
				debuglog += f"Error while scanning {fpath}: {err}" + "\n"
except KeyboardInterrupt:
	# if the user cancels it
	sys.exit()
except Exception as err:
	print(f"Error while scanning: {err}. Exiting...")
	sys.exit(1)

# output the data
print(f"\n\n{filescanned} files scanned\n{totalfail} files were unable to be scanned\n{len(detected)} threats detected")
print("Threats detected: ")
for d in detected:
	# print each detected file
	print(d)
if len(detected) > 0:
	if input("Remove threats? (y/n) ") == "y":
		removalrun = True
		for d in detected:
			try:
				os.remove(d)
			except:
				removefailed.append(d)
		if len(removefailed) > 0:
			print("\nThe following files could not be removed: ")
			for t in removefailed:
				print(t)

debuglog += "End run log\n\n"
if "--debug" in sys.argv:
	print("Saving debug log...")
	outdebuglog = open("mbz_debug_log.log",'a')
	outdebuglog.write(debuglog)
	outdebuglog.close()

try:
	logfilename = input("Input the location to save the log file, or press enter: ")
except:
	sys.exit()
if logfilename == "":
	sys.exit()

logcontent = f"""MalwareBazaar AV (unofficial) {VERSION} log file (ran at {datetime.datetime.now().isoformat()} by {getpass.getuser()})
{filescanned} files scanned. {len(detected)} files detected. {totalfail} files could not be scanned. {len(removefailed)} files could not be removed.
Threat removal run: {removalrun}
Files detected:
"""
for detectedfile in detected:
	hashmessage = detected[detectedfile]
	if hashmessage in ADDED_HASHES:
		hashmessage += " - added hash"
	logcontent += "{} ({})".format(detectedfile,hashmessage)
	if detectedfile in removefailed:
		logcontent += " <-- removal failed"
	logcontent += "\n"

try:
	logfilehandle = open(logfilename,'w',encoding="UTF-8")
	logfilehandle.write(logcontent)
	logfilehandle.close()
	print(f"Log saved to {logfilename}")
except Exception as err:
	print(f"Failed to save log file: {err}")