Callback on include

.gitmodules

@@ -1,3 +1,4 @@
 [submodule "yara"]
 	path = yara
-url=https://[email protected]/VirusTotal/yara.git
+	url=https://[email protected]/edhoedt/yara
+	branch = callback_on_include

yara-python.c

@@ -68,6 +68,10 @@ This module allows you to apply YARA rules to files or strings.\n\
 For complete documentation please visit:\n\
 https://plusvic.github.io/yara\n"
 
+#if defined(_WIN32) || defined(__CYGWIN__)
+#include <string.h>
+#define strdup _strdup
+#endif
 
 // Match object
 
@@ -1700,14 +1704,99 @@ void raise_exception_on_error_or_warning(
 
 ////////////////////////////////////////////////////////////////////////////////
 
+const char* yara_include_callback(
+  const char* include_name,
+  const char* calling_rule_filename,
+  const char* calling_rule_namespace,
+  void* user_data)
+{
+  PyObject* callback = (PyObject*) user_data;
+  PyObject* py_incl_name = NULL;
+  PyObject* py_calling_fn = NULL;
+  PyObject* py_calling_ns = NULL;
+
+  PyGILState_STATE gil_state = PyGILState_Ensure();
+
+  if (include_name != NULL)
+  {
+    py_incl_name = PY_STRING(include_name);
+  }
+  else //safeguard: should never happen for 'include_name'
+  {
+    py_incl_name = Py_None;
+  }
+  if (calling_rule_filename != NULL)
+  {
+    py_calling_fn = PY_STRING(calling_rule_filename);
+  }
+  else
+  {
+    py_calling_fn = Py_None;
+  }
+  if (calling_rule_namespace != NULL)
+  {
+    py_calling_ns = PY_STRING(calling_rule_namespace);
+  }
+  else
+  {
+    py_calling_ns = Py_None;
+  }
+
+  Py_INCREF(callback);
+  PyObject *type=NULL, *value=NULL, *traceback=NULL;
+  PyErr_Fetch(&type, &value, &traceback);
+  PyObject* result =  PyObject_CallFunctionObjArgs(callback,
+                                                   py_incl_name,
+                                                   py_calling_fn,
+                                                   py_calling_ns,
+                                                   NULL);
+  PyErr_Restore(type, value, traceback);
+
+  Py_DECREF(callback);
+  Py_DECREF(py_incl_name);
+  Py_XDECREF(py_calling_fn);
+  Py_XDECREF(py_calling_ns);
+
+  const char* cstring_result = NULL;
+  if (result != NULL && result != Py_None && PY_STRING_CHECK(result))
+  {
+    //transferring string ownership to C code
+    cstring_result = strdup(PY_STRING_TO_C(result));
+  }
+  else
+  {
+    if(PyErr_Occurred() == NULL)
+    {
+      PyErr_Format(PyExc_TypeError, "'include_callback' callback function must return a yara rule or rules set formated as a single ascii or unicode string");
+    }
+    cstring_result = NULL;
+  }
+  Py_XDECREF(result);
+  PyGILState_Release(gil_state);
+
+  return cstring_result;
+}
+
+void yara_include_free(
+  const char* result_ptr,
+  void* user_data)
+{
+  if(result_ptr != NULL)
+  {
+    free((void*)result_ptr);
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
 static PyObject* yara_compile(
     PyObject* self,
     PyObject* args,
     PyObject* keywords)
 {
   static char *kwlist[] = {
     "filepath", "source", "file", "filepaths", "sources",
-    "includes", "externals", "error_on_warning", NULL};
+    "includes", "externals", "error_on_warning", "include_callback", NULL};
 
   YR_COMPILER* compiler;
   YR_RULES* yara_rules;
@@ -1724,6 +1813,7 @@ static PyObject* yara_compile(
   PyObject* includes = NULL;
   PyObject* externals = NULL;
   PyObject* error_on_warning = NULL;
+  PyObject* include_callback = NULL;
 
   Py_ssize_t pos = 0;
 
@@ -1737,7 +1827,7 @@ static PyObject* yara_compile(
   if (PyArg_ParseTupleAndKeywords(
         args,
         keywords,
-        "|ssOOOOOO",
+        "|ssOOOOOOO",
         kwlist,
         &filepath,
         &source,
@@ -1746,8 +1836,10 @@ static PyObject* yara_compile(
         &sources_dict,
         &includes,
         &externals,
-        &error_on_warning))
+        &error_on_warning,
+        &include_callback))
   {
+
     error = yr_compiler_create(&compiler);
 
     if (error != ERROR_SUCCESS)
@@ -1792,6 +1884,21 @@ static PyObject* yara_compile(
       }
     }
 
+    if (include_callback != NULL)
+    {
+      if (!PyCallable_Check(include_callback))
+      {
+        return PyErr_Format(
+            PyExc_TypeError,
+            "'include_callback' must be callable");
+      }
+      Py_INCREF(include_callback);
+      yr_compiler_set_include_callback(compiler,
+                                       yara_include_callback,
+                                       yara_include_free,
+                                       include_callback);
+    }
+
     if (externals != NULL && externals != Py_None)
     {
       if (PyDict_Check(externals))
@@ -1961,7 +2068,7 @@ static PyObject* yara_compile(
 
     yr_compiler_destroy(compiler);
   }
-
+  Py_XDECREF(include_callback);
   return result;
 }