From 2fbc33b0f0e999522c41eda66372fc2935aa8228 Mon Sep 17 00:00:00 2001
From: rdbo <rdbodev@gmail.com>
Date: Sun, 22 Jan 2023 07:10:53 -0300
Subject: [PATCH] implemented VMT APIs for Python

---
 libmem-py/src/libmem-py/libmem-py.c |  11 +++
 libmem-py/src/libmem-py/types.h     | 122 ++++++++++++++++++++++++++++
 libmem-py/tests/tests.py            |  14 ++++
 3 files changed, 147 insertions(+)

diff --git a/libmem-py/src/libmem-py/libmem-py.c b/libmem-py/src/libmem-py/libmem-py.c
index 9bb3dba0..d40725d4 100644
--- a/libmem-py/src/libmem-py/libmem-py.c
+++ b/libmem-py/src/libmem-py/libmem-py.c
@@ -1435,6 +1435,9 @@ PyInit_libmem(void)
 	if (PyType_Ready(&py_lm_inst_t) < 0)
 		goto ERR_PYMOD;
 
+	if (PyType_Ready(&py_lm_vmt_t) < 0)
+		goto ERR_PYMOD;
+
 	pymod = PyModule_Create(&libmem_mod);
 	if (!pymod)
 		goto ERR_PYMOD;
@@ -1475,6 +1478,11 @@ PyInit_libmem(void)
 			       (PyObject *)&py_lm_inst_t) < 0)
 		goto ERR_INST;
 
+	Py_INCREF(&py_lm_vmt_t);
+	if (PyModule_AddObject(pymod, "lm_vmt_t",
+			       (PyObject *)&py_lm_vmt_t) < 0)
+		goto ERR_VMT;
+
 	/* global variables */
 	DECL_GLOBAL_PROT(LM_PROT_X);
 	DECL_GLOBAL_PROT(LM_PROT_R);
@@ -1487,6 +1495,9 @@ PyInit_libmem(void)
 
 	goto EXIT; /* no errors */
 
+ERR_VMT:
+	Py_DECREF(&py_lm_vmt_t);
+	Py_DECREF(pymod);
 ERR_INST:
 	Py_DECREF(&py_lm_inst_t);
 	Py_DECREF(pymod);
diff --git a/libmem-py/src/libmem-py/types.h b/libmem-py/src/libmem-py/types.h
index cf53bd3e..534a6350 100644
--- a/libmem-py/src/libmem-py/types.h
+++ b/libmem-py/src/libmem-py/types.h
@@ -384,3 +384,125 @@ static PyTypeObject py_lm_inst_t = {
 	.tp_repr = py_lm_inst_str
 };
 
