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,138 @@ 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,
@@ -2092,6 +2225,9 @@ static PyObject* yara_process_memory_iterator(
     "pid", NULL};
 
   unsigned int pid = UINT_MAX;
+  int err;
+
+  ProcessMemoryIterator *result;
 
   if (!PyArg_ParseTupleAndKeywords(
         args,
@@ -2104,7 +2240,29 @@ static PyObject* yara_process_memory_iterator(
         PyExc_TypeError,
         "Error parsing arguments.");
   }
-  return Py_BuildValue("I", pid);
+
+  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)