added afl_custom_fuzz_count

1 files changed, 44 insertions, 0 deletions

diff --git a/src/afl-fuzz-python.c b/src/afl-fuzz-python.c
index e540f548..68540dd7 100644
--- a/src/afl-fuzz-python.c
+++ b/src/afl-fuzz-python.c
@@ -347,6 +347,12 @@ struct custom_mutator *load_custom_mutator_py(afl_state_t *afl,
 
   }
 
+  if (py_functions[PY_FUNC_FUZZ_COUNT]) {
+
+    mutator->afl_custom_fuzz_count = fuzz_count_py;
+
+  }
+
   if (py_functions[PY_FUNC_POST_TRIM]) {
 
     mutator->afl_custom_post_trim = post_trim_py;
@@ -477,6 +483,44 @@ s32 init_trim_py(void *py_mutator, u8 *buf, size_t buf_size) {
 
 }
 
+u32 fuzz_count_py(void *py_mutator, const u8 *buf, size_t buf_size) {
+
+  PyObject *py_args, *py_value;
+
+  py_args = PyTuple_New(1);
+  py_value = PyByteArray_FromStringAndSize(buf, buf_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_mutator_t *)py_mutator)->py_functions[PY_FUNC_FUZZ_COUNT], py_args);
+  Py_DECREF(py_args);
+
+  if (py_value != NULL) {
+
+  #if PY_MAJOR_VERSION >= 3
+    u32 retcnt = (u32)PyLong_AsLong(py_value);
+  #else
+    u32 retcnt = PyInt_AsLong(py_value);
+  #endif
+    Py_DECREF(py_value);
+    return retcnt;
+
+  } else {
+
+    PyErr_Print();
+    FATAL("Call failed");
+
+  }
+
+}
+
 s32 post_trim_py(void *py_mutator, u8 success) {
 
   PyObject *py_args, *py_value;