Add console_callback to match call. (VirusTotal#194)

wxsBSD · web-flow · commit 4379581cd67b · 2022-03-04T10:58:10.000+01:00
* Add console_callback to match call. This provides the user with an interface to handle console.log() messages. If not provided the log message is printed to stdout (limited to 1000 bytes based upon https://docs.python.org/3/c-api/sys.html?highlight=stdout#c.PySys_WriteStdout). Tested with the following (along with the updated tests): wxs@wxs-mbp yara-python % cat test.py import yara r = """ import "console" rule a { condition: console.log("Hello from Python!") } """ def console(message): print(f"Callback: {message}") rules = yara.compile(source=r) rules.match("/bin/ls", console_callback=console) rules.match("/bin/ls") wxs@wxs-mbp yara-python % PYTHONPATH=build/lib.macosx-10.14-arm64-3.8 python3 test.py Callback: Hello from Python! Hello from Python! wxs@wxs-mbp yara-python % * Add comment.
diff --git a/tests.py b/tests.py
@@ -940,6 +940,19 @@ def callback(requested_filename, filename, namespace):
         r = yara.compile(source='include "foo" rule r { condition: included }', include_callback=callback)
         self.assertTrue(r.match(data='dummy'))
 
+    def testConsoleCallback(self):
+        global called
+        called = False
+
+        def callback(message):
+            global called
+            called = True
+            return yara.CALLBACK_CONTINUE
+
+        r = yara.compile(source='import "console" rule r { condition: console.log("AXSERS") }')
+        r.match(data='dummy', console_callback=callback)
+        self.assertTrue(called)
+
     def testCompare(self):
 
         r = yara.compile(sources={
diff --git a/yara-python.c b/yara-python.c
@@ -418,6 +418,7 @@ typedef struct _CALLBACK_DATA
   PyObject* modules_data;
   PyObject* modules_callback;
   PyObject* warnings_callback;
+  PyObject* console_callback;
   int which;
 
 } CALLBACK_DATA;
@@ -691,6 +692,58 @@ static int handle_module_imported(
 }
 
 
+static int handle_console_log(
+    void* message_data,
+    CALLBACK_DATA* data)
+{
+  PyGILState_STATE gil_state = PyGILState_Ensure();
+  int result = CALLBACK_CONTINUE;
+
+  if (data->console_callback == NULL)
+  {
+    // If the user does not specify a console callback we dump to stdout.
+    // If we want to support 3.2 and newer only we can use
+    // https://docs.python.org/3/c-api/sys.html?highlight=stdout#c.PySys_FormatStdout
+    // instead of this call with the limit.
+    PySys_WriteStdout("%.1000s\n", (char*) message_data);
+  }
+  else
+  {
+    PyObject* log_string = PY_STRING((char*) message_data);
+    Py_INCREF(data->console_callback);
+
+    PyObject* callback_result = PyObject_CallFunctionObjArgs(
+        data->console_callback,
+        log_string,
+        NULL);
+
+    if (callback_result != NULL)
+    {
+      #if PY_MAJOR_VERSION >= 3
+      if (PyLong_Check(callback_result))
+      #else
+      if (PyLong_Check(callback_result) || PyInt_Check(callback_result))
+      #endif
+      {
+        result = (int) PyLong_AsLong(callback_result);
+      }
+    }
+    else
+    {
+      result = CALLBACK_ERROR;
+    }
+
+    Py_DECREF(log_string);
+    Py_XDECREF(callback_result);
+    Py_DECREF(data->console_callback);
+  }
+
+  PyGILState_Release(gil_state);
+
+  return result;
+}
+
+
 static int handle_too_many_matches(
     YR_SCAN_CONTEXT* context,
     YR_STRING* string,
@@ -870,6 +923,10 @@ int yara_callback(
     if (callback == NULL ||
         (which & CALLBACK_NON_MATCHES) != CALLBACK_NON_MATCHES)
       return CALLBACK_CONTINUE;
+    break;
+
+  case CALLBACK_MSG_CONSOLE_LOG:
+    return handle_console_log(message_data, user_data);
   }
 
   // At this point we have handled all the other cases of when this callback
@@ -1551,7 +1608,8 @@ static PyObject* Rules_match(
   static char* kwlist[] = {
       "filepath", "pid", "data", "externals",
       "callback", "fast", "timeout", "modules_data",
-      "modules_callback", "which_callbacks", "warnings_callback", NULL
+      "modules_callback", "which_callbacks", "warnings_callback",
+      "console_callback", NULL
       };
 
   char* filepath = NULL;
@@ -1574,12 +1632,13 @@ static PyObject* Rules_match(
   callback_data.modules_data = NULL;
   callback_data.modules_callback = NULL;
   callback_data.warnings_callback = NULL;
+  callback_data.console_callback = NULL;
   callback_data.which = CALLBACK_ALL;
 
   if (PyArg_ParseTupleAndKeywords(
         args,
         keywords,
-        "|sis*OOOiOOiO",
+        "|sis*OOOiOOiOO",
         kwlist,
         &filepath,
         &pid,
@@ -1591,7 +1650,8 @@ static PyObject* Rules_match(
         &callback_data.modules_data,
         &callback_data.modules_callback,
         &callback_data.which,
-        &callback_data.warnings_callback))
+        &callback_data.warnings_callback,
+        &callback_data.console_callback))
   {
     if (filepath == NULL && data.buf == NULL && pid == -1)
     {
@@ -1633,6 +1693,17 @@ static PyObject* Rules_match(
       }
     }
 
+    if (callback_data.console_callback != NULL)
+    {
+      if (!PyCallable_Check(callback_data.console_callback))
+      {
+        PyBuffer_Release(&data);
+        return PyErr_Format(
+            PyExc_TypeError,
+            "'console_callback' must be callable");
+      }
+    }
+
     if (callback_data.modules_data != NULL)
     {
       if (!PyDict_Check(callback_data.modules_data))
