Remove generate_callback_stub (only IATHook are impacted)

TODO

@@ -12,19 +12,6 @@ TODO:
         - Test !! (bp, BP_HX, bp on only on process, bp_hx on only one thread..)
 		- test breakpoint with specific target
 
-    - Threading
-        - Quid IAT hook stub ? just einit threads and remove this ?
-        - Continue test with new generate_callback_stub and remove commented code if it works
-
-        - rewrite/REMOVE? generate_stub_64[32] : it's a non-sens to not save stuff on the stack..
-            I can re-copy the args on stack..
-
-    - Injection
-        - code generated by generate_python_exec_shellcode_64[32] may be reused
-            Just need to pass the address of the python string as argument
-
-
-
     - Winproxy:
         - rethink OptionalExport ? not useful with lazy resolution (or we need to force resolution..)
 

windows/hooks.py

@@ -43,6 +43,7 @@ for func in winfuncs.functions:
 
 class IATHook(object):
     """Look at my hook <3"""
+    yolo = []
 
     def __init__(self, IAT_entry, callback, types=None):
         if types is None:
@@ -53,9 +54,11 @@ class IATHook(object):
         self.callback_types = self.transform_arguments(self.original_types)
         self.entry = IAT_entry
         self.callback = callback
-        self.stub = native_exec.generate_callback_stub(self.hook_callback, self.callback_types)
+        self.stub = ctypes.WINFUNCTYPE(*self.callback_types)(self.hook_callback)
+        self.stub_addr = ctypes.cast(self.stub, PVOID).value
         self.realfunction = ctypes.WINFUNCTYPE(*types)(IAT_entry.nonhookvalue)
         self.is_enable = False
+        #IATHook.yolo.append(self)
 
     def transform_arguments(self, types):
         res = []
@@ -67,9 +70,9 @@ class IATHook(object):
         return res
 
     def enable(self):
-        """Enable the IAT hook"""
+        """Enable the IAT hook: you MUST keep a reference to the IATHook while the hook is enabled"""
         with utils.VirtualProtected(self.entry.addr, ctypes.sizeof(PVOID), PAGE_EXECUTE_READWRITE):
-            self.entry.value = self.stub
+            self.entry.value = self.stub_addr
         self.is_enable = True
 
     def disable(self):
@@ -93,3 +96,7 @@ class IATHook(object):
                 args = adapted_args
             return self.realfunction(*args)
         return self.callback(*adapted_args, real_function=real_function)
+
+    # Use this tricks to prevent garbage collection of hook ?
+    #def __del__(self):
+    #    pass

windows/native_exec/__init__.py

@@ -1,3 +1,3 @@
-from .native_function import generate_callback_stub, create_function
+from .native_function import create_function
 
-__all__ = ["generate_callback_stub", "create_function"]
+__all__ = ["create_function"]

windows/native_exec/native_function.py

@@ -268,22 +268,16 @@ allocator = CustomAllocator()
 #    return code
 
 
-def generate_callback_stub(callback, types):
-    func_type = ctypes.WINFUNCTYPE(*types)
-    c_callable = func_type(callback)
-
-
-    stub = c_callable
-    stub_addr = ctypes.cast(c_callable, ctypes.c_void_p).value
-    #    if windows.current_process.bitness == 32:
-    #        stub = generate_stub_32(c_callable)
-    #    else:
-    #        stub = generate_stub_64(c_callable)
-    #    stub_addr = allocator.write_code(stub.get_code())
-    generate_callback_stub.l.append((stub, c_callable))
-    return stub_addr
-
-generate_callback_stub.l = []
+#def generate_callback_stub(callback, types):
+#    func_type = ctypes.WINFUNCTYPE(*types)
+#    c_callable = func_type(callback)
+#
+#    stub = c_callable
+#    stub_addr = ctypes.cast(c_callable, ctypes.c_void_p).value
+#    generate_callback_stub.l.append((stub, c_callable))
+#    return stub_addr
+#
+#generate_callback_stub.l = []
 
 
 def create_function(code, types):

windows/pe_parse.py

@@ -150,6 +150,7 @@ def GetPEFile(baseaddr, target=None, force_bitness=None):
 
         def set_hook(self, callback, types=None):
             """Setup a hook on the entry and return it.
+            You MUST keep a reference to the hook while the hook is enabled.
 
             :param callback: the hook
 

windows/test/test_hooks.py

@@ -20,7 +20,7 @@ class HookTestCase(unittest.TestCase):
             phkResult[0] = 12345678
             return 0
 
-        RegOpenKeyExA.set_hook(open_reg_hook)
+        x = RegOpenKeyExA.set_hook(open_reg_hook)
         import _winreg
         open_args = (0x12345678, "MY_KEY_VALUE")
         k = _winreg.OpenKey(*open_args)