Merge pull request #106 from lostsnow/feature/concat-hook

Feature/concat hook

Author: jinghao1

dongtai_agent_python/assess/c_api_hook.py

@@ -9,6 +9,9 @@
     'builtins.bytes.__new__': 'callback_bytes_cast',
     'builtins.bytearray.__init__': 'callback_bytearray_cast',
     'builtins.str.__new__': 'callback_unicode_cast',
+    'builtins.bytes.__add__': 'callback_bytes_concat',
+    'builtins.bytearray.__add__': 'callback_bytearray_concat',
+    'builtins.str.__add__': 'callback_unicode_concat',
 }
 
 

dongtai_agent_python/assess_ext/include/patch.h

@@ -21,6 +21,7 @@ PyObject *str_origin(PyObject *self, PyObject *args);
 int apply_cformat_patch(funchook_t *funchook);
 int apply_fstring_patch(funchook_t *funchook);
 int apply_cast_patch(funchook_t *funchook);
+int apply_concat_patch(funchook_t *funchook);
 
 #define BUILD_NEW_BINARYFUNC(NAME)                                            \
     static PyObject *NAME##_new(PyObject *self, PyObject *args) {             \

dongtai_agent_python/assess_ext/patch.c

@@ -69,6 +69,7 @@ PyObject *enable_patches(PyObject *self, PyObject *arg) {
     apply_patch(apply_cformat_patch, funchook);
     apply_patch(apply_fstring_patch, funchook);
     apply_patch(apply_cast_patch, funchook);
+    apply_patch(apply_concat_patch, funchook);
 
     Py_RETURN_NONE;
 }
@@ -90,6 +91,10 @@ PyObject *install(PyObject *self, PyObject *arg) {
 }
 
 void patch_string_callback(char *prop_method_name, PyObject *source, PyObject *target, PyObject *hook_args, PyObject *hook_kwargs) {
+    if (!PyObject_HasAttrString(patch_module, prop_method_name)) {
+        return;
+    }
+
     PyObject *result;
     PyObject *prop_hook_args;
     int free_hook_args = 0;

dongtai_agent_python/assess_ext/patch/concat.c

@@ -0,0 +1,96 @@
+#define PY_SSIZE_T_CLEAN
+#include <Python.h>
+
+#include <funchook.h>
+#include <logger.h>
+#include <utils.h>
+#include <patch.h>
+
+binaryfunc bytes_concat_origin;
+binaryfunc bytearray_concat_origin;
+binaryfunc bytearray_inplace_concat_origin;
+binaryfunc unicode_concat_origin;
+void (*unicode_append_origin)(PyObject **l, PyObject *r);
+
+PyObject *bytes_concat_new(PyObject *l, PyObject *r) {
+    PyObject *result = bytes_concat_origin(l, r);
+
+    if (result == NULL) {
+        return result;
+    }
+
+    patch_string_callback("callback_bytes_concat", l, result, r, NULL);
+
+    return result;
+}
+
+PyObject *bytearray_concat_new(PyObject *l, PyObject *r) {
+    PyObject *result = bytearray_concat_origin(l, r);
+
+    if (result == NULL) {
+        return result;
+    }
+
+    patch_string_callback("callback_bytearray_concat", l, result, r, NULL);
+
+    return result;
+}
+
+PyObject *bytearray_inplace_concat_new(PyObject *l, PyObject *r) {
+    PyObject *result = bytearray_inplace_concat_origin(l, r);
+
+    if (result == NULL) {
+        return result;
+    }
+
+    patch_string_callback("callback_bytearray_concat", l, result, r, NULL);
+
+    return result;
+}
+
+PyObject *unicode_concat_new(PyObject *l, PyObject *r) {
+    PyObject *result = unicode_concat_origin(l, r);
+
+    if (result == NULL) {
+        return result;
+    }
+
+    patch_string_callback("callback_unicode_concat", l, result, r, NULL);
+
+    return result;
+}
+
+void unicode_append_new(PyObject **l, PyObject *r) {
+    PyObject *origin_l = *l;
+    Py_XINCREF(origin_l);
+    unicode_append_origin(l, r);
+
+    if (*l == NULL) {
+        Py_XDECREF(origin_l);
+        return;
+    }
+
+    patch_string_callback("callback_unicode_concat", origin_l, *l, r, NULL);
+    Py_XDECREF(origin_l);
+}
+
+int apply_concat_patch(funchook_t *funchook) {
+    bytes_concat_origin = PyBytes_Type.tp_as_sequence->sq_concat;
+    funchook_prepare_wrapper(funchook, &bytes_concat_origin, bytes_concat_new);
+
+    bytearray_concat_origin = PyByteArray_Concat;
+    funchook_prepare_wrapper(funchook, &bytearray_concat_origin, bytearray_concat_new);
+
+    bytearray_inplace_concat_origin = PyByteArray_Type.tp_as_sequence->sq_inplace_concat;
+    funchook_prepare_wrapper(funchook, &bytearray_inplace_concat_origin, bytearray_inplace_concat_new);
+
+    unicode_concat_origin = PyUnicode_Concat;
+    funchook_prepare_wrapper(funchook, &unicode_concat_origin, unicode_concat_new);
+
+    unicode_append_origin = PyUnicode_Append;
+    funchook_prepare_wrapper(funchook, &unicode_append_origin, unicode_append_new);
+
+    log_debug("------c_patch------------------ concat");
+
+    return 0;
+}

dongtai_agent_python/policy_api.json

@@ -543,7 +543,7 @@
           "source": "P",
           "track": "true",
           "target": "R",
-          "value": "builtins.str.__format__",
+          "value": "builtins.str.__add__",
           "inherit": "false"
         },
         {
@@ -581,13 +581,6 @@
           "value": "builtins.str.cformat",
           "inherit": "false"
         },
-        {
-          "source": "P",
-          "track": "true",
-          "target": "R",
-          "value": "builtins.str.concat",
-          "inherit": "false"
-        },
         {
           "source": "P1",
           "track": "true",
@@ -756,6 +749,13 @@
           "value": "builtins.str.zfill",
           "inherit": "false"
         },
+        {
+          "source": "P",
+          "track": "true",
+          "target": "R",
+          "value": "builtins.bytes.__add__",
+          "inherit": "false"
+        },
         {
           "source": "P1,source",
           "track": "true",
@@ -777,6 +777,13 @@
           "value": "builtins.bytes.decode",
           "inherit": "false"
         },
+        {
+          "source": "P",
+          "track": "true",
+          "target": "R",
+          "value": "builtins.bytearray.__add__",
+          "inherit": "false"
+        },
         {
           "source": "P1,source",
           "track": "true",

dongtai_agent_python/setting/const.py

@@ -28,6 +28,9 @@
     'builtins.str.__new__',
     'builtins.bytes.__new__',
     'builtins.bytearray.__init__',
+    'builtins.str.__add__',
+    'builtins.bytes.__add__',
+    'builtins.bytearray.__add__',
 ]
 
 CRYPTO_BAD_CIPHER_NEW = [