From f43749aa07e5ca036a1067064501cc1884c21f83 Mon Sep 17 00:00:00 2001
From: Clement Rouault <project@hakril.net>
Date: Fri, 4 Aug 2017 17:19:27 +0200
Subject: [PATCH] [WIP] new hook method + winfunc APIs in winproxy

---
 windows/hooks.py          | 24 ++++++++++++++++++++++++
 windows/utils/winutils.py |  1 +
 windows/winproxy.py       | 25 +++++++++++++++++++++++++
 3 files changed, 50 insertions(+)

diff --git a/windows/hooks.py b/windows/hooks.py
index 9630c23..80cf523 100644
--- a/windows/hooks.py
+++ b/windows/hooks.py
@@ -1,6 +1,7 @@
 import sys
 import ctypes
 
+import windows
 import windows.utils as utils
 from . import native_exec
 from .generated_def import winfuncs
@@ -100,3 +101,26 @@ class IATHook(object):
     # Use this tricks to prevent garbage collection of hook ?
     #def __del__(self):
     #    pass
+
+
+## New simple hook API based on winproxy
+
+def setup_hook(target, hook, dll_to_hook):
+    "TODO: Test and doc :D"
+    dll_to_hook = "python27.dll"
+    dll_name, api_name = windows.winproxy.get_target(target)
+    prototype = target.prototype
+    hook._types_info = (prototype._restype_,) + prototype._argtypes_
+
+    if not dll_name.endswith(".dll"):
+        dll_name += ".dll"
+    # Get the peb of our process
+    peb = windows.current_process.peb
+    # Get the dll_to_hook
+    module_to_hook = [m for m in peb.modules if m.name.lower() == dll_to_hook.lower()][0]
+    # Get the iat entries for DLL dll_name
+    adv_imports = module_to_hook.pe.imports[dll_name]
+    # Get RegOpenKeyExA iat entry
+    iat = [n for n in adv_imports if n.name == api_name][0]
+    iat.set_hook(hook)
+    return iat
\ No newline at end of file
diff --git a/windows/utils/winutils.py b/windows/utils/winutils.py
index 580405d..ad79f03 100644
--- a/windows/utils/winutils.py
+++ b/windows/utils/winutils.py
@@ -13,6 +13,7 @@ from ..generated_def.winstructs import *
 
 
 # Function resolution !
+# should be in winproxy ?
 def get_func_addr(dll_name, func_name):
         # Load the DLL
         ctypes.WinDLL(dll_name)
diff --git a/windows/winproxy.py b/windows/winproxy.py
index 8d22e07..7433057 100644
--- a/windows/winproxy.py
+++ b/windows/winproxy.py
@@ -1,6 +1,8 @@
 import ctypes
 import functools
 
+import windows
+
 from ctypes.wintypes import *
 from windows.generated_def.winstructs import *
 from windows.generated_def.windef import *
@@ -8,6 +10,7 @@ import windows.generated_def.winfuncs as winfuncs
 from windows.generated_def.ntstatus import NtStatusException
 from windows.dbgprint import dbgprint
 
+## winfuncs helpers
 
 def is_implemented(winfunc):
     try:
@@ -16,6 +19,28 @@ def is_implemented(winfunc):
         return False
     return True
 
+
+def get_target(winfunc):
+    """POC for new hook"""
+    return winfunc.target_dll, winfunc.target_func
+
+
+# Rip-of windows.utils: should be removed from the other place ?
+def _get_func_addr(dll_name, func_name):
+        # Load the DLL
+        ctypes.WinDLL(dll_name)
+        modules = windows.current_process.peb.modules
+        if not dll_name.lower().endswith(".dll"):
+            dll_name += ".dll"
+        mod = [x for x in modules if x.name == dll_name][0]
+        return mod.pe.exports[func_name]
+
+
+def resolve(winfunc):
+    winfunc.force_resolution()
+    return _get_func_addr(*get_target(winfunc))
+
+
 class Kernel32Error(WindowsError):
     def __new__(cls, func_name):
         win_error = ctypes.WinError()