#!/usr/bin/python

'''
antivirus - Version 0.1a (2016-06-06).
Latest version at https://www.itsec.se

Note: You will need a VirusTotal Public API key for this script
to work!
https://www.virustotal.com/en/documentation/public-api/#getting-started

A very simple antivirus to find known malware on your system.
The script uses VirusTotal (https://virustotal.com/), an
online service that aggregate the output of several different
antivirus engines.

The script searches through a specified path for all
executable files (using extensions commonly used by malware).
When an executable files is identified, a MD5-hash is calculated
and uploaded to VirusTotals online service. The MD5-hash is
compared against known malware signatures. If a malware is 
detected/suspected, additional information is provided. Since
the script uses the Public API, a 15 second wait period is
enforced after each request to VirusTotal. Only MD5-hashes
are submitted to VirusTotal, no sensitive information is 
included.

This is an alpha release with additional features being considered:
- progress bar (for 15 sec wait)
- option to use magic signature, rather than file signatures
- option to include additional file formats/extensions
- option to check content of compressed files (e.g. zip, tar)
- CSV format (to replace the current verbose output)
- support for multithreading
- Single file support
- catch IndexError for path parameter
- add threshold (e.g. at least 2 detections)

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the named License,
or any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.

Copyright (C) 2016 by Gustav Nordenskjold, Sweden.
'''

import os, sys, datetime
import urllib, urllib2, hashlib, json, time

APIKEY = '' # 4 requests/min
APIPUBLIC = True
URL = "https://www.virustotal.com/vtapi/v2/file/report"


FILE_EXT = [ 'bat', 'chm', 'cmd', 'com', 'cpl', 'dll', 'docm', 'dotm', 'drv', 'exe',
			'gadget', 'hta', 'jar', 'msc', 'msh', 'msh1', 'msh1xml', 'msh2', 'msh2xml',
			'mshxml', 'msi', 'msp', 'ocx', 'pif', 'potm', 'ppam', 'ppsm', 'pptm',
			'ps1', 'ps1xml', 'ps2', 'ps2xml', 'psc1', 'psc2', 'scf', 'scr', 'sldm',
			'swf', 'vb', 'vbe', 'vbs', 'vxd', 'ws', 'wsc', 'wsf', 'wsh', 'xlam', 'xlsm',
			'xltm' ]

def usage():
	print 'Usage: python ' + sys.argv[0] + ' [-p path] [-v] [-h]'
	print '-p path\t\tPath to include in file search (default is filesystem root).'
	print '-v\t\tVerbose output (print the path of all binaries).'
	print '-h\t\tThis help text.'
	print ''

def checksum_md5(file, hasher, blocksize=65536):
	buf = file.read(blocksize)
	while len(buf) > 0:
		hasher.update(buf)
		buf = file.read(blocksize)
	return hasher.hexdigest()

if __name__ == '__main__':

	print 'antivirus - check executable files against Virustotal API'
	print 'Version 0.1a (2016-06-06). Latest version at http://www.itsec.se'
		
	VERBOSE = False
	ROOT = ''
	
	for i in range(len(sys.argv)):
		if sys.argv[i][:2] == '-v':
			VERBOSE = True
		elif sys.argv[i][:2] == '-h':
			usage()
			sys.exit(0)
		elif sys.argv[i][:2] == '-p':
			ROOT = sys.argv[i+1]
		
	if not ROOT:
		path = sys.executable
		while os.path.split(path)[1]:
			path = os.path.split(path)[0]
		ROOT = path
	
	if not os.path.exists(ROOT):
		print 'Error: Path not valid.'
		sys.exit(0)
	else:
		print 'Scaning path: ' + os.path.abspath(ROOT)
	
	if APIKEY == '':
		print '\nError: You will need a VirusTotal Public API key for this script to work!'
		print 'Visit https://www.virustotal.com/en/documentation/public-api/#getting-started\n'
		sys.exit(0)
	
	for dirpath, dirnames, filenames in os.walk(ROOT):
		for filename in filenames:
				
			try:
				if filename.split(".")[-1].strip().lower() in FILE_EXT:
					
					filepath = os.path.join(dirpath, filename)
					
					if VERBOSE:
						print filepath
				
					md5 = checksum_md5( open(filepath, 'rb'), hashlib.md5() )
					parameters = {"resource": md5, "apikey": APIKEY}
					data = urllib.urlencode(parameters)
					req = urllib2.Request(URL, data)
					response = urllib2.urlopen(req)
					vt_json = json.loads(response.read())

					if vt_json['response_code'] == 1:
						if vt_json['positives'] != 0:
							print ''
							print 'Possible malware: ' + filename
							print 'Complete path:    ' + os.path.abspath(filepath)
							print 'More info:        ' + vt_json.get('permalink', 'Error - Unknown response_code from VirusTotal')
							print ''
						
					if APIPUBLIC:
						time.sleep(15) # The script uses a Public API Key that is limited to 4 req/min
				
			except Exception as e:
				print ''
				print 'Exception: ' + str(e)
				sys.exit(0)