从virustotal上下载病毒样本

 

#!/usr/bin/env python

import os
import csv
#import Queue
import zipfile
import requests
import argparse
import multiprocessing

# TODO: Don't hardcode the relative path?
samples_path = "gym_malware/envs/utils/samples/"
hashes_path = "gym_malware/envs/utils/sample_hashes.csv"
vturl = "https://www.virustotal.com/intelligence/download"


def get_sample_hashes():
    hash_rows = []
    with open(hashes_path) as csvfile:
        for row in csv.DictReader(csvfile):
            hash_rows.append(row)
    return hash_rows


def vt_download_sample(sha256, sample_path, vtapikey):
    tries = 0
    success = False
    while not success and tries < 10:
        resp = requests.get(vturl, params={"hash": sha256, "apikey": vtapikey})

        if not resp.ok:
            tries += 1
            continue

        else:
            success = True

    if not success:
        return False

    with open(sample_path, "wb") as ofile:
        ofile.write(resp.content)

    return True


def download_worker_function(download_queue, vtapikey):
    while True:
        try:
            sha256 = download_queue.get()
        except queue.Empty:
            continue

        if sha256 == "STOP":
            download_queue.task_done()