1 files changed, 329 insertions, 0 deletions
diff --git a/src/afl-fuzz-python.c b/src/afl-fuzz-python.c
new file mode 100644
index 00000000..ed158e6c
--- /dev/null
+++ b/src/afl-fuzz-python.c
@@ -0,0 +1,329 @@
+/*
+   american fuzzy lop - fuzzer code
+   --------------------------------
+
+   Written and maintained by Michal Zalewski <lcamtuf@google.com>
+
+   Forkserver design by Jann Horn <jannhorn@googlemail.com>
+
+   Copyright 2013, 2014, 2015, 2016, 2017 Google Inc. All rights reserved.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at:
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+   This is the real deal: the program takes an instrumented binary and
+   attempts a variety of basic fuzzing tricks, paying close attention to
+   how they affect the execution path.
+
+ */
+
+#include "afl-fuzz.h"
+
+/* Python stuff */
+#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_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();
+            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);
+            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.");
+      }
+
+      PyObject *py_args, *py_value;
+
+      /* Provide the init function a seed for the Python RNG */
+      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);
+
+      py_value = PyObject_CallObject(py_functions[PY_FUNC_INIT], py_args);
+
+      Py_DECREF(py_args);
+
+      if (py_value == NULL) {
+        PyErr_Print();
+        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) {
+
+  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);
+
+    py_value = PyObject_CallObject(py_functions[PY_FUNC_FUZZ], py_args);
+
+    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");
+      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);
+
+  py_value = PyObject_CallObject(py_functions[PY_FUNC_INIT_TRIM], py_args);
+  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);
+
+  py_value = PyObject_CallObject(py_functions[PY_FUNC_POST_TRIM], py_args);
+  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);
+  py_value = PyObject_CallObject(py_functions[PY_FUNC_TRIM], py_args);
+  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) {
+
+  static u8 tmp[64];
+  static u8 clean_trace[MAP_SIZE];
+
+  u8  needs_write = 0, fault = 0;
+  u32 trim_exec = 0;
+  u32 orig_len = q->len;
+
+  stage_name = tmp;
+  bytes_trim_in += q->len;
+
+  /* Initialize trimming in the Python module */
+  stage_cur = 0;
+  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);
+
+  while(stage_cur < stage_max) {
+    sprintf(tmp, "ptrim %s", DI(trim_exec));
+
+    u32 cksum;
+
+    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");
+
+    write_to_testcase(retbuf, retlen);
+
+    fault = run_target(argv, exec_tmout);
+    ++trim_execs;
+
+    if (stop_soon || fault == FAULT_ERROR) goto abort_trimming;
+
+    cksum = hash32(trace_bits, MAP_SIZE, HASH_CONST);
+
+    if (cksum == q->exec_cksum) {
+
+      q->len = retlen;
+      memcpy(in_buf, retbuf, retlen);
+
+      /* Let's save a clean trace, which will be needed by
+         update_bitmap_score once we're done with the trimming stuff. */
+
+      if (!needs_write) {
+
+        needs_write = 1;
+        memcpy(clean_trace, trace_bits, MAP_SIZE);
+
+      }
+
+      /* Tell the Python module that the trimming was successful */
+      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);
+    } 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);
+    }
+
+      /* Since this can be slow, update the screen every now and then. */
+
+      if (!(trim_exec++ % stats_update_freq)) show_stats();
+  }
+
+  if (not_on_tty && debug)
+    SAYF("[Python Trimming] DONE: %u bytes -> %u bytes", orig_len, q->len);
+
+  /* If we have made changes to in_buf, we also need to update the on-disk
+     version of the test case. */
+
+  if (needs_write) {
+
+    s32 fd;
+
+    unlink(q->fname); /* ignore errors */
+
+    fd = open(q->fname, O_WRONLY | O_CREAT | O_EXCL, 0600);
+
+    if (fd < 0) PFATAL("Unable to create '%s'", q->fname);
+
+    ck_write(fd, in_buf, q->len, q->fname);
+    close(fd);
+
+    memcpy(trace_bits, clean_trace, MAP_SIZE);
+    update_bitmap_score(q);
+
+  }
+
+
+
+abort_trimming:
+
+  bytes_trim_out += q->len;
+  return fault;
+
+}
+
+#endif /* USE_PYTHON */