first commit

LICENSE

@@ -0,0 +1,22 @@
+MIT License
+
+Copyright (c) 2017 Hyrum S. Anderson
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+

README.md

@@ -0,0 +1,43 @@
+# dqeaf-gym-malware 
+
+A malware mutation agent trained using deep reinforcement learning
+
+Overview
+======
+The agent uses neural networks to mutate the malware by take the malware features as input and getting the action to perform as output. 
+
+
+Setup
+======
+
+The project built on Python3.6 we recommend first creating a virtualenv (details can be found [here]) with Python3.6 then performing the following actions ensure you have the correct python libraries:
+
+[here]: https://docs.python.org/3/tutorial/venv.html
+```sh
+pip install -r requirements.txt
+```
+It also uses LIEF to parse and modify malware. 
+
+[LIEF]: https://github.com/lief-project/LIEF
+
+Linux
+```
+pip install https://github.com/lief-project/LIEF/releases/download/0.7.0/linux_lief-0.7.0_py3.6.tar.gz
+```
+
+OSX
+```
+pip install https://github.com/lief-project/LIEF/releases/download/0.7.0/osx_lief-0.7.0_py3.6.tar.gz
+```
+
+Gym-Malware Environment
+====
+
+The gym-malware environment (https://github.com/endgameinc/gym-malware) was modified to extract only 518 out of 2350 features for the training of the agent i.e. byte histogram normalized to sum to unity and two-dimensional entropy histogram. Additionaly only 4 actions are used for the mutation i.e. append random bytes, append import, append section and remove signature.
+
+
+Resources
+======
+
+Gym-Malware Environment : https://github.com/endgameinc/gym-malware
+Deep Reinforcement Learning : https://ieeexplore.ieee.org/stamp/stamp.jsp?arnumber=8676031

download_samples.py

@@ -0,0 +1,163 @@
+#!/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/vtapi/v2/file/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)