Unverified Commit 4379581c by Wesley Shields Committed by GitHub

Add console_callback to match call. (#194)

* 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.
parent 08372e19
......@@ -940,6 +940,19 @@ class TestYara(unittest.TestCase):
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={
......
......@@ -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))
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment