From 6792a96216370ac3a922d9857b1288f1ef286783 Mon Sep 17 00:00:00 2001
From: Travis <trav.matth@gmail.com>
Date: Thu, 19 Mar 2020 04:48:18 +0000
Subject: [PATCH] combine afl crash dirs into one, closes #328

---
 bin/deepstate/executors/fuzz/afl.py | 38 +++++++++++++++++++--
 tests/test_afl_combine_crashes.py   | 51 +++++++++++++++++++++++++++++
 2 files changed, 87 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_afl_combine_crashes.py

diff --git a/bin/deepstate/executors/fuzz/afl.py b/bin/deepstate/executors/fuzz/afl.py
index 3eca512c..58dd6503 100644
--- a/bin/deepstate/executors/fuzz/afl.py
+++ b/bin/deepstate/executors/fuzz/afl.py
@@ -16,6 +16,10 @@
 import os
 import logging
 import argparse
+import re
+import shutil
+from collections import defaultdict
+
 
 from typing import List, Dict, Optional
 
@@ -174,14 +178,44 @@ def _sync_seeds(self, src, dest, excludes=[]) -> None:
     super()._sync_seeds(src, dest, excludes=excludes)
 
 
+  def consolidate_crash_dirs(self) -> None:
+    regex = re.compile(".*crashes.*")
+    crash_files = defaultdict(set)
+    crash_files_path = os.path.join(self.output_test_dir, "the_fuzzer")
+    for directory in os.listdir(crash_files_path):
+        crash_dir = os.path.join(crash_files_path, directory)
+        L.debug("Inspecting AFL crash dir: %s", crash_dir)
+        if not os.path.isdir(crash_dir) or not regex.match(crash_dir):
+            continue
+        for f in os.listdir(crash_dir):
+            crash_path = os.path.join(crash_dir, f)
+            if f == "README.txt" or not os.path.isfile(crash_path):
+                continue
+            crash_files[f].add(directory)
+    dirs_to_delete = set()
+    for name, paths in crash_files.items():
+        for p in paths:
+            if p == "crashes":
+                continue
+            dirs_to_delete.add(os.path.join(crash_files_path, p))
+            old = os.path.join(crash_files_path, p, name)
+            new = os.path.join(crash_files_path, "crashes", name)
+            L.debug("Moving crash report %s to %s", old, new)
+            os.rename(old, new)
+    for d in dirs_to_delete:
+        L.debug("Deleting crash directory %s", d)
+        shutil.rmtree(d, onerror=lambda err: print(err))
+
+
   def post_exec(self) -> None:
     """
     AFL post_exec outputs last updated fuzzer stats,
     and (TODO) performs crash triaging with seeds from
-    both sync_dir and local queue.
+    both sync_dir and local queue. Merges
+    output_test_dir/the_fuzzer/crashes* into one dir
     """
-    # TODO: merge output_test_dir/the_fuzzer/crashes* into one dir
     super().post_exec()
+    self.consolidate_crash_dirs()
 
 
 def main():
diff --git a/tests/test_afl_combine_crashes.py b/tests/test_afl_combine_crashes.py
new file mode 100644
index 00000000..7bd6581b
--- /dev/null
+++ b/tests/test_afl_combine_crashes.py
@@ -0,0 +1,51 @@
+import os
+import shutil
+import logging
+from unittest import TestCase
+from deepstate.executors.fuzz.afl import AFL
+
+
+from typing import Text
+
+
+L = logging.getLogger(__name__)
+
+
+expected_contents = set(["0", "1", "2", "3",
+                        "4", "5", "6", "7",
+                        "8", "README.txt"])
+
+
+def create_crash_dirs(path: Text) -> None:
+  dir_i, crash_i = 0, 0
+  while dir_i < 3:
+    i, crash_dir = 0, os.path.join(path, "crashes")
+    if dir_i:
+      crash_dir += ".dir_" + str(dir_i)
+    os.mkdir(crash_dir, 0o777)
+    crash_file = os.path.join(crash_dir, "README.txt")
+    with open(crash_file, "w") as f:
+      f.write(crash_file)
+    while i < 3:
+      crash_file = os.path.join(crash_dir, str(crash_i))
+      with open(crash_file, "w") as f:
+        f.write(crash_file)
+      i, crash_i = i + 1, crash_i + 1
+    dir_i += 1
+
+
+class AFLCombineCrashDirsTest(TestCase):
+  def test_combine_crash_directories(self):
+    afl = AFL("deepstate-afl")
+    afl.output_test_dir = os.path.join(os.getcwd(), "tests")
+    crash_out_dir = os.path.join(afl.output_test_dir, "the_fuzzer")
+    os.mkdir(crash_out_dir, 0o777)
+    create_crash_dirs(crash_out_dir)
+    afl.consolidate_crash_dirs()
+    contents = set()
+    for directory in os.listdir(crash_out_dir):
+        crash_dir = os.path.join(crash_out_dir, directory)
+        for f in os.listdir(crash_dir):
+          contents.add(f)
+    shutil.rmtree(crash_out_dir)
+    self.assertEqual(len(expected_contents - contents), 0)