Uniform API for both Python and custom mutator

7 files changed, 322 insertions, 271 deletions

diff --git a/include/afl-fuzz.h b/include/afl-fuzz.h
index b8b6d5e3..e6aaee68 100644
--- a/include/afl-fuzz.h
+++ b/include/afl-fuzz.h
@@ -472,12 +472,12 @@ struct custom_mutator {
    *
    * (Optional)
    */
-  u32 (*afl_custom_init)(void);
+  void (*afl_custom_init)(unsigned int seed);
 
   /**
    * Perform custom mutations on a given input
    *
-   * (Required)
+   * (Optional for now. Required in the future)
    *
    * @param[in] data Input data to be mutated
    * @param[in] size Size of input data
@@ -498,46 +498,67 @@ struct custom_mutator {
    * (Optional) If this functionality is not needed, simply don't define this
    * function.
    *
-   * @param[in] data Buffer containing the test case to be executed.
-   * @param[in] size Size of the test case.
+   * @param[in] data Buffer containing the test case to be executed
+   * @param[in] size Size of the test case
    * @param[out] new_data Buffer to store the test case after processing
-   * @return Size of data after processing.
+   * @return Size of data after processing
    */
   size_t (*afl_custom_pre_save)(u8* data, size_t size, u8** new_data);
 
   /**
-   * TODO: figure out what `trim` is
+   * This method is called at the start of each trimming operation and receives
+   * the initial buffer. It should return the amount of iteration steps possible
+   * on this input (e.g. if your input has n elements and you want to remove
+   * them one by one, return n, if you do a binary search, return log(n),
+   * and so on...).
+   *
+   * If your trimming algorithm doesn't allow you to determine the amount of
+   * (remaining) steps easily (esp. while running), then you can alternatively
+   * return 1 here and always return 0 in post_trim until you are finished and
+   * no steps remain. In that case, returning 1 in post_trim will end the
+   * trimming routine. The whole current index/max iterations stuff is only used
+   * to show progress.
    *
    * (Optional)
+   *
+   * @param data Buffer containing the test case
+   * @param size Size of the test case
+   * @return The amount of possible iteration steps to trim the input
    */
-  u32 (*afl_custom_init_trim)(u8*, size_t);
+  u32 (*afl_custom_init_trim)(u8* data, size_t size);
 
   /**
-   * TODO: figure out how `trim` works
+   * This method is called for each trimming operation. It doesn't have any
+   * arguments because we already have the initial buffer from init_trim and we
+   * can memorize the current state in global variables. This can also save
+   * reparsing steps for each iteration. It should return the trimmed input
+   * buffer, where the returned data must not exceed the initial input data in
+   * length. Returning anything that is larger than the original data (passed
+   * to init_trim) will result in a fatal abort of AFLFuzz.
    *
    * (Optional)
    *
-   * @param[out] ret (TODO: finish here)
-   * @param[out] ret_len (TODO: finish here)
+   * @param[out] ret Buffer containing the trimmed test case
+   * @param[out] ret_len Size of the trimmed test case
    */
   void (*afl_custom_trim)(u8** ret, size_t* ret_len);
 
   /**
-   * A post-processing function for the last trim operation.
+   * This method is called after each trim operation to inform you if your
+   * trimming step was successful or not (in terms of coverage). If you receive
+   * a failure here, you should reset your input to the last known good state.
    *
    * (Optional)
    *
    * @param success Indicates if the last trim operation was successful.
+   * @return The next trim iteration index (from 0 to the maximum amount of
+   *     steps returned in init_trim)
    */
   u32 (*afl_custom_post_trim)(u8 success);
 };
 
 extern struct custom_mutator* mutator;
 
-size_t (*custom_mutator)(u8* data, size_t size, u8* mutated_out,
-                         size_t max_size, unsigned int seed);
-size_t (*pre_save_handler)(u8* data, size_t size, u8** new_data);
-
 /* Interesting values, as per config.h */
 
 extern s8  interesting_8[INTERESTING_8_LEN];
