Add two new hooks for the custom mutator

- `afl_custom_queue_get` and `afl_custom_queue_new_entry`
- Update the corresponding document and examples
This commit is contained in:
h1994st
2020-03-07 16:28:48 -05:00
parent dc0b2dda5e
commit 8f93cf5c55
8 changed files with 313 additions and 20 deletions

View File

@ -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 */