Refactoring fuzz_py API

include/afl-fuzz.h

@@ -627,7 +627,9 @@ 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*);
+/* TODO: unify fuzz interface for custom mutator and Python mutator */
+size_t fuzz_py(u8*, size_t, u8*, size_t, unsigned int);
+void   fuzz_py_original(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(u8*, size_t);
 u32    post_trim_py(u8);

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;
 

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);
 

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,7 +159,8 @@ void init_py(unsigned int seed) {
   }
 }
 
-void fuzz_py(char* buf, size_t buflen, char* add_buf, size_t add_buflen,
+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;

src/afl-fuzz.c

@@ -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();