c_api_patch.py

import ctypes
import os

from dongtai_agent_python.utils import scope

class CAPIPatch(object):
    def __init__(self):
        base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        self.fh_lib = ctypes.cdll.LoadLibrary(os.path.join(base_dir, 'assess_ext/libfunchook.so.1'))

        # define signatures
        self.funchook_create = self.fh_lib.funchook_create
        self.funchook_create.restype = ctypes.c_void_p
        self.funchook_create.argtypes = []

        self.funchook_prepare = self.fh_lib.funchook_prepare
        self.funchook_prepare.restype = ctypes.c_ssize_t
        self.funchook_prepare.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p]

        self.funchook_install = self.fh_lib.funchook_install
        self.funchook_install.restype = ctypes.c_ssize_t
        self.funchook_install.argtypes = [ctypes.c_void_p, ctypes.c_int]

    def create(self):
        self.funchook = self.funchook_create()

    def install(self):
        ret = self.funchook_install(self.funchook, 0)
        assert not ret, 'ret is ' + str(ret)

    def cformat_patch(self):
        PyUnicode_Format = ctypes.pythonapi.PyUnicode_Format
        PyUnicode_Format.restype = ctypes.py_object
        PyUnicode_Format.argtypes = [ctypes.py_object, ctypes.py_object]

        # must keep those references alive, or stuff will be GC'd and weird errors will occur

        # create hook (this function will replace the original function)
        hook_type = ctypes.PYFUNCTYPE(ctypes.py_object, ctypes.py_object, ctypes.py_object)

        self.cformat_hook = hook_type(self.cformat_impl)

        # create a pointer object with the function address
        self.cformat_orig_ptr = ctypes.c_void_p(ctypes.c_void_p.from_address(ctypes.addressof(PyUnicode_Format)).value)
        # orig_write_ptr.value will get a ptr to the original PySys_WriteStdout
        # and PySys_WriteStdout will now point to the hook
        ret = self.funchook_prepare(self.funchook, ctypes.addressof(self.cformat_orig_ptr), self.cformat_hook)
        assert not ret, 'ret is ' + str(ret)

        self.cformat_orig = hook_type.from_address(ctypes.addressof(self.cformat_orig_ptr))
        PyUnicode_Format('aaa %d', 1111)

    def cformat_impl(self, msg, args):
        with scope.scope(scope.SCOPE_AGENT):
            print('hook cformat: ' + str(msg) + ' ' + str(args))
        return self.cformat_orig(msg, args)