Exposing the process memory iterator interface to Python.

yara-python.c

@@ -28,6 +28,7 @@ limitations under the License.
 
 #include <time.h>
 #include <yara.h>
+#include <yara/proc.h>
 
 #if PY_VERSION_HEX < 0x02050000 && !defined(PY_SSIZE_T_MIN)
 typedef int Py_ssize_t;
@@ -2082,6 +2083,187 @@ static PyObject* yara_load(
   return (PyObject*) rules;
 }
 
+typedef struct
+{
+  PyObject_HEAD
+  PyObject* externals;
+  YR_MEMORY_BLOCK_ITERATOR* block_iterator;
+  YR_MEMORY_BLOCK* block;
+} ProcessMemoryIterator;
+
+static PyObject* ProcessMemoryIterator_getattro(
+    PyObject* self,
+    PyObject* name)
+{
+  return PyObject_GenericGetAttr(self, name);
+}
+
+static void ProcessMemoryIterator_dealloc(PyObject* self);
+
+static PyObject* ProcessMemoryIterator_next(PyObject* self);
+
+static PyTypeObject ProcessMemoryIterator_Type = {
+  PyVarObject_HEAD_INIT(NULL, 0)
+  "yara.ProcessMemoryIterator",  /*tp_name*/
+  sizeof(ProcessMemoryIterator), /*tp_basicsize*/
+  0,                          /*tp_itemsize*/
+  (destructor) ProcessMemoryIterator_dealloc, /*tp_dealloc*/
+  0,                          /*tp_print*/
+  0,                          /*tp_getattr*/
+  0,                          /*tp_setattr*/
+  0,                          /*tp_compare*/
+  0,                          /*tp_repr*/
+  0,                          /*tp_as_number*/
+  0,                          /*tp_as_sequence*/
+  0,                          /*tp_as_mapping*/
+  0,                          /*tp_hash */
+  0,                          /*tp_call*/
+  0,                          /*tp_str*/
+  ProcessMemoryIterator_getattro, /*tp_getattro*/
+  0,                          /*tp_setattro*/
+  0,                          /*tp_as_buffer*/
+  Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /*tp_flags*/
+  "ProcessMemoryIterator",    /* tp_doc */
+  0,                          /* tp_traverse */
+  0,                          /* tp_clear */
+  0,                          /* tp_richcompare */
+  0,                          /* tp_weaklistoffset */
+  PyObject_SelfIter,          /* tp_iter */
+  (iternextfunc) ProcessMemoryIterator_next,  /* tp_iternext */
+  0,                          /* tp_methods */ // TODO????
+  0,                          /* tp_members */
+  0,                          /* tp_getset */
+  0,                          /* tp_base */
+  0,                          /* tp_dict */
+  0,                          /* tp_descr_get */
+  0,                          /* tp_descr_set */
+  0,                          /* tp_dictoffset */
+  0,                          /* tp_init */
+  0,                          /* tp_alloc */
+  0,                          /* tp_new */
+};
+
+static ProcessMemoryIterator* ProcessMemoryIterator_NEW(void)
+{
+  ProcessMemoryIterator* it = PyObject_NEW(ProcessMemoryIterator, &ProcessMemoryIterator_Type);
+  if (it == NULL)
+    return NULL;
+
+  it->block_iterator = NULL;
+  it->block = NULL;
+
+  return it;
+}
+
+static void ProcessMemoryIterator_dealloc(
+    PyObject* self)
+{
+  ProcessMemoryIterator* it = (ProcessMemoryIterator*) self;
+
+  if (it->block_iterator != NULL)
+  {
+    yr_process_close_iterator(it->block_iterator);
+    PyMem_Free(it->block_iterator);
+    it->block_iterator = NULL;
+  }
+  PyObject_Del(self);
+}
+
+static PyObject* ProcessMemoryIterator_next(
+    PyObject* self)
+{
+  ProcessMemoryIterator* it = (ProcessMemoryIterator*) self;
+  int err;
+
+  // This indicates that the iterator has been used.
+  if (it->block_iterator == NULL)
+  {
+    PyErr_SetNone(PyExc_StopIteration);
+    return NULL;
+  }
+
+  // During the first invocation, we need to use get_first_memory_block.
+  if (it->block == NULL)
+    it->block = yr_process_get_first_memory_block(it->block_iterator);
+  else
+    it->block = yr_process_get_next_memory_block(it->block_iterator);
+
+  if (it->block == NULL)
+  {
+    PyErr_SetNone(PyExc_StopIteration);
+    return NULL;
+  }
+
+  uint8_t* data_ptr = yr_process_fetch_memory_block_data(it->block);
+  if (data_ptr == NULL)
+  {
+    // This is how we are notified that the process is done.
+    it->block = NULL;
+    err = yr_process_close_iterator(it->block_iterator);
+    PyMem_Free(it->block_iterator);
+    it->block_iterator = NULL;
+    if (err != 0)
+    {
+      return handle_error(err, NULL);
+    }
+
+    PyErr_SetNone(PyExc_StopIteration);
+    return NULL;
+  }
+
+  return PyBytes_FromStringAndSize(
+      (const char*) data_ptr,
+      it->block->size);
+}
+
+static PyObject* yara_process_memory_iterator(
+    PyObject* self,
+    PyObject* args,
+    PyObject* keywords)
+{
+  static char *kwlist[] = {
+    "pid", NULL};
+
+  unsigned int pid = UINT_MAX;
+  int err;
+
+  ProcessMemoryIterator *result;
+
+  if (!PyArg_ParseTupleAndKeywords(
+        args,
+        keywords,
+        "|I",
+        kwlist,
+        &pid))
+  {
+    return PyErr_Format(
+        PyExc_TypeError,
+        "Error parsing arguments.");
+  }
+
+  result = ProcessMemoryIterator_NEW();
+
+  result->block_iterator = PyMem_Malloc(sizeof(YR_MEMORY_BLOCK_ITERATOR));
+  if (result->block_iterator == NULL)
+    return PyErr_NoMemory();
+
+  // Fail early if we can't access the process with the given pid.
+  err = yr_process_open_iterator(pid, result->block_iterator);
+  if (err != 0)
+  {
+    PyMem_Free(result->block_iterator);
+    return handle_error(err, NULL);
+  }
+
+  result->block = yr_process_get_first_memory_block(result->block_iterator);
+  if (result->block == NULL)
+  {
+    PyMem_Free(result->block_iterator);
+    result->block_iterator = NULL;
+    return PyErr_NoMemory();
+  }
+  return (PyObject *) result;
+}
 
 void finalize(void)
 {
@@ -2102,6 +2284,13 @@ static PyMethodDef yara_methods[] = {
     METH_VARARGS | METH_KEYWORDS,
     "Loads a previously saved YARA rules file and returns an instance of class Rules"
   },
+  {
+    "process_memory_iterator",
+    (PyCFunction) yara_process_memory_iterator,
+    METH_VARARGS | METH_KEYWORDS,
+    "Returns an iterator over blocks of memory of a process.\n"
+    "Signature: process_memory_iterator(pid=None)"
+  },
   { NULL, NULL }
 };