#!/usr/bin/env python
import os
import sys
import json
import urllib
import urllib2
import hashlib
import argparse
 
VIRUSTOTAL_FILE_URL = 'https://www.virustotal.com/vtapi/v2/file/report'
API_KEY = ''
 
TPL_SECTION = "[*] ({0}):"
TPL_MATCH = "\t\_ Results: {0}/{1} {2}\n\t   SHA256: {3}\n\t   Scan Date: {4}"
TPL_SIGNATURES = "\t   Signatures:\n\t\t{0}"
 
def color(text, color_code):
    if sys.platform == "win32" and os.getenv("TERM") != "xterm":
        return text
 
    return '\x1b[%dm%s\x1b[0m' % (color_code, text)
 
def red(text):
    return color(text, 31)
 
def blink(text):
    return color(text, 05)
    
def Green(text):
    return color(text, 32)
 
def yellow(text):
    return color(text, 34)
 
class Hash(object):
    def __init__(self, path):
        self.path = path
        self.md5 = ''
        self.sha256 = ''
 
    def get_chunks(self):
        fd = open(self.path, 'rb')
        while True:
            chunk = fd.read(16 * 1024)
            if not chunk:
                break
 
            yield chunk
        fd.close()
 
    def calculate(self):
        md5 = hashlib.md5()
        sha256 = hashlib.sha256()
 
        for chunk in self.get_chunks():
            md5.update(chunk)
            sha256.update(chunk)
 
        self.md5 = md5.hexdigest()
        self.sha256 = sha256.hexdigest()
 
class Scanner(object):
    def __init__(self, key, path):
        self.key = key
        self.path = path
        self.list = []
 
    def populate(self):
        paths = []
 
        if os.path.isfile(self.path):
            paths.append(self.path)
        else:
            for root, folders, files in os.walk(self.path):
                for file_name in files:
                    # Skip hidden files, might need an option for this.
                    if file_name.startswith('.'):
                        continue
 
                    file_path = os.path.join(root, file_name)
                    if os.path.exists(file_path):
                        paths.append(file_path)
 
        for path in paths:
            hashes = Hash(path)
            hashes.calculate()
 
            self.list.append({
                'path' : path,
                'md5' : hashes.md5,
                'sha256' : hashes.sha256
                })
 
    def scan(self):
        hashes = []
        for entry in self.list:
            if entry['sha256'] not in hashes:
                hashes.append(entry['sha256'])
 
        data = urllib.urlencode({
            'resource' : ','.join(hashes),
            'apikey' : self.key
            })
 
        try:
            request = urllib2.Request(VIRUSTOTAL_FILE_URL, data)
            response = urllib2.urlopen(request)
            report = json.loads(response.read())
        except Exception as e:
            print(red("[!] ERROR: Cannot obtain results from VirusTotal: {0}\n".format(e)))
            return
 
        results = []
	CHECK = 0
	VIRUS = 0
	CLEAN = 0
 
        if type(report) is dict:
            results.append(report)
        elif type(report) is list:
            results = report
 
        for entry in results:
 
            sha256 = entry['resource']
 
            entry_paths = []
            for item in self.list:
                if item['sha256'] == sha256:
                    if item['path'] not in entry_paths:
                        entry_paths.append(item['path'])
 
            print(TPL_SECTION.format('\n     '.join(entry_paths))),
	    
            if entry['response_code'] == 0:
		CHECK = CHECK + 1
                print(yellow('< VirusTotal Hasnt Seen This File! Manually Verify!'))
 
            elif entry['positives'] < 1:
		CLEAN = CLEAN + 1		
		continue
 
            elif entry['positives'] > 0:
		VIRUS = VIRUS +1 
                print (red(blink('< This is a virus!!!!!!')))
	
	
	# PRINT STATS
	print ('\n')
	print ('Scanned %0.f files with virustotal.com:') % (CLEAN+CHECK+VIRUS)	
	print (Green('%0.f files are virus free.')) % CLEAN
	print (yellow('%0.f files need to be manually verified.')) % CHECK
	print (red('%0.f files are infected.')) % VIRUS
	
    def run(self):
        if not self.key:
            print(red("[!] ERROR: You didn't specify a valid VirusTotal API key.\n"))
            return
 
        if not os.path.exists(self.path):
            print(red("[!] ERROR: The target path {0} does not exist.\n".format(self.path)))
            return
 
        self.populate()
        self.scan()
 
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('path', type=str, help='Path to the file or folder to lookup on VirusTotal')
    parser.add_argument('--key', type=str, action='store', default=API_KEY, help='VirusTotal API key')
 
    try:
        args = parser.parse_args()
    except IOError as e:
        parser.error(e)
        sys.exit()
 
    scan = Scanner(args.key, args.path)
    scan.run()