From faedb3fb29186c29a4f0cf28daa5d07350ed8094 Mon Sep 17 00:00:00 2001
From: yangzao <yangzaocn@outlook.com>
Date: Sat, 25 Nov 2023 21:18:32 -0700
Subject: update python module

---
 src/afl-fuzz-python.c | 32 ++++++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

(limited to 'src/afl-fuzz-python.c')

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) {
 
-- 
cgit 1.4.1