about summary refs log tree commit diff
diff options
context:
space:
mode:
authoryangzao <yangzaocn@outlook.com>2023-11-25 21:18:32 -0700
committeryangzao <yangzaocn@outlook.com>2023-11-25 21:18:32 -0700
commitfaedb3fb29186c29a4f0cf28daa5d07350ed8094 (patch)
tree2d69a3ab094de3ef59e91c225abb17ac50b35e63
parent8af74bcaeebbe2407006333024d8803baacdb4e2 (diff)
downloadafl++-faedb3fb29186c29a4f0cf28daa5d07350ed8094.tar.gz
update python module
-rw-r--r--custom_mutators/examples/example.py5
-rw-r--r--include/afl-fuzz.h2
-rw-r--r--src/afl-fuzz-python.c32
3 files changed, 39 insertions, 0 deletions
diff --git a/custom_mutators/examples/example.py b/custom_mutators/examples/example.py
index 3a6d22e4..830f302f 100644
--- a/custom_mutators/examples/example.py
+++ b/custom_mutators/examples/example.py
@@ -133,6 +133,11 @@ def fuzz(buf, add_buf, max_size):
 #     @return: The buffer containing the test case after
 #     '''
 #     return buf
+# def post_run():
+#     '''
+#     Called after each time the execution of the target program by AFL++
+#     '''
+#     pass
 #
 # def havoc_mutation(buf, max_size):
 #     '''
diff --git a/include/afl-fuzz.h b/include/afl-fuzz.h
index 94f48009..f1813df6 100644
--- a/include/afl-fuzz.h
+++ b/include/afl-fuzz.h
@@ -345,6 +345,7 @@ enum {
   /* 13 */ PY_FUNC_DESCRIBE,
   /* 14 */ PY_FUNC_FUZZ_SEND,
   /* 15 */ PY_FUNC_SPLICE_OPTOUT,
+  /* 16 */ PY_FUNC_POST_RUN,
   PY_FUNC_COUNT
 
 };
@@ -1085,6 +1086,7 @@ void                   finalize_py_module(void *);
 
 u32         fuzz_count_py(void *, const u8 *, size_t);
 void        fuzz_send_py(void *, const u8 *, size_t);
+void        post_run_py(void *);
 size_t      post_process_py(void *, u8 *, size_t, u8 **);
 s32         init_trim_py(void *, u8 *, size_t);
 s32         post_trim_py(void *, u8);
diff --git a/src/afl-fuzz-python.c b/src/afl-fuzz-python.c
index 7dad0770..1b287405 100644
--- a/src/afl-fuzz-python.c
+++ b/src/afl-fuzz-python.c
@@ -249,6 +249,8 @@ static py_mutator_t *init_py_module(afl_state_t *afl, u8 *module_name) {
         PyObject_GetAttrString(py_module, "queue_get");
     py_functions[PY_FUNC_FUZZ_SEND] =
         PyObject_GetAttrString(py_module, "fuzz_send");
+    py_functions[PY_FUNC_POST_RUN] =
+        PyObject_GetAttrString(py_module, "post_run");
     py_functions[PY_FUNC_SPLICE_OPTOUT] =
         PyObject_GetAttrString(py_module, "splice_optout");
     if (py_functions[PY_FUNC_SPLICE_OPTOUT]) { afl->custom_splice_optout = 1; }
@@ -468,6 +470,12 @@ struct custom_mutator *load_custom_mutator_py(afl_state_t *afl,
 
   }
 
+  if (py_functions[PY_FUNC_POST_RUN]) {
+
+    mutator->afl_custom_post_run = post_run_py;
+
+  }
+
   if (py_functions[PY_FUNC_SPLICE_OPTOUT]) {
 
     mutator->afl_custom_splice_optout = splice_optout_py;
@@ -925,6 +933,30 @@ void fuzz_send_py(void *py_mutator, const u8 *buf, size_t buf_size) {
 
 }
 
+void post_run_py(void *py_mutator) {
+
+  PyObject *py_args, *py_value;
+
+  py_args = PyTuple_New(0);
+  py_value = PyObject_CallObject(
+      ((py_mutator_t *)py_mutator)
+          ->py_functions[PY_FUNC_POST_RUN],
+      py_args);
+  Py_DECREF(py_args);
+
+  if (py_value != NULL) {
+
+    Py_DECREF(py_value);
+
+  } else {
+
+    PyErr_Print();
+    FATAL("Call failed");
+
+  }
+
+}
+
 u8 queue_new_entry_py(void *py_mutator, const u8 *filename_new_queue,
                       const u8 *filename_orig_queue) {