Add two new hooks for the custom mutator

- `afl_custom_queue_get` and `afl_custom_queue_new_entry`
- Update the corresponding document and examples

1 files changed, 118 insertions, 2 deletions

diff --git a/src/afl-fuzz-python.c b/src/afl-fuzz-python.c
index 32f9f6ab..8ceb6957 100644
--- a/src/afl-fuzz-python.c
+++ b/src/afl-fuzz-python.c
@@ -55,8 +55,14 @@ int init_py_module(u8* module_name) {
     py_functions[PY_FUNC_POST_TRIM] =
         PyObject_GetAttrString(py_module, "post_trim");
     py_functions[PY_FUNC_TRIM] = PyObject_GetAttrString(py_module, "trim");
-    py_functions[PY_FUNC_HAVOC_MUTATION] = PyObject_GetAttrString(py_module, "havoc_mutation");
-    py_functions[PY_FUNC_HAVOC_MUTATION_PROBABILITY] = PyObject_GetAttrString(py_module, "havoc_mutation_probability");
+    py_functions[PY_FUNC_HAVOC_MUTATION] =
+        PyObject_GetAttrString(py_module, "havoc_mutation");
+    py_functions[PY_FUNC_HAVOC_MUTATION_PROBABILITY] =
+        PyObject_GetAttrString(py_module, "havoc_mutation_probability");
+    py_functions[PY_FUNC_QUEUE_GET] =
+        PyObject_GetAttrString(py_module, "queue_get");
+    py_functions[PY_FUNC_QUEUE_NEW_ENTRY] =
+        PyObject_GetAttrString(py_module, "queue_new_entry");
 
     for (py_idx = 0; py_idx < PY_FUNC_COUNT; ++py_idx) {
 
@@ -73,6 +79,12 @@ int init_py_module(u8* module_name) {
           if (PyErr_Occurred()) PyErr_Print();
           py_notrim = 1;
 
+        } else if ((py_idx >= PY_FUNC_HAVOC_MUTATION) &&
+                   (py_idx <= PY_FUNC_QUEUE_NEW_ENTRY)) {
+
+          // Implenting the havoc and queue API is optional for now
+          if (PyErr_Occurred()) PyErr_Print();
+
         } else {
 
           if (PyErr_Occurred()) PyErr_Print();
@@ -442,5 +454,109 @@ u8 havoc_mutation_probability_py(void) {
 
 }
 
+u8 queue_get_py(const u8* filename) {
+
+  PyObject *py_args, *py_value;
+
+  py_args = PyTuple_New(1);
+
+  // File name
+#if PY_MAJOR_VERSION >= 3
+  py_value = PyUnicode_FromString(filename);
+#else
+  py_value = PyString_FromString(filename);
+#endif
+  if (!py_value) {
+
+    Py_DECREF(py_args);
+    FATAL("Failed to convert arguments");
+
+  }
+
+  PyTuple_SetItem(py_args, 0, py_value);
+
+  // Call Python function
+  py_value = PyObject_CallObject(py_functions[PY_FUNC_QUEUE_GET], py_args);
+  Py_DECREF(py_args);
+
+  if (py_value != NULL) {
+
+    int ret = PyObject_IsTrue(py_value);
+    Py_DECREF(py_value);
+    
+    if (ret == -1) {
+
+      PyErr_Print();
+      FATAL("Failed to convert return value");
+
+    }
+
+    return (u8) ret & 0xFF;
+
+  } else {
+    
+    PyErr_Print();
+    FATAL("Call failed");
+
+  }
+
+}
+
+void queue_new_entry_py(const u8* filename_new_queue,
+                        const u8* filename_orig_queue) {
+
+  PyObject *py_args, *py_value;
+
+  py_args = PyTuple_New(2);
+
+  // New queue
+#if PY_MAJOR_VERSION >= 3
+  py_value = PyUnicode_FromString(filename_new_queue);
+#else
+  py_value = PyString_FromString(filename_new_queue);
+#endif
+  if (!py_value) {
+  
+    Py_DECREF(py_args);
+    FATAL("Failed to convert arguments");
+
+  }
+
+  PyTuple_SetItem(py_args, 0, py_value);
+
+  // Orig queue
+  py_value = Py_None;
+  if (filename_orig_queue) {
+
+#if PY_MAJOR_VERSION >= 3
+    py_value = PyUnicode_FromString(filename_orig_queue);
+#else
+    py_value = PyString_FromString(filename_orig_queue);
+#endif
+    if (!py_value) {
+ 
+      Py_DECREF(py_args);
+      FATAL("Failed to convert arguments");
+
+    }
+
+  }
+
+  PyTuple_SetItem(py_args, 1, py_value);
+
+  // Call
+  py_value = PyObject_CallObject(py_functions[PY_FUNC_QUEUE_NEW_ENTRY],
+                                 py_args);
+  Py_DECREF(py_args);
+
+  if (py_value == NULL) {
+
+    PyErr_Print();
+    FATAL("Call failed");
+
+  }
+
+}
+
 #endif                                                        /* USE_PYTHON */