1 files changed, 106 insertions, 0 deletions
diff --git a/custom_mutators/aflpp_tritondse/aflpp_tritondse.py b/custom_mutators/aflpp_tritondse/aflpp_tritondse.py
new file mode 100644
index 00000000..33bf8a9f
--- /dev/null
+++ b/custom_mutators/aflpp_tritondse/aflpp_tritondse.py
@@ -0,0 +1,106 @@
+import sys
+import os
+import logging
+
+from tritondse import Config
+from tritondse import CoverageStrategy
+from tritondse import ProcessState
+from tritondse import Program
+from tritondse import Seed
+from tritondse import SeedFormat
+from tritondse import SymbolicExecutor
+from tritondse import SymbolicExplorator
+
+
+#logging.basicConfig(level=logging.INFO)
+
+is_debug = False
+out_path = "out/tritondse/queue"
+input_file = None
+prog = None
+config = None
+dse = None
+cycle = 0
+count = 0
+hashes = set()
+
+def pre_exec_hook(se: SymbolicExecutor, state: ProcessState):
+    #logging.info(f"[PRE-EXEC] Processing seed: {se.seed.hash}, \
+    #                ({repr(se.seed.content)})")
+    global count
+    global hasshes
+    if se.seed.hash not in hashes:
+        hashes.add(se.seed.hash)
+        filename = out_path + "/id:" + f"{count:06}" + "," + se.seed.hash
+        if not os.path.exists(filename):
+            with open(filename, 'wb') as file:
+                file.write(se.seed.content)
+                count += 1
+    if input_file:
+        with open(input_file, 'wb') as file:
+            file.write(se.seed.content)
+
+
+def init(seed):
+    global prog
+    global config
+    global dse
+    global input_file
+    global is_debug
+    # Load the program (LIEF-based program loader).
+    prog = Program(os.environ['TRITON_DSE_TARGET'])
+    # Set the configuration.
+    argv = None
+    try:
+        foo = os.environ['AFL_DEBUG']
+        is_debug = True
+    except KeyError:
+        pass
+    try:
+        argv_list = os.environ['TRITON_DSE_TARGET_ARGV']
+        argv = argv_list.split()
+    except KeyError:
+        pass
+    try:
+        foo = os.environ['TRITON_DSE_TARGET_INPUT']
+        input_file = foo
+    except KeyError:
+        pass
+    config = Config(coverage_strategy = CoverageStrategy.PATH,
+                    debug = is_debug,
+                    pipe_stdout = is_debug,
+                    pipe_stderr = is_debug,
+                    execution_timeout = 1,
+                    program_argv = argv,
+                    smt_timeout= 50,
+                    seed_format = SeedFormat.RAW)
+    # Create an instance of the Symbolic Explorator
+    dse = SymbolicExplorator(config, prog)
+    # Add callbacks.
+    dse.callback_manager.register_pre_execution_callback(pre_exec_hook)
+    # Create the output directory
+    os.makedirs(out_path, exist_ok=True)
+
+
+#def fuzz(buf, add_buf, max_size):
+#    return b""
+
+
+def queue_new_entry(filename_new_queue, filename_orig_queue):
+    global dse
+    global cycle
+    # Add seed to the worklist.
+    with open(filename_new_queue, "rb") as file:
+        seed = file.read()
+    seed = Seed(seed)
+    dse.add_input_seed(seed)
+    if is_debug:
+        print("NEW FILE " + filename_new_queue + " count " + str(cycle))
+        cycle += 1
+    # Start exploration!
+    #dse.step()
+    dse.explore()
+    pass
+
+def splice_optout():
+    pass