@@ -598,17 +619,19 @@ void setup_custom_mutator(void);
 void destroy_custom_mutator(void);
 void load_custom_mutator(const char*);
 void load_custom_mutator_py(const char*);
+u8   trim_case_custom(char** argv, struct queue_entry* q, u8* in_buf);
 
 /* Python */
 #ifdef USE_PYTHON
-int  init_py();
-void finalize_py();
-void fuzz_py(char*, size_t, char*, size_t, char**, size_t*);
+int    init_py_module(u8*);
+void   finalize_py_module();
+
+void   init_py(unsigned int seed);
+void   fuzz_py(char*, size_t, char*, size_t, char**, size_t*);
 size_t pre_save_py(u8* data, size_t size, u8** new_data);
-u32  init_trim_py(char*, size_t);
-u32  post_trim_py(char);
-void trim_py(char**, size_t*);
-u8   trim_case_python(char**, struct queue_entry*, u8*);
+u32    init_trim_py(u8*, size_t);
+u32    post_trim_py(u8);
+void   trim_py(u8**, size_t*);
 #endif
 
 /* Queue */
diff --git a/src/afl-fuzz-globals.c b/src/afl-fuzz-globals.c
index 1fd4b26d..a054499f 100644
--- a/src/afl-fuzz-globals.c
+++ b/src/afl-fuzz-globals.c
@@ -259,11 +259,6 @@ s32 cmplog_child_pid, cmplog_forksrv_pid;
 /* Custom mutator */
 struct custom_mutator* mutator;
 
-/* hooks for the custom mutator function */
-size_t (*custom_mutator)(u8 *data, size_t size, u8 *mutated_out,
-                         size_t max_size, unsigned int seed);
-size_t (*pre_save_handler)(u8 *data, size_t size, u8 **new_data);
-
 /* Interesting values, as per config.h */
 
 s8  interesting_8[] = {INTERESTING_8};
diff --git a/src/afl-fuzz-mutators.c b/src/afl-fuzz-mutators.c
index 0e77a690..22ca1384 100644
--- a/src/afl-fuzz-mutators.c
+++ b/src/afl-fuzz-mutators.c
@@ -42,11 +42,22 @@ void setup_custom_mutator(void) {
   }
 
 #ifdef USE_PYTHON
-  if (init_py()) FATAL("Failed to initialize Python module");
+  u8* module_name = getenv("AFL_PYTHON_MODULE");
 
-  // u8* module_name = getenv("AFL_PYTHON_MODULE");
-  // if (py_module && module_name)
-  //   load_custom_mutator_py(module_name);
+  if (module_name) {
+
+    if (limit_time_sig)
+      FATAL(
+          "MOpt and Python mutator are mutually exclusive. We accept pull "
+          "requests that integrates MOpt with the optional mutators "
+          "(custom/radamsa/redquenn/...).");
+
+    if (init_py_module(module_name))
+      FATAL("Failed to initialize Python module");
+
+    load_custom_mutator_py(module_name);
+
+  }
 #else
   if (getenv("AFL_PYTHON_MODULE"))
     FATAL("Your AFL binary was built without Python support");
