diff options
Diffstat (limited to 'src/afl-fuzz-python.c')
| -rw-r--r-- | src/afl-fuzz-python.c | 117 |
1 files changed, 95 insertions, 22 deletions
diff --git a/src/afl-fuzz-python.c b/src/afl-fuzz-python.c index ed158e6c..e22291b5 100644 --- a/src/afl-fuzz-python.c +++ b/src/afl-fuzz-python.c @@ -26,45 +26,62 @@ #ifdef USE_PYTHON int init_py() { + Py_Initialize(); u8* module_name = getenv("AFL_PYTHON_MODULE"); if (module_name) { + PyObject* py_name = PyString_FromString(module_name); py_module = PyImport_Import(py_name); Py_DECREF(py_name); if (py_module != NULL) { + u8 py_notrim = 0; py_functions[PY_FUNC_INIT] = PyObject_GetAttrString(py_module, "init"); py_functions[PY_FUNC_FUZZ] = PyObject_GetAttrString(py_module, "fuzz"); - py_functions[PY_FUNC_INIT_TRIM] = PyObject_GetAttrString(py_module, "init_trim"); - py_functions[PY_FUNC_POST_TRIM] = PyObject_GetAttrString(py_module, "post_trim"); + py_functions[PY_FUNC_INIT_TRIM] = + PyObject_GetAttrString(py_module, "init_trim"); + py_functions[PY_FUNC_POST_TRIM] = + PyObject_GetAttrString(py_module, "post_trim"); py_functions[PY_FUNC_TRIM] = PyObject_GetAttrString(py_module, "trim"); for (u8 py_idx = 0; py_idx < PY_FUNC_COUNT; ++py_idx) { + if (!py_functions[py_idx] || !PyCallable_Check(py_functions[py_idx])) { + if (py_idx >= PY_FUNC_INIT_TRIM && py_idx <= PY_FUNC_TRIM) { + // Implementing the trim API is optional for now - if (PyErr_Occurred()) - PyErr_Print(); + if (PyErr_Occurred()) PyErr_Print(); py_notrim = 1; + } else { - if (PyErr_Occurred()) - PyErr_Print(); - fprintf(stderr, "Cannot find/call function with index %d in external Python module.\n", py_idx); + + if (PyErr_Occurred()) PyErr_Print(); + fprintf(stderr, + "Cannot find/call function with index %d in external " + "Python module.\n", + py_idx); return 1; + } + } } if (py_notrim) { + py_functions[PY_FUNC_INIT_TRIM] = NULL; py_functions[PY_FUNC_POST_TRIM] = NULL; py_functions[PY_FUNC_TRIM] = NULL; - WARNF("Python module does not implement trim API, standard trimming will be used."); + WARNF( + "Python module does not implement trim API, standard trimming will " + "be used."); + } PyObject *py_args, *py_value; @@ -73,9 +90,11 @@ int init_py() { py_args = PyTuple_New(1); py_value = PyInt_FromLong(UR(0xFFFFFFFF)); if (!py_value) { + Py_DECREF(py_args); fprintf(stderr, "Cannot convert argument\n"); return 1; + } PyTuple_SetItem(py_args, 0, py_value); @@ -85,51 +104,68 @@ int init_py() { Py_DECREF(py_args); if (py_value == NULL) { + PyErr_Print(); - fprintf(stderr,"Call failed\n"); + fprintf(stderr, "Call failed\n"); return 1; + } + } else { + PyErr_Print(); fprintf(stderr, "Failed to load \"%s\"\n", module_name); return 1; + } + } return 0; + } void finalize_py() { + if (py_module != NULL) { + u32 i; for (i = 0; i < PY_FUNC_COUNT; ++i) Py_XDECREF(py_functions[i]); Py_DECREF(py_module); + } Py_Finalize(); + } -void fuzz_py(char* buf, size_t buflen, char* add_buf, size_t add_buflen, char** ret, size_t* retlen) { +void fuzz_py(char* buf, size_t buflen, char* add_buf, size_t add_buflen, + char** ret, size_t* retlen) { if (py_module != NULL) { + PyObject *py_args, *py_value; py_args = PyTuple_New(2); py_value = PyByteArray_FromStringAndSize(buf, buflen); if (!py_value) { + Py_DECREF(py_args); fprintf(stderr, "Cannot convert argument\n"); return; + } PyTuple_SetItem(py_args, 0, py_value); py_value = PyByteArray_FromStringAndSize(add_buf, add_buflen); if (!py_value) { + Py_DECREF(py_args); fprintf(stderr, "Cannot convert argument\n"); return; + } PyTuple_SetItem(py_args, 1, py_value); @@ -139,26 +175,35 @@ void fuzz_py(char* buf, size_t buflen, char* add_buf, size_t add_buflen, char** Py_DECREF(py_args); if (py_value != NULL) { + *retlen = PyByteArray_Size(py_value); *ret = malloc(*retlen); memcpy(*ret, PyByteArray_AsString(py_value), *retlen); Py_DECREF(py_value); + } else { + PyErr_Print(); - fprintf(stderr,"Call failed\n"); + fprintf(stderr, "Call failed\n"); return; + } + } + } u32 init_trim_py(char* buf, size_t buflen) { + PyObject *py_args, *py_value; py_args = PyTuple_New(1); py_value = PyByteArray_FromStringAndSize(buf, buflen); if (!py_value) { + Py_DECREF(py_args); FATAL("Failed to convert arguments"); + } PyTuple_SetItem(py_args, 0, py_value); @@ -167,24 +212,32 @@ u32 init_trim_py(char* buf, size_t buflen) { Py_DECREF(py_args); if (py_value != NULL) { + u32 retcnt = PyInt_AsLong(py_value); Py_DECREF(py_value); return retcnt; + } else { + PyErr_Print(); FATAL("Call failed"); + } + } u32 post_trim_py(char success) { + PyObject *py_args, *py_value; py_args = PyTuple_New(1); py_value = PyBool_FromLong(success); if (!py_value) { + Py_DECREF(py_args); FATAL("Failed to convert arguments"); + } PyTuple_SetItem(py_args, 0, py_value); @@ -193,16 +246,22 @@ u32 post_trim_py(char success) { Py_DECREF(py_args); if (py_value != NULL) { + u32 retcnt = PyInt_AsLong(py_value); Py_DECREF(py_value); return retcnt; + } else { + PyErr_Print(); FATAL("Call failed"); + } + } void trim_py(char** ret, size_t* retlen) { + PyObject *py_args, *py_value; py_args = PyTuple_New(0); @@ -210,14 +269,19 @@ void trim_py(char** ret, size_t* retlen) { Py_DECREF(py_args); if (py_value != NULL) { + *retlen = PyByteArray_Size(py_value); *ret = malloc(*retlen); memcpy(*ret, PyByteArray_AsString(py_value), *retlen); Py_DECREF(py_value); + } else { + PyErr_Print(); FATAL("Call failed"); + } + } u8 trim_case_python(char** argv, struct queue_entry* q, u8* in_buf) { @@ -237,20 +301,24 @@ u8 trim_case_python(char** argv, struct queue_entry* q, u8* in_buf) { stage_max = init_trim_py(in_buf, q->len); if (not_on_tty && debug) - SAYF("[Python Trimming] START: Max %d iterations, %u bytes", stage_max, q->len); + SAYF("[Python Trimming] START: Max %d iterations, %u bytes", stage_max, + q->len); + + while (stage_cur < stage_max) { - while(stage_cur < stage_max) { sprintf(tmp, "ptrim %s", DI(trim_exec)); u32 cksum; - char* retbuf = NULL; + char* retbuf = NULL; size_t retlen = 0; trim_py(&retbuf, &retlen); if (retlen > orig_len) - FATAL("Trimmed data returned by Python module is larger than original data"); + FATAL( + "Trimmed data returned by Python module is larger than original " + "data"); write_to_testcase(retbuf, retlen); @@ -280,17 +348,23 @@ u8 trim_case_python(char** argv, struct queue_entry* q, u8* in_buf) { stage_cur = post_trim_py(1); if (not_on_tty && debug) - SAYF("[Python Trimming] SUCCESS: %d/%d iterations (now at %u bytes)", stage_cur, stage_max, q->len); + SAYF("[Python Trimming] SUCCESS: %d/%d iterations (now at %u bytes)", + stage_cur, stage_max, q->len); + } else { + /* Tell the Python module that the trimming was unsuccessful */ stage_cur = post_trim_py(0); if (not_on_tty && debug) - SAYF("[Python Trimming] FAILURE: %d/%d iterations", stage_cur, stage_max); + SAYF("[Python Trimming] FAILURE: %d/%d iterations", stage_cur, + stage_max); + } - /* Since this can be slow, update the screen every now and then. */ + /* Since this can be slow, update the screen every now and then. */ + + if (!(trim_exec++ % stats_update_freq)) show_stats(); - if (!(trim_exec++ % stats_update_freq)) show_stats(); } if (not_on_tty && debug) @@ -303,7 +377,7 @@ u8 trim_case_python(char** argv, struct queue_entry* q, u8* in_buf) { s32 fd; - unlink(q->fname); /* ignore errors */ + unlink(q->fname); /* ignore errors */ fd = open(q->fname, O_WRONLY | O_CREAT | O_EXCL, 0600); @@ -317,8 +391,6 @@ u8 trim_case_python(char** argv, struct queue_entry* q, u8* in_buf) { } - - abort_trimming: bytes_trim_out += q->len; @@ -327,3 +399,4 @@ abort_trimming: } #endif /* USE_PYTHON */ + |