Added the possibility to access pre-registers in a post-hook

addons/source-python/packages/source-python/memory/hooks.py

@@ -26,7 +26,8 @@
            'PreHook',
            'set_hooks_disabled',
            'get_hooks_disabled',
-           'hooks_disabled'
+           'hooks_disabled',
+           'use_pre_registers',
            )
 
 
@@ -86,6 +87,56 @@ class PostHook(_Hook):
 # =============================================================================
 # >> FUNCTIONS
 # =============================================================================
+@contextmanager
+def use_pre_registers(stack_data, value=True):
+    """Temporarily set ``StackData.use_pre_registers`` to the given value.
+    When the context ends, the previous value is restored.
+
+    Some functions overwrite CPU registers during their execution with values
+    from their internal calculations. In a post-hook, you have access to the
+    modified CPU registers, but in some cases you might want to access the
+    registers that were saved before the pre-hook was called. In that case you
+    can use this context manager to get access to the previous state of the
+    registers. On Windows this is often required when hooking THISCALL
+    functions, because the this-pointer is saved in the CPU register ``ECX``,
+    but gets overwritten during the execution of the hooked function. So, in a
+    post-hook you won't have access to the this-pointer anymore.
+
+    Example (CS:S/Windows):
+
+    .. code:: python
+
+        from entities.hooks import EntityCondition
+        from entities.hooks import EntityPostHook
+        from entities.hooks import EntityPreHook
+
+        from memory.hooks import use_pre_registers
+
+        @EntityPreHook(EntityCondition.is_player, 'drop_weapon')
+        def pre_on_drop_weapon(stack_data):
+            print(f'PRE: this = {stack_data[0].address}')
+
+        @EntityPostHook(EntityCondition.is_player, 'drop_weapon')
+        def post_on_drop_weapon(stack_data, ret_val):
+            print(f'POST FALSE: this = {stack_data[0].address}')
+            with use_pre_registers(stack_data):
+                print(f'POST CORRECT: this = {stack_data[0].address}')
+
+    Output:
+
+    .. code::
+
+        PRE: this = 546778280
+        POST FALSE: this = 16439007
+        POST CORRECT: this = 546778280
+    """
+    old = stack_data.use_pre_registers
+    stack_data.use_pre_registers = value
+    try:
+        yield
+    finally:
+        stack_data.use_pre_registers = old
+
 @contextmanager
 def hooks_disabled(disabled=True):
     """Temporarily disable or enable all hook callbacks. By default hooks are

src/core/modules/memory/memory_hooks.cpp

@@ -173,9 +173,10 @@ object CStackData::GetItem(unsigned int iIndex)
 		BOOST_RAISE_EXCEPTION(PyExc_IndexError, "Index out of range.")
 
 	// Argument already cached?
-	object retval = m_mapCache[iIndex];
-	if (retval)
-		return retval;
+	object retval;
+	//object retval = m_mapCache[iIndex];
+	//if (retval)
+	//	return retval;
 
 	switch(m_pHook->m_pCallingConvention->m_vecArgTypes[iIndex])
 	{
@@ -196,7 +197,7 @@ object CStackData::GetItem(unsigned int iIndex)
 		case DATA_TYPE_STRING:		retval = GetArgument<const char *>(m_pHook, iIndex); break;
 		default: BOOST_RAISE_EXCEPTION(PyExc_TypeError, "Unknown type.") break;
 	}
-	m_mapCache[iIndex] = retval;
+	//m_mapCache[iIndex] = retval;
 	return retval;
 }
 
@@ -206,7 +207,7 @@ void CStackData::SetItem(unsigned int iIndex, object value)
 		BOOST_RAISE_EXCEPTION(PyExc_IndexError, "Index out of range.")
 
 	// Update cache
-	m_mapCache[iIndex] = value;
+	//m_mapCache[iIndex] = value;
 	switch(m_pHook->m_pCallingConvention->m_vecArgTypes[iIndex])
 	{
 		case DATA_TYPE_BOOL:		SetArgument<bool>(m_pHook, iIndex, value); break;

src/core/modules/memory/memory_hooks.h

@@ -51,20 +51,30 @@ class CStackData
 	void		SetItem(unsigned int iIndex, object value);
 
 	CRegisters* GetRegisters()
-	{ return m_pHook->m_pRegisters; }
+	{ return m_pHook->GetRegisters(); }
 
 	str	__repr__()
 	{ return str(tuple(ptr(this))); }
 
 	void* GetReturnAddress()
 	{
-		void* pESP = m_pHook->m_pRegisters->m_esp->GetValue<void*>();
+		void* pESP = m_pHook->GetRegisters()->m_esp->GetValue<void*>();
 		if (m_pHook->m_RetAddr.count(pESP) == 0) {
 			return NULL;
 		}
 		return m_pHook->m_RetAddr[pESP].back();
 	}
 
+	bool GetUsePreRegister()
+	{
+		return m_pHook->m_bUsePreRegisters;
+	}
+
+	void SetUsePreRegisters(bool value)
+	{
+		m_pHook->m_bUsePreRegisters = value;
+	}
+
 protected:
 	CHook*                m_pHook;
 	std::map<int, object> m_mapCache;

src/core/modules/memory/memory_wrap.cpp

@@ -635,6 +635,11 @@ void export_stack_data(scope _memory)
 		.add_property("registers",
 			make_function(&CStackData::GetRegisters, reference_existing_object_policy())
 		)
+
+		.add_property(
+			"use_pre_registers",
+			&CStackData::GetUsePreRegister,
+			&CStackData::SetUsePreRegisters)
 	;
 }
 

src/thirdparty/DynamicHooks/include/conventions/x86MsCdecl.h

@@ -77,6 +77,7 @@ class x86MsCdecl: public ICallingConvention
 
 private:
 	void* m_pReturnBuffer;
+	int* m_pOffsets;
 };
 
 #endif // _X86_MS_CDECL_H

src/thirdparty/DynamicHooks/include/conventions/x86MsFastcall.h

@@ -78,6 +78,7 @@ class x86MsFastcall: public ICallingConvention
 
 private:
 	void* m_pReturnBuffer;
+	int* m_pOffsets;
 };
 
 #endif // _X86_MS_FASTCALL_H