Merge pull request #221 from h1994st/master

Uniform Python and custom mutator API

1 files changed, 165 insertions, 237 deletions

diff --git a/src/afl-fuzz-python.c b/src/afl-fuzz-python.c
index 42286527..c22e4402 100644
--- a/src/afl-fuzz-python.c
+++ b/src/afl-fuzz-python.c
@@ -28,195 +28,244 @@
 /* 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 (!module_name) return 1;
 
-    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/...).");
+  Py_Initialize();
 
 #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) {
 
-    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");
 
-      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_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");
+    for (py_idx = 0; py_idx < PY_FUNC_COUNT; ++py_idx) {
 
-      for (py_idx = 0; py_idx < PY_FUNC_COUNT; ++py_idx) {
+      if (!py_functions[py_idx] || !PyCallable_Check(py_functions[py_idx])) {
 
-        if (!py_functions[py_idx] || !PyCallable_Check(py_functions[py_idx])) {
+        if (py_idx == PY_FUNC_PRE_SAVE) {
 
-          if (py_idx >= PY_FUNC_INIT_TRIM && py_idx <= PY_FUNC_TRIM) {
+          // Implenting the pre_save API is optional for now
+          if (PyErr_Occurred()) PyErr_Print();
 
-            // Implementing the trim API is optional for now
-            if (PyErr_Occurred()) PyErr_Print();
-            py_notrim = 1;
+        } else 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.");
+    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;
+
+  }
+
+  return 0;
+
+}
 
-      PyObject *py_args, *py_value;
+void finalize_py_module() {
 
-      /* Provide the init function a seed for the Python RNG */
-      py_args = PyTuple_New(1);
+  if (py_module != NULL) {
+
+    u32 i;
+    for (i = 0; i < PY_FUNC_COUNT; ++i)
+      Py_XDECREF(py_functions[i]);
+
+    Py_DECREF(py_module);
+
+  }
+
+  Py_Finalize();
+
+}
+
+void init_py(unsigned int seed) {
+  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));
+  py_value = PyLong_FromLong(seed);
 #else
-      py_value = PyInt_FromLong(UR(0xFFFFFFFF));
+  py_value = PyInt_FromLong(seed);
 #endif
 
-      if (!py_value) {
+  if (!py_value) {
+
+    Py_DECREF(py_args);
+    fprintf(stderr, "Cannot convert argument\n");
+    return;
 
-        Py_DECREF(py_args);
-        fprintf(stderr, "Cannot convert argument\n");
-        return 1;
+  }
 
-      }
+  PyTuple_SetItem(py_args, 0, py_value);
 
-      PyTuple_SetItem(py_args, 0, py_value);
+  py_value = PyObject_CallObject(py_functions[PY_FUNC_INIT], py_args);
 
-      py_value = PyObject_CallObject(py_functions[PY_FUNC_INIT], py_args);
+  Py_DECREF(py_args);
 
-      Py_DECREF(py_args);
+  if (py_value == NULL) {
 
-      if (py_value == NULL) {
+    PyErr_Print();
+    fprintf(stderr, "Call failed\n");
+    return;
 
-        PyErr_Print();
-        fprintf(stderr, "Call failed\n");
-        return 1;
+  }
+}
 
-      }
+size_t fuzz_py(u8* buf, size_t buf_size,
+               u8* add_buf, size_t add_buf_size,
+               u8* mutated_out, size_t max_size) {
 
-    } else {
+  size_t mutated_size;
+  PyObject *py_args, *py_value;
+  py_args = PyTuple_New(3);
 
-      PyErr_Print();
-      fprintf(stderr, "Failed to load \"%s\"\n", module_name);
-      return 1;
+  /* buf */
+  py_value = PyByteArray_FromStringAndSize(buf, buf_size);
+  if (!py_value) {
 
-    }
+    Py_DECREF(py_args);
+    FATAL("Failed to convert arguments");
 
   }
 
-  return 0;
+  PyTuple_SetItem(py_args, 0, py_value);
 
-}
+  /* add_buf */
+  py_value = PyByteArray_FromStringAndSize(add_buf, add_buf_size);
+  if (!py_value) {
+
+    Py_DECREF(py_args);
+    FATAL("Failed to convert arguments");
 
-void finalize_py() {
+  }
 
-  if (py_module != NULL) {
+  PyTuple_SetItem(py_args, 1, py_value);
 
-    u32 i;
-    for (i = 0; i < PY_FUNC_COUNT; ++i)
-      Py_XDECREF(py_functions[i]);
+  /* max_size */
+#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_module);
+    Py_DECREF(py_args);
+    FATAL("Failed to convert arguments");
 
   }
 
-  Py_Finalize();
+  PyTuple_SetItem(py_args, 2, py_value);
 
-}
+  py_value = PyObject_CallObject(py_functions[PY_FUNC_FUZZ], py_args);
 
