fix tritondse

1 files changed, 65 insertions, 3 deletions

diff --git a/custom_mutators/aflpp_tritondse/aflpp_tritondse.py b/custom_mutators/aflpp_tritondse/aflpp_tritondse.py
index e0219f0b..48367bc7 100644
--- a/custom_mutators/aflpp_tritondse/aflpp_tritondse.py
+++ b/custom_mutators/aflpp_tritondse/aflpp_tritondse.py
@@ -22,14 +22,17 @@ config = None
 dse = None
 cycle = 0
 count = 0
+finding = 0
 hashes = set()
 format = SeedFormat.RAW
 
 def pre_exec_hook(se: SymbolicExecutor, state: ProcessState):
     global count
     global hashes
+    global finding
     if se.seed.hash not in hashes:
         hashes.add(se.seed.hash)
+        finding = 1
         filename = out_path + "/id:" + f"{count:06}" + "," + se.seed.hash
         if not os.path.exists(filename):
             if is_debug:
@@ -47,6 +50,59 @@ def pre_exec_hook(se: SymbolicExecutor, state: ProcessState):
     #        file.write(se.seed.content)
 
 
+#def rtn_open(se: SymbolicExecutor, pstate: ProcessState, pc):
+#    """
+#    The open behavior.
+#    """
+#    logging.debug('open hooked')
+#
+#    # Get arguments
+#    arg0 = pstate.get_argument_value(0)  # const char *pathname
+#    flags = pstate.get_argument_value(1)  # int flags
+#    mode = pstate.get_argument_value(2)  # int mode
+#    arg0s = pstate.memory.read_string(arg0)
+#
+#    # Concretize the whole path name
+#    pstate.concretize_memory_bytes(arg0, len(arg0s)+1)  # Concretize the whole string + \0
+#
+#    # We use flags as concrete value
+#    pstate.concretize_argument(1)
+#
+#    # Use the flags to open the file in the write mode.
+#    mode = ""
+#    if (flags & 0xFF) == 0x00:   # O_RDONLY
+#        mode = "r"
+#    elif (flags & 0xFF) == 0x01: # O_WRONLY
+#        mode = "w"
+#    elif (flags & 0xFF) == 0x02: # O_RDWR
+#        mode = "r+"
+#
+#    if (flags & 0x0100): # O_CREAT
+#        mode += "x"
+#    if (flags & 0x0200): # O_APPEND
+#        mode = "a"  # replace completely value
+#
+#    if se.seed.is_file_defined(arg0s) and "r" in mode:  # input file and opened in reading
+#        logging.info(f"opening an input file: {arg0s}")
+#        # Program is opening an input
+#        data = se.seed.get_file_input(arg0s)
+#        filedesc = pstate.create_file_descriptor(arg0s, io.BytesIO(data))
+#        fd = filedesc.id
+#    else:
+#        # Try to open it as a regular file
+#        try:
+#            fd = open(arg0s, mode)  # use the mode here
+#            filedesc = pstate.create_file_descriptor(arg0s, fd)
+#            fd = filedesc.id
+#        except Exception as e:
+#            logging.debug(f"Failed to open {arg0s} {e}")
+#            fd = pstate.minus_one
+#
+#    pstate.write_register("rax", fd)  # write the return value
+#    pstate.cpu.program_counter = pstate.pop_stack_value()  # pop the return value
+#    se.skip_instruction()  # skip the current instruction so that the engine go straight fetching the next instruction
+
+
 def init(seed):
     global config
     global dse
@@ -115,10 +171,16 @@ def init(seed):
     dse = SymbolicExplorator(config, prog)
     # Add callbacks.
     dse.callback_manager.register_pre_execution_callback(pre_exec_hook)
+    #dse.callback_manager.register_function_callback("open", rtn_open)
 
 
-#def fuzz(buf, add_buf, max_size):
-#    return b""
+def fuzz(buf, add_buf, max_size):
+    global finding
+    finding = 1
+    while finding == 1:
+      finding = 0
+      dse.step()
+    return b""
 
 
 def queue_new_entry(filename_new_queue, filename_orig_queue):
@@ -141,7 +203,7 @@ def queue_new_entry(filename_new_queue, filename_orig_queue):
         dse.add_input_seed(seed)
         # Start exploration!
         #dse.step()
-        dse.explore()
+        #dse.explore()
     pass
 
 def splice_optout():