#!/usr/bin/python
# -*- coding: utf-8 -*-

import requests, json
import sys, os.path
import hashlib
import time

apikey = '9bb5fb16ced08cccd6adce0efcb162e9e23e0f8acb35ea114d68fb4dea9f1b54'

def phelp():
    print("""
- pyrustotal -

Usage: ./pyrustotal.py [OPTION] [FILE]

Mandatory arguments:
    help        prints the help (hint, it's this one)
    search      searches the file's hash on VirusTotal for previous scans
                and retrieves the scan report
    scan        same as `search' except that it scans the file it if it's not
                found or rescans it if it's found
""")

def hashfile():
    f = open(sys.argv[2], 'rb')
    shash = hashlib.sha256()
    s = None;
    while s is None or len(s) > 0:
        s = f.read(1024)
        shash.update(s)
    return shash.hexdigest()

def search():
    if len(sys.argv) == 2 or not os.path.isfile(sys.argv[2]):
        print("You have to provide a file, exiting.")
        sys.exit(1)
    url = 'https://www.virustotal.com/vtapi/v2/file/report'
    params = {
            "resource": hashfile(),
            "apikey": apikey,
            }
    req = requests.post(url, params)
    if req.status_code != requests.codes.ok:
        print("Exceeded public API request rate, please try again in a",
                "minute, or switch to private API.")
        sys.exit(1)
    data = json.loads(req.text)
    print(data['verbose_msg'])
    if data['response_code']:
        print("\nScan date: %s\n\nPositives: %s\nTotal: %s" %
                (data['scan_date'], data['positives'], data['total']))
        return True

def scan():
    if search():
        opt = input("Do you want to rescan the file? [y/N] ")
        if opt == 'n' or opt == 'N' or opt == '':
            sys.exit(0)
        else:
            rescan()
            return
    url = 'https://www.virustotal.com/vtapi/v2/file/scan'
    params = {
            "file": open(sys.argv[2], 'rb'),
            "apikey": apikey,
            }
    req = requests.post(url, files=params)
    data = json.loads(req.text)
    if data['response_code']:
        print(data['verbose_msg'], "\nResource ID: %s" % data['resource'])

def rescan():
    url = 'https://www.virustotal.com/vtapi/v2/file/rescan'
    params = {
            "resource": hashfile(),
            "apikey": apikey,
            }
    req = requests.post(url, params)
    data = json.loads(req.text)
    print("\nWaiting for the result...")
    time.sleep(10)
    search()

args = {
        'help': phelp,
        'search': search,
        'scan': scan,
        }

if len(sys.argv) == 1:
    print("No arguments given.\nYou should try the `help' one.")
else:
    try:
        args[sys.argv[1]]()
    except KeyError as e:
        print("No function called %s." % e)