import requests
import argparse
import csv
import os
import time

url = 'https://www.virustotal.com/vtapi/v2/file/report'
wanted = ['scan_id','md5','sha1','sha256','scan_date','total','positives']

md5s = []
def write_to_csv(csv_file,csv_columns,dict_data):
    with open(csv_file, 'w') as csvfile:
        writer = csv.DictWriter(csvfile, fieldnames=csv_columns)
        writer.writeheader()
        for data in dict_data:
            writer.writerow(data)
    return

def get_md5(md5_filepath):
    with open(md5_filepath) as f:
        md5s = [md.strip() for md in f.readlines()]
    return md5s

def get_report(apikey, md5):
    params = {'apikey': apikey, 
              'resource': md5}
    response = requests.get(url, params=params)
    if int(response.status_code) != 200:
        time.sleep(60)
        response = requests.get(url, params=params)
        if int(response.status_code) != 200:
            raise
    return response.json()

def parse_report(report_json):
    return {k:report_json[k] for k in wanted}

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--apikey",
        type=str,
        default=None,
        help="Provide VirusTotal api key for querying reports")

    parser.add_argument(
        "--md5file",
        type=str,
        default=None,
        help="Provide file location for md5 list")
    args = parser.parse_args()
    output_file = os.path.join(os.getcwd(),args.md5file+'_report.csv')
    md5s = get_md5(args.md5file)
    reports = []
    for md5 in md5s:
        report_j = get_report(args.apikey, md5)
        if report_j:
            print(md5)
        try:
            reports.append(parse_report(report_j))
        except:
            pass
        time.sleep(0)
    write_to_csv(output_file, wanted, reports)

