从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()
            return True

        print("{} downloading".format(sha256))
        sample_path = os.path.join(samples_path, sha256)
        success = vt_download_sample(sha256, sample_path, vtapikey)

        if not success:
            print("{} had a problem".format(sha256))

        print("{} done".format(sha256))
        download_queue.task_done()

def use_virustotal(args):
    """
    Use Virustotal to download the environment malware
    """
    m = multiprocessing.Manager()
    download_queue = m.JoinableQueue(args.nconcurrent)

    archive_procs = [
        multiprocessing.Process(
            target=download_worker_function,
            args=(download_queue, args.vtapikey))
        for i in range(args.nconcurrent)
    ]
    for w in archive_procs:
        w.start()

    for row in get_sample_hashes():
        download_queue.put(row["sha256"])

    for i in range(args.narchiveprocs):
        download_queue.put("STOP")

    download_queue.join()
    for w in archive_procs:
        w.join()

def use_virusshare(args):
    """
    Use VirusShare zip files as the source for the envirnment malware
    """
    pwd = bytes(args.zipfilepassword, "ascii")
    md5_to_sha256_dict = {d["md5"]: d["sha256"] for d in get_sample_hashes()}

    for path in args.zipfile:
        z = zipfile.ZipFile(path)
        for f in z.namelist():
            z_object_md5 = f.split("_")[1]
            if z_object_md5 in md5_to_sha256_dict:
                sample_bytez = z.open(f, "r", pwd).read()
                with open(md5_to_sha256_dict[z_object_md5], "wb") as ofile:
                    ofile.write(sample_bytez)
                print("Extracted {}".format(md5_to_sha256_dict[z_object_md5]))

if __name__ == '__main__':
    prog = "download_samples"
    descr = "Download the samples that define the malware gym environment"
    parser = argparse.ArgumentParser(prog=prog, description=descr)
    parser.add_argument(
        "--virustotal",
        default=False,
        action="store_true",
        help="Use Virustotal to download malware samples")
    parser.add_argument(
        "--vtapikey", type=str, default=None, help="Virustotal API key")
    parser.add_argument(
        "--nconcurrent",
        type=int,
        default=6,
        help="Maximum concurrent downloads from Virustotal")
    parser.add_argument(
        "--virusshare",
        default=False,
        action="store_true",
        help="Use malware samples from VirusShare torrents")
    parser.add_argument(
        "--zipfile",
        type=str,
        nargs="+",
        help="The path of VirusShare zipfile 290 or 291")
    parser.add_argument(
        "--zipfilepassword",
        type=str,
        default=None,
        help="Password for the VirusShare zipfiles 290 or 291")
    args = parser.parse_args()

    if not args.virustotal and not args.virusshare:
        parser.error("Must use either Virustotal or VirusShare")

    if args.virusshare:
        if len(args.zipfile) == 0:
            parser.error("Must the paths for one or more Virusshare zip files")

        if args.zipfilepassword is None:
            parser.error("Must enter a password for the VirusShare zip files")

        use_virusshare(args)

    if args.virustotal:
        if args.vtapikey is None:
            parser.error("Must enter a VirusTotal API key")

        use_virustotal(args)

使用方法：

python download_samples.py  --virustotal --vtapikey 1a7b7440ceca037b88fd160ef6c8e04b69ba434bdd76ef2ab0ab52a567xxxxx

csv文件格式：

sha256,sha1,md5
0007df5e92070f8d12411078070bdcafb24df81c837d8113a1e047ed7ac9fba1,e760b4ae027975928735024273a4240995442e2b,002e5581fabb21af4d4e7ec070561d38
0026b14f896934c621eccca48474353fff08f592ebc2949dde4b881f2353e3d2,f5cc8bd5accc281a8a41a9b13d870734361ec26b,292bd61f51ef0983b058a3b0f16ad263
00341b912ae7a9fc5bd25ac544bb2525cdc10f7dfcf51e6d96e9221a9ca06525,0329a4316eea3cf6d1376ea1eab5e2806258193b,c0370cb71216559beef7fe943b52003e
007792005ee9d835d5d0d4e0d6f7b886605272252a202e97a04bbc30bbbe12ae,f8190fe3936eff91a011901e30d66d0ad96e7e0c,64fed9d345dec9156090832c2b768982
009868767950256d823b0e9c6a89b8a7b2cef63424adc1840d1350ffa0bd3e42,50d4083adcd17910c2889842daf0d5e6ec41ab40,2f7a71e7abfd8536b9dee243656e0a8a
00a52b54695bac31830bdecf1c0e71b10da9bf3e9ff3d52cf1fc90f110458475,26f7549b66b2578112a77ceda7be7647ce5bc763,84b7490cae7fb84010863e006988951a
00c00e802109d0a3cb122c90168380ca23dfd3c28b1f03711b6218a8b1800f7c,eca4d84561c6440975ca64402e92ab01cf1bf4c8,d40cf7ae9174d5dc79c2e9db8cdb1bbb
0105e7aadf4e069b10aea00a43d90b753acfdd81c8db6e37df2c5b563162c30b,310a76a010cb58b510da8eb743f53ff517e441a9,2d4dcc983545014af6c8994ffa478488
0106fb2d96d5643f7ccb4a3e9fe8f3bb34c7d65d03333370648915991a3b200d,ad073b1ada3bb0aff0cae2edd7d41f6f09816cbe,f3e5b6b8c47211d54c2031d7a9a8f54f
012244e5a30708451b0b8b36a45e7d36fc8694f999adec739ef21efbc5f8e922,ef0ce82ca912e79a4fe64879ceb7fc30605367bb,ef24712cdbd8bd210e44d1546f5b91ab
014b392af2230b6275acf08a1384b1dd578e7fa3e7aba70c1b5b2ea6956c2108,43336578eb0efb1f9096836ff420fba635527020,d8f99268a5727a64bbac9a149b169afc
01640574490f32ba3d84bef60bdc30794edacf32932e93bada4d068dc5e27457,320c06678b0253ef5e30933d341a981744702c49,06c7dcdcdc887e052c2b6ee0dc88a2a6