#!/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