-void fuzz_py(char* buf, size_t buflen, char* add_buf, size_t add_buflen,
-             char** ret, size_t* retlen) {
+  Py_DECREF(py_args);
 
-  if (py_module != NULL) {
+  if (py_value != NULL) {
 
-    PyObject *py_args, *py_value;
-    py_args = PyTuple_New(2);
-    py_value = PyByteArray_FromStringAndSize(buf, buflen);
-    if (!py_value) {
+    mutated_size = PyByteArray_Size(py_value);
+    memcpy(mutated_out, PyByteArray_AsString(py_value), mutated_size);
+    Py_DECREF(py_value);
+    return mutated_size;
 
-      Py_DECREF(py_args);
-      fprintf(stderr, "Cannot convert argument\n");
-      return;
+  } else {
 
-    }
+    PyErr_Print();
+    FATAL("Call failed");
 
-    PyTuple_SetItem(py_args, 0, py_value);
+  }
 
-    py_value = PyByteArray_FromStringAndSize(add_buf, add_buflen);
-    if (!py_value) {
+}
 
-      Py_DECREF(py_args);
-      fprintf(stderr, "Cannot convert argument\n");
-      return;
+size_t pre_save_py(u8* buf, size_t buf_size, u8** out_buf) {
 
-    }
+  size_t out_buf_size;
+  PyObject *py_args, *py_value;
+  py_args = PyTuple_New(2);
+  py_value = PyByteArray_FromStringAndSize(buf, buf_size);
+  if (!py_value) {
 
-    PyTuple_SetItem(py_args, 1, py_value);
+    Py_DECREF(py_args);
+    FATAL("Failed to convert arguments");
 
-    py_value = PyObject_CallObject(py_functions[PY_FUNC_FUZZ], py_args);
+  }
 
-    Py_DECREF(py_args);
+  PyTuple_SetItem(py_args, 0, py_value);
 
-    if (py_value != NULL) {
+  py_value = PyObject_CallObject(py_functions[PY_FUNC_PRE_SAVE], py_args);
 
-      *retlen = PyByteArray_Size(py_value);
-      *ret = malloc(*retlen);
-      memcpy(*ret, PyByteArray_AsString(py_value), *retlen);
-      Py_DECREF(py_value);
+  Py_DECREF(py_args);
 
-    } else {
+  if (py_value != NULL) {
 
-      PyErr_Print();
-      fprintf(stderr, "Call failed\n");
-      return;
+    out_buf_size = PyByteArray_Size(py_value);
+    *out_buf = malloc(out_buf_size);
+    memcpy(*out_buf, PyByteArray_AsString(py_value), out_buf_size);
+    Py_DECREF(py_value);
+    return out_buf_size;
 
-    }
+  } else {
+
+    PyErr_Print();
+    FATAL("Call failed");
 
   }
 
 }
 
-u32 init_trim_py(char* buf, size_t buflen) {
+u32 init_trim_py(u8* buf, size_t buf_size) {
 
   PyObject *py_args, *py_value;
 
   py_args = PyTuple_New(1);
-  py_value = PyByteArray_FromStringAndSize(buf, buflen);
+  py_value = PyByteArray_FromStringAndSize(buf, buf_size);
   if (!py_value) {
 
     Py_DECREF(py_args);
@@ -248,7 +297,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;
 
@@ -286,7 +335,7 @@ u32 post_trim_py(char success) {
 
 }
 
-void trim_py(char** ret, size_t* retlen) {
+void trim_py(u8** out_buf, size_t* out_buf_size) {
 
   PyObject *py_args, *py_value;
 
@@ -296,9 +345,9 @@ void trim_py(char** ret, size_t* retlen) {
 
   if (py_value != NULL) {
 
-    *retlen = PyByteArray_Size(py_value);
-    *ret = malloc(*retlen);
-    memcpy(*ret, PyByteArray_AsString(py_value), *retlen);
+    *out_buf_size = PyByteArray_Size(py_value);
+    *out_buf = malloc(*out_buf_size);
+    memcpy(*out_buf, PyByteArray_AsString(py_value), *out_buf_size);
     Py_DECREF(py_value);
 
   } else {
@@ -310,126 +359,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 */