4 files changed, 82 insertions, 8 deletions

diff --git a/src/afl-fuzz-mutators.c b/src/afl-fuzz-mutators.c
index 22ca1384..9365d487 100644
--- a/src/afl-fuzz-mutators.c
+++ b/src/afl-fuzz-mutators.c
@@ -281,7 +281,6 @@ void load_custom_mutator_py(const char* module_name) {
   mutator->name = module_name;
   ACTF("Loading Python mutator library from '%s'...", module_name);
 
-  /* TODO: unify "init" and "fuzz" */
   if (py_functions[PY_FUNC_INIT])
     mutator->afl_custom_init = init_py;
 
diff --git a/src/afl-fuzz-one.c b/src/afl-fuzz-one.c
index 24ea76f3..1e6dd45d 100644
--- a/src/afl-fuzz-one.c
+++ b/src/afl-fuzz-one.c
@@ -1647,7 +1647,7 @@ python_stage:
     ck_read(fd, new_buf, target->len, target->fname);
     close(fd);
 
-    fuzz_py(out_buf, len, new_buf, target->len, &retbuf, &retlen);
+    fuzz_py_original(out_buf, len, new_buf, target->len, &retbuf, &retlen);
 
     ck_free(new_buf);
 
diff --git a/src/afl-fuzz-python.c b/src/afl-fuzz-python.c
index 30156fa6..c8caa4c1 100644
--- a/src/afl-fuzz-python.c
+++ b/src/afl-fuzz-python.c
@@ -32,6 +32,8 @@ int init_py_module(u8* module_name) {
 
   if (!module_name) return 1;
 
+  Py_Initialize();
+
 #if PY_MAJOR_VERSION >= 3
   PyObject* py_name = PyUnicode_FromString(module_name);
 #else
@@ -58,7 +60,12 @@ int init_py_module(u8* module_name) {
 
       if (!py_functions[py_idx] || !PyCallable_Check(py_functions[py_idx])) {
 
-        if (py_idx >= PY_FUNC_INIT_TRIM && py_idx <= PY_FUNC_TRIM) {
+        if (py_idx == PY_FUNC_PRE_SAVE) {
+
+          // Implenting the pre_save API is optional for now
+          if (PyErr_Occurred()) PyErr_Print();
+
+        } else if (py_idx >= PY_FUNC_INIT_TRIM && py_idx <= PY_FUNC_TRIM) {
 
           // Implementing the trim API is optional for now
           if (PyErr_Occurred()) PyErr_Print();
@@ -152,8 +159,9 @@ void init_py(unsigned int seed) {
   }
 }
 
-void fuzz_py(char* buf, size_t buflen, char* add_buf, size_t add_buflen,
-             char** ret, size_t* retlen) {
+void fuzz_py_original(char* buf, size_t buflen,
+                      char* add_buf, size_t add_buflen,
+                      char** ret, size_t* retlen) {
 
   if (py_module != NULL) {
 
@@ -204,6 +212,72 @@ void fuzz_py(char* buf, size_t buflen, char* add_buf, size_t add_buflen,
 
 }
 
+size_t fuzz_py(u8* data, size_t size, u8* mutated_out, size_t max_size,
+               unsigned int seed) {
+
+  size_t out_size;
+  PyObject *py_args, *py_value;
+  py_args = PyTuple_New(3);
+
+  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);
+
+#if PY_MAJOR_VERSION >= 3
+  py_value = PyLong_FromLong(max_size);
+#else
+  py_value = PyInt_FromLong(max_size);
+#endif
+  if (!py_value) {
+
+    Py_DECREF(py_args);
+    FATAL("Failed to convert arguments");
+
+  }
+
+  PyTuple_SetItem(py_args, 1, py_value);
+
+#if PY_MAJOR_VERSION >= 3
+  py_value = PyLong_FromLong(seed);
+#else
+  py_value = PyInt_FromLong(seed);
+#endif
+  if (!py_value) {
+
+    Py_DECREF(py_args);
+    FATAL("Failed to convert arguments");
+
+  }
+
+  PyTuple_SetItem(py_args, 2, py_value);
+
+  py_value = PyObject_CallObject(py_functions[PY_FUNC_FUZZ], py_args);
+
+  Py_DECREF(py_args);
+
+  if (py_value != NULL) {
+
+    out_size = PyByteArray_Size(py_value);
+    memcpy(mutated_out, PyByteArray_AsString(py_value), out_size);
+    Py_DECREF(py_value);
+
+    return out_size;
+
+  } else {
+
+    PyErr_Print();
+    FATAL("Call failed");
+
+  }
+
+}
+
 size_t pre_save_py(u8* data, size_t size, u8** new_data) {
 
   size_t new_size;
diff --git a/src/afl-fuzz.c b/src/afl-fuzz.c
index 2f0043ab..d329a20e 100644
--- a/src/afl-fuzz.c
+++ b/src/afl-fuzz.c
@@ -752,9 +752,9 @@ int main(int argc, char** argv, char** envp) {
   if (get_afl_env("AFL_FAST_CAL")) fast_cal = 1;
 
   if (get_afl_env("AFL_AUTORESUME")) {
-    
+
     autoresume = 1;
-    if (in_place_resume) 
+    if (in_place_resume)
       SAYF("AFL_AUTORESUME has no effect for '-i -'");
 
   }
@@ -862,7 +862,6 @@ int main(int argc, char** argv, char** envp) {
   check_cpu_governor();
 
   setup_post();
-  setup_custom_mutator();
   setup_shm(dumb_mode);
 
   if (!in_bitmap) memset(virgin_bits, 255, MAP_SIZE);
@@ -873,6 +872,8 @@ int main(int argc, char** argv, char** envp) {
 
   setup_dirs_fds();
 
+  setup_custom_mutator();
+
   setup_cmdline_file(argv + optind);
 
   read_testcases();