+/****************************************/
+
+/* lm_process_t */
+typedef struct {
+	PyObject_HEAD
+	lm_vmt_t vmt;
+} py_lm_vmt_obj;
+
+PyObject *
+py_lm_vmt_str(PyObject *self)
+{
+	py_lm_vmt_obj *pyvmt = (py_lm_vmt_obj *)self;
+	return PyUnicode_FromFormat("<lm_vmt_t(vtable = 0x%zx)>", pyvmt->vmt.vtable);
+}
+
+int
+py_lm_vmt_init(PyObject *self,
+	       PyObject *args,
+	       PyObject *kwds)
+{
+	lm_address_t *vtable;
+	py_lm_vmt_obj *pyvmt = (py_lm_vmt_obj *)self;
+
+	if (!PyArg_ParseTuple(args, "k", &vtable))
+		return -1;
+
+	LM_VmtNew(vtable, &pyvmt->vmt);
+
+	return 0;
+}
+
+void
+py_lm_vmt_del(PyObject *self)
+{
+	py_lm_vmt_obj *pyvmt = (py_lm_vmt_obj *)self;
+	
+	LM_VmtFree(&pyvmt->vmt);
+}
+
+PyObject *
+py_lm_vmt_hook(PyObject *self,
+	       PyObject *args)
+{
+	py_lm_vmt_obj *pyvmt = (py_lm_vmt_obj *)self;
+	lm_size_t index;
+	lm_address_t dst;
+
+	if (!PyArg_ParseTuple(args, "kk", &index, &dst))
+		return NULL;
+
+	LM_VmtHook(&pyvmt->vmt, index, dst);
+
+	return Py_BuildValue("");
+}
+
+PyObject *
+py_lm_vmt_unhook(PyObject *self,
+		 PyObject *args)
+{
+	py_lm_vmt_obj *pyvmt = (py_lm_vmt_obj *)self;
+	lm_size_t index;
+
+	if (!PyArg_ParseTuple(args, "k", &index))
+		return NULL;
+
+	LM_VmtUnhook(&pyvmt->vmt, index);
+
+	return Py_BuildValue("");
+}
+
+PyObject *
+py_lm_vmt_get_original(PyObject *self,
+		       PyObject *args)
+{
+	py_lm_vmt_obj *pyvmt = (py_lm_vmt_obj *)self;
+	lm_size_t index;
+	lm_address_t orig_func;
+
+	if (!PyArg_ParseTuple(args, "k", &index))
+		return NULL;
+
+	orig_func = LM_VmtGetOriginal(&pyvmt->vmt, index);
+
+	return PyLong_FromSize_t(orig_func);
+}
+
+PyObject *
+py_lm_vmt_reset(PyObject *self,
+		PyObject *args)
+{
+	py_lm_vmt_obj *pyvmt = (py_lm_vmt_obj *)self;
+
+	LM_VmtReset(&pyvmt->vmt);
+
+	return Py_BuildValue("");
+}
+
+static PyMethodDef py_lm_vmt_methods[] = {
+	{ "hook", py_lm_vmt_hook, METH_VARARGS, "Hooks a VMT function at an index" },
+	{ "unhook", py_lm_vmt_unhook, METH_VARARGS, "Unhooks a VMT function at an index" },
+	{ "get_original", py_lm_vmt_get_original, METH_VARARGS, "Gets the original VMT function at an index" },
+	{ "reset", py_lm_vmt_reset, METH_VARARGS, "Resets the original VMT" },
+	{ NULL }
+};
+
+static PyTypeObject py_lm_vmt_t = {
+	PyVarObject_HEAD_INIT(NULL, 0)
+	.tp_name = "libmem.lm_vmt_t",
+	.tp_doc = "Manages a Virtual Method Table",
+	.tp_basicsize = sizeof(py_lm_vmt_obj),
+	.tp_itemsize = 0,
+	.tp_flags = Py_TPFLAGS_DEFAULT,
+	.tp_new = PyType_GenericNew,
+	.tp_init = py_lm_vmt_init,
+	.tp_finalize = py_lm_vmt_del,
+	.tp_methods = py_lm_vmt_methods,
+	.tp_str = py_lm_vmt_str,
+	.tp_repr = py_lm_vmt_str
+};
+
+/****************************************/
+
diff --git a/libmem-py/tests/tests.py b/libmem-py/tests/tests.py
index 041d9451..c6945613 100644
--- a/libmem-py/tests/tests.py
+++ b/libmem-py/tests/tests.py
@@ -226,4 +226,18 @@ def separator():
 separator()
 
 # TODO: Add tests for 'LM_CodeLengthEx'
+# separator()
+
+print("[*] VMT Hooking")
+
+vtable = ctypes.c_ulonglong(0x1020304050607080)
+vmt = lm_vmt_t(ctypes.addressof(vtable))
+print(f"[*] Original Function: {hex(vmt.get_original(0))}")
+vmt.hook(0, 0xdeadbeef)
+print(f"[*] VMT After Hook: {hex(vtable.value)}")
+vmt.unhook(0)
+# vmt.reset()
+print(f"[*] VMT After Unhook: {hex(vtable.value)}")
+
+separator()