@@ -62,7 +73,7 @@ void destroy_custom_mutator(void) {
     else {
       /* Python mutator */
 #ifdef USE_PYTHON
-      finalize_py();
+      finalize_py_module();
 #endif
     }
 
@@ -80,11 +91,11 @@ void load_custom_mutator(const char* fn) {
   ACTF("Loading custom mutator library from '%s'...", fn);
 
   dh = dlopen(fn, RTLD_NOW);
-  if (!mutator->dh) FATAL("%s", dlerror());
+  if (!dh) FATAL("%s", dlerror());
   mutator->dh = dh;
 
   /* Mutator */
-  /* "afl_custom_init", optional */
+  /* "afl_custom_init", optional for backward compatibility */
   mutator->afl_custom_init = dlsym(dh, "afl_custom_init");
   if (!mutator->afl_custom_init)
     WARNF("Symbol 'afl_custom_init' not found.");
@@ -92,13 +103,14 @@ void load_custom_mutator(const char* fn) {
   /* "afl_custom_fuzz" or "afl_custom_mutator", required */
   mutator->afl_custom_fuzz = dlsym(dh, "afl_custom_fuzz");
   if (!mutator->afl_custom_fuzz) {
+
     /* Try "afl_custom_mutator" for backward compatibility */
     WARNF("Symbol 'afl_custom_fuzz' not found. Try 'afl_custom_mutator'.");
 
     mutator->afl_custom_fuzz = dlsym(dh, "afl_custom_mutator");
-    if (!mutator->afl_custom_fuzz) {
+    if (!mutator->afl_custom_fuzz)
       FATAL("Symbol 'afl_custom_mutator' not found.");
-    }
+
   }
 
   /* "afl_custom_pre_save", optional */
@@ -106,6 +118,7 @@ void load_custom_mutator(const char* fn) {
   if (!mutator->afl_custom_pre_save)
     WARNF("Symbol 'afl_custom_pre_save' not found.");
 
+  u8 notrim = 0;
   /* "afl_custom_init_trim", optional */
   mutator->afl_custom_init_trim = dlsym(dh, "afl_custom_init_trim");
   if (!mutator->afl_custom_init_trim)
@@ -121,29 +134,177 @@ void load_custom_mutator(const char* fn) {
   if (!mutator->afl_custom_post_trim)
     WARNF("Symbol 'afl_custom_post_trim' not found.");
 
+  if (notrim) {
+
+    mutator->afl_custom_init_trim = NULL;
+    mutator->afl_custom_trim = NULL;
+    mutator->afl_custom_post_trim = NULL;
+    WARNF(
+        "Custom mutator does not implement all three trim APIs, standard "
+        "trimming will be used.");
+
+  }
+
   OKF("Custom mutator '%s' installed successfully.", fn);
 
   /* Initialize the custom mutator */
   if (mutator->afl_custom_init)
-    mutator->afl_custom_init();
+    mutator->afl_custom_init(UR(0xFFFFFFFF));
+
+}
+
+u8 trim_case_custom(char** argv, struct queue_entry* q, u8* in_buf) {
+
+  static u8 tmp[64];
+  static u8 clean_trace[MAP_SIZE];
+
+  u8  needs_write = 0, fault = 0;
+  u32 trim_exec = 0;
+  u32 orig_len = q->len;
+
+  stage_name = tmp;
+  bytes_trim_in += q->len;
+
+  /* Initialize trimming in the custom mutator */
+  stage_cur = 0;
+  stage_max = mutator->afl_custom_init_trim(in_buf, q->len);
+
+  if (not_on_tty && debug)
+    SAYF("[Custom Trimming] START: Max %d iterations, %u bytes", stage_max,
+         q->len);
+
+  while (stage_cur < stage_max) {
+
+    sprintf(tmp, "ptrim %s", DI(trim_exec));
+
+    u32 cksum;
+
+    u8*    retbuf = NULL;
+    size_t retlen = 0;
+
+    mutator->afl_custom_trim(&retbuf, &retlen);
+
+    if (retlen > orig_len)
+      FATAL(
+          "Trimmed data returned by custom mutator is larger than original "
+          "data");
+
+    write_to_testcase(retbuf, retlen);
+
+    fault = run_target(argv, exec_tmout);
+    ++trim_execs;
+
+    if (stop_soon || fault == FAULT_ERROR) {
+
+      free(retbuf);
+      goto abort_trimming;
+
+    }
+
+    cksum = hash32(trace_bits, MAP_SIZE, HASH_CONST);
+
+    if (cksum == q->exec_cksum) {
+
+      q->len = retlen;
+      memcpy(in_buf, retbuf, retlen);
+
+      /* Let's save a clean trace, which will be needed by
+         update_bitmap_score once we're done with the trimming stuff. */
+
+      if (!needs_write) {
+
+        needs_write = 1;
+        memcpy(clean_trace, trace_bits, MAP_SIZE);
+
+      }
+
+      /* Tell the custom mutator that the trimming was successful */
+      stage_cur = mutator->afl_custom_post_trim(1);
+
+      if (not_on_tty && debug)
+        SAYF("[Custom Trimming] SUCCESS: %d/%d iterations (now at %u bytes)",
+             stage_cur, stage_max, q->len);
+
+    } else {
+
+      /* Tell the custom mutator that the trimming was unsuccessful */
+      stage_cur = mutator->afl_custom_post_trim(0);
+      if (not_on_tty && debug)
+        SAYF("[Custom Trimming] FAILURE: %d/%d iterations", stage_cur,
+             stage_max);
+
+    }
+
+    free(retbuf);
+
+    /* Since this can be slow, update the screen every now and then. */
+
+    if (!(trim_exec++ % stats_update_freq)) show_stats();
+
+  }
+
+  if (not_on_tty && debug)
+    SAYF("[Custom Trimming] DONE: %u bytes -> %u bytes", orig_len, q->len);
+
+  /* If we have made changes to in_buf, we also need to update the on-disk
+     version of the test case. */
+
+  if (needs_write) {
+
+    s32 fd;
+
+    unlink(q->fname);                                      /* ignore errors */
+
+    fd = open(q->fname, O_WRONLY | O_CREAT | O_EXCL, 0600);
+
+    if (fd < 0) PFATAL("Unable to create '%s'", q->fname);
+
+    ck_write(fd, in_buf, q->len, q->fname);
+    close(fd);
+
+    memcpy(trace_bits, clean_trace, MAP_SIZE);
+    update_bitmap_score(q);
+
+  }
+
+abort_trimming:
+
+  bytes_trim_out += q->len;
+  return fault;
 
 }
 
-// void load_custom_mutator_py(const char* module_name) {
+void load_custom_mutator_py(const char* module_name) {
 
-//   mutator = ck_alloc(sizeof(struct custom_mutator));
+  mutator = ck_alloc(sizeof(struct custom_mutator));
 
-//   mutator->name = module_name;
-//   ACTF("Loading Python mutator library from '%s'...", module_name);
+  mutator->name = module_name;
+  ACTF("Loading Python mutator library from '%s'...", module_name);
 
-//   /* Initialize of the Python mutator has been invoked in "init_py()" */
-//   mutator->afl_custom_init = NULL;
-//   mutator->afl_custom_fuzz = fuzz_py;
-//   mutator->afl_custom_pre_save = pre_save_py;
-//   mutator->afl_custom_init_trim = init_trim_py;
-//   mutator->afl_custom_trim = trim_py;
-//   mutator->afl_custom_post_trim = post_trim_py;
+  /* TODO: unify "init" and "fuzz" */
+  if (py_functions[PY_FUNC_INIT])
+    mutator->afl_custom_init = init_py;
 
-//   OKF("Python mutator '%s' installed successfully.", module_name);
+  /* "afl_custom_fuzz" should not be NULL, but the interface of Python mutator
+     is quite different from the custom mutator. */
+  mutator->afl_custom_fuzz = NULL;
 
-// }
+  if (py_functions[PY_FUNC_PRE_SAVE])
+    mutator->afl_custom_pre_save = pre_save_py;
+
+  if (py_functions[PY_FUNC_INIT_TRIM])
+    mutator->afl_custom_init_trim = init_trim_py;
+
+  if (py_functions[PY_FUNC_POST_TRIM])
+    mutator->afl_custom_post_trim = post_trim_py;
+
+  if (py_functions[PY_FUNC_TRIM])
+    mutator->afl_custom_trim = trim_py;
+
+  OKF("Python mutator '%s' installed successfully.", module_name);
+
+  /* Initialize the custom mutator */
+  if (mutator->afl_custom_init)
+    mutator->afl_custom_init(UR(0xFFFFFFFF));
+
+}
diff --git a/src/afl-fuzz-one.c b/src/afl-fuzz-one.c
index f1efe2df..24ea76f3 100644
--- a/src/afl-fuzz-one.c
+++ b/src/afl-fuzz-one.c
@@ -449,7 +449,7 @@ u8 fuzz_one_original(char** argv) {
    * TRIMMING *
    ************/
 
-  if (!dumb_mode && !queue_cur->trim_done && !custom_mutator && !disable_trim) {
+  if (!dumb_mode && !queue_cur->trim_done && !disable_trim) {
 
     u8 res = trim_case(argv, queue_cur, in_buf);
 
@@ -484,7 +484,7 @@ u8 fuzz_one_original(char** argv) {
 
   // custom_stage:	// not used - yet
 
-  if (custom_mutator) {
+  if (mutator->afl_custom_fuzz) {
 
     stage_short = "custom";
     stage_name = "custom mutator";
@@ -499,8 +499,9 @@ u8 fuzz_one_original(char** argv) {
     for (stage_cur = 0; stage_cur < stage_max; ++stage_cur) {
 
       size_t orig_size = (size_t)len;
-      size_t mutated_size = custom_mutator(in_buf, orig_size, mutated_buf,
-                                           max_seed_size, UR(UINT32_MAX));
+      size_t mutated_size = mutator->afl_custom_fuzz(in_buf, orig_size,
+                                                     mutated_buf, max_seed_size,
+                                                     UR(UINT32_MAX));
       if (mutated_size > 0) {
 
         out_buf = ck_realloc(out_buf, mutated_size);
diff --git a/src/afl-fuzz-python.c b/src/afl-fuzz-python.c
index be86ebba..da478cc2 100644
--- a/src/afl-fuzz-python.c
+++ b/src/afl-fuzz-python.c
@@ -28,116 +28,73 @@
 /* Python stuff */
 #ifdef USE_PYTHON
 
-int init_py() {
+int init_py_module(u8* module_name) {
 
-  Py_Initialize();
-  u8* module_name = getenv("AFL_PYTHON_MODULE");
-
-  if (module_name) {
-
-    if (limit_time_sig)
-      FATAL(
-          "MOpt and Python mutator are mutually exclusive. We accept pull "
-          "requests that integrates MOpt with the optional mutators "
-          "(custom/radamsa/redquenn/...).");
+  if (!module_name) return 1;
 
 #if PY_MAJOR_VERSION >= 3
-    PyObject* py_name = PyUnicode_FromString(module_name);
+  PyObject* py_name = PyUnicode_FromString(module_name);
 #else
-    PyObject* py_name = PyString_FromString(module_name);
+  PyObject* py_name = PyString_FromString(module_name);
 #endif
 
-    py_module = PyImport_Import(py_name);
-    Py_DECREF(py_name);
+  py_module = PyImport_Import(py_name);
+  Py_DECREF(py_name);
 
-    if (py_module != NULL) {
-
-      u8 py_notrim = 0, py_idx;
-      py_functions[PY_FUNC_INIT] = PyObject_GetAttrString(py_module, "init");
-      py_functions[PY_FUNC_FUZZ] = PyObject_GetAttrString(py_module, "fuzz");
-      py_functions[PY_FUNC_PRE_SAVE] =
-          PyObject_GetAttrString(py_module, "pre_save");
-      py_functions[PY_FUNC_INIT_TRIM] =
-          PyObject_GetAttrString(py_module, "init_trim");
-      py_functions[PY_FUNC_POST_TRIM] =
-          PyObject_GetAttrString(py_module, "post_trim");
-      py_functions[PY_FUNC_TRIM] = PyObject_GetAttrString(py_module, "trim");
+  if (py_module != NULL) {
 
-      for (py_idx = 0; py_idx < PY_FUNC_COUNT; ++py_idx) {
+    u8 py_notrim = 0, py_idx;
+    py_functions[PY_FUNC_INIT] = PyObject_GetAttrString(py_module, "init");
+    py_functions[PY_FUNC_FUZZ] = PyObject_GetAttrString(py_module, "fuzz");
+    py_functions[PY_FUNC_PRE_SAVE] =
+        PyObject_GetAttrString(py_module, "pre_save");
+    py_functions[PY_FUNC_INIT_TRIM] =
+        PyObject_GetAttrString(py_module, "init_trim");
+    py_functions[PY_FUNC_POST_TRIM] =
+        PyObject_GetAttrString(py_module, "post_trim");
+    py_functions[PY_FUNC_TRIM] = PyObject_GetAttrString(py_module, "trim");
 
-        if (!py_functions[py_idx] || !PyCallable_Check(py_functions[py_idx])) {
+    for (py_idx = 0; py_idx < PY_FUNC_COUNT; ++py_idx) {
 
-          if (py_idx >= PY_FUNC_INIT_TRIM && py_idx <= PY_FUNC_TRIM) {
+      if (!py_functions[py_idx] || !PyCallable_Check(py_functions[py_idx])) {
 
-            // Implementing the trim API is optional for now
-            if (PyErr_Occurred()) PyErr_Print();
-            py_notrim = 1;
+        if (py_idx >= PY_FUNC_INIT_TRIM && py_idx <= PY_FUNC_TRIM) {
 
-          } else {
+          // Implementing the trim API is optional for now
+          if (PyErr_Occurred()) PyErr_Print();
+          py_notrim = 1;
 
-            if (PyErr_Occurred()) PyErr_Print();
-            fprintf(stderr,
-                    "Cannot find/call function with index %d in external "
-                    "Python module.\n",
-                    py_idx);
-            return 1;
+        } else {
 
-          }
+          if (PyErr_Occurred()) PyErr_Print();
+          fprintf(stderr,
+                  "Cannot find/call function with index %d in external "
+                  "Python module.\n",
+                  py_idx);
+          return 1;
 
         }
 
       }
 
-      if (py_notrim) {
-
-        py_functions[PY_FUNC_INIT_TRIM] = NULL;
-        py_functions[PY_FUNC_POST_TRIM] = NULL;
-        py_functions[PY_FUNC_TRIM] = NULL;
-        WARNF(
-            "Python module does not implement trim API, standard trimming will "
-            "be used.");
-
-      }
-
-      PyObject *py_args, *py_value;
-
-      /* Provide the init function a seed for the Python RNG */
-      py_args = PyTuple_New(1);
-#if PY_MAJOR_VERSION >= 3
-      py_value = PyLong_FromLong(UR(0xFFFFFFFF));
-#else
-      py_value = PyInt_FromLong(UR(0xFFFFFFFF));
-#endif
-
-      if (!py_value) {
-
-        Py_DECREF(py_args);
-        fprintf(stderr, "Cannot convert argument\n");
-        return 1;
-
-      }
-
-      PyTuple_SetItem(py_args, 0, py_value);
-
-      py_value = PyObject_CallObject(py_functions[PY_FUNC_INIT], py_args);
-
-      Py_DECREF(py_args);
-
-      if (py_value == NULL) {
+    }
 
-        PyErr_Print();
-        fprintf(stderr, "Call failed\n");
-        return 1;
+    if (py_notrim) {
 
-      }
+      py_functions[PY_FUNC_INIT_TRIM] = NULL;
+      py_functions[PY_FUNC_POST_TRIM] = NULL;
+      py_functions[PY_FUNC_TRIM] = NULL;
+      WARNF(
+          "Python module does not implement trim API, standard trimming will "
+          "be used.");
 
-    } else {
+    }
 
-      PyErr_Print();
-      fprintf(stderr, "Failed to load \"%s\"\n", module_name);
-      return 1;
+  } else {
 
-    }
+    PyErr_Print();
+    fprintf(stderr, "Failed to load \"%s\"\n", module_name);
+    return 1;
 
   }
 
@@ -145,7 +102,7 @@ int init_py() {
 
 }
 
-void finalize_py() {
+void finalize_py_module() {
 
   if (py_module != NULL) {
 
@@ -213,7 +170,43 @@ void fuzz_py(char* buf, size_t buflen, char* add_buf, size_t add_buflen,
 
 }
 
-u32 init_trim_py(char* buf, size_t buflen) {
+size_t pre_save_py(u8* data, size_t size, u8** new_data) {
+
+  size_t new_size;
+  PyObject *py_args, *py_value;
+  py_args = PyTuple_New(2);
+  py_value = PyByteArray_FromStringAndSize(data, size);
+  if (!py_value) {
+
+    Py_DECREF(py_args);
+    FATAL("Failed to convert arguments");
+
+  }
+
+  PyTuple_SetItem(py_args, 0, py_value);
+
+  py_value = PyObject_CallObject(py_functions[PY_FUNC_PRE_SAVE], py_args);
+
+  Py_DECREF(py_args);
+
+  if (py_value != NULL) {
+
+    new_size = PyByteArray_Size(py_value);
+    *new_data = malloc(new_size);
+    memcpy(*new_data, PyByteArray_AsString(py_value), new_size);
+    Py_DECREF(py_value);
+    return new_size;
+
+  } else {
+
+    PyErr_Print();
+    FATAL("Call failed");
+
+  }
+
+}
+
+u32 init_trim_py(u8* buf, size_t buflen) {
 
   PyObject *py_args, *py_value;
 
@@ -250,7 +243,7 @@ u32 init_trim_py(char* buf, size_t buflen) {
 
 }
 
-u32 post_trim_py(char success) {
+u32 post_trim_py(u8 success) {
 
   PyObject *py_args, *py_value;
 
@@ -288,7 +281,7 @@ u32 post_trim_py(char success) {
 
 }
 
-void trim_py(char** ret, size_t* retlen) {
+void trim_py(u8** ret, size_t* retlen) {
 
   PyObject *py_args, *py_value;
 
@@ -312,126 +305,5 @@ void trim_py(char** ret, size_t* retlen) {
 
 }
 
-u8 trim_case_python(char** argv, struct queue_entry* q, u8* in_buf) {
-
-  static u8 tmp[64];
-  static u8 clean_trace[MAP_SIZE];
-
-  u8  needs_write = 0, fault = 0;
-  u32 trim_exec = 0;
-  u32 orig_len = q->len;
-
-  stage_name = tmp;
-  bytes_trim_in += q->len;
-
-  /* Initialize trimming in the Python module */
-  stage_cur = 0;
-  stage_max = init_trim_py(in_buf, q->len);
-
-  if (not_on_tty && debug)
-    SAYF("[Python Trimming] START: Max %d iterations, %u bytes", stage_max,
-         q->len);
-
-  while (stage_cur < stage_max) {
-
-    sprintf(tmp, "ptrim %s", DI(trim_exec));
-
-    u32 cksum;
-
-    char*  retbuf = NULL;
-    size_t retlen = 0;
-
-    trim_py(&retbuf, &retlen);
-
-    if (retlen > orig_len)
-      FATAL(
-          "Trimmed data returned by Python module is larger than original "
-          "data");
-
-    write_to_testcase(retbuf, retlen);
-
-    fault = run_target(argv, exec_tmout);
-    ++trim_execs;
-
-    if (stop_soon || fault == FAULT_ERROR) {
-
-      free(retbuf);
-      goto abort_trimming;
-
-    }
-
-    cksum = hash32(trace_bits, MAP_SIZE, HASH_CONST);
-
-    if (cksum == q->exec_cksum) {
-
-      q->len = retlen;
-      memcpy(in_buf, retbuf, retlen);
-
-      /* Let's save a clean trace, which will be needed by
-         update_bitmap_score once we're done with the trimming stuff. */
-
-      if (!needs_write) {
-
-        needs_write = 1;
-        memcpy(clean_trace, trace_bits, MAP_SIZE);
-
-      }
-
-      /* Tell the Python module that the trimming was successful */
-      stage_cur = post_trim_py(1);
-
-      if (not_on_tty && debug)
-        SAYF("[Python Trimming] SUCCESS: %d/%d iterations (now at %u bytes)",
-             stage_cur, stage_max, q->len);
-
-    } else {
-
-      /* Tell the Python module that the trimming was unsuccessful */
-      stage_cur = post_trim_py(0);
-      if (not_on_tty && debug)
-        SAYF("[Python Trimming] FAILURE: %d/%d iterations", stage_cur,
-             stage_max);
-
-    }
-
-    free(retbuf);
-
-    /* Since this can be slow, update the screen every now and then. */
-
-    if (!(trim_exec++ % stats_update_freq)) show_stats();
-
-  }
-
-  if (not_on_tty && debug)
-    SAYF("[Python Trimming] DONE: %u bytes -> %u bytes", orig_len, q->len);
-
-  /* If we have made changes to in_buf, we also need to update the on-disk
-     version of the test case. */
-
-  if (needs_write) {
-
-    s32 fd;
-
-    unlink(q->fname);                                      /* ignore errors */
-
-    fd = open(q->fname, O_WRONLY | O_CREAT | O_EXCL, 0600);
-
-    if (fd < 0) PFATAL("Unable to create '%s'", q->fname);
-
-    ck_write(fd, in_buf, q->len, q->fname);
-    close(fd);
-
-    memcpy(trace_bits, clean_trace, MAP_SIZE);
-    update_bitmap_score(q);
-
-  }
-
-abort_trimming:
-
-  bytes_trim_out += q->len;
-  return fault;
-
-}
-
 #endif                                                        /* USE_PYTHON */
 
diff --git a/src/afl-fuzz-run.c b/src/afl-fuzz-run.c
index 95c2c5d4..a1d10387 100644
--- a/src/afl-fuzz-run.c
+++ b/src/afl-fuzz-run.c
@@ -309,10 +309,10 @@ void write_to_testcase(void* mem, u32 len) {
 
     lseek(fd, 0, SEEK_SET);
 
-  if (pre_save_handler) {
+  if (mutator->afl_custom_pre_save) {
 
     u8*    new_data;
-    size_t new_size = pre_save_handler(mem, len, &new_data);
+    size_t new_size = mutator->afl_custom_pre_save(mem, len, &new_data);
     ck_write(fd, new_data, new_size, out_file);
 
   } else {
@@ -678,9 +678,8 @@ void sync_fuzzers(char** argv) {
 
 u8 trim_case(char** argv, struct queue_entry* q, u8* in_buf) {
 
-#ifdef USE_PYTHON
-  if (py_functions[PY_FUNC_TRIM]) return trim_case_python(argv, q, in_buf);
-#endif
+  /* Custom mutator trimmer */
+  if (mutator->afl_custom_trim) return trim_case_custom(argv, q, in_buf);
 
   static u8 tmp[64];
   static u8 clean_trace[MAP_SIZE];
diff --git a/src/afl-fuzz-stats.c b/src/afl-fuzz-stats.c
index c1aa8315..1b763c01 100644
--- a/src/afl-fuzz-stats.c
+++ b/src/afl-fuzz-stats.c
@@ -655,7 +655,7 @@ void show_stats(void) {
 
   }
 
-  if (custom_mutator) {
+  if (mutator) {
 
     sprintf(tmp, "%s/%s", DI(stage_finds[STAGE_CUSTOM_MUTATOR]),
             DI(stage_cycles[STAGE_CUSTOM_MUTATOR]));