# Source code for windows.hooks

    import sys
    import ctypes

    import windows
    import windows.utils as utils
    from . import native_exec
    from .generated_def import winfuncs
    from .generated_def.windef import PAGE_EXECUTE_READWRITE
    from .generated_def.winstructs import *

    # TODO Not a big fan of importing 'meta' every load
    # Should do an Hook API that take the winproxy function (not generate every hook possible)
    import windows.generated_def.meta

    class Callback(object):
        """Give type information to hook callback"""
        def __init__(self, *types):
            self.types = types

        def __call__(self, func):
            func._types_info = self.types
            return func

    class KnownCallback(object):
        types = ()

        def __call__(self, func):
            func._types_info = self.types
            return func

    def add_callback_to_module(callback):
        setattr(sys.modules[__name__], type(callback).__name__, callback)

    # Generate IATCallback decorator for all known functions

    for func in windows.generated_def.meta.functions:
        prototype = getattr(winfuncs, func + "Prototype")
        callback_name = func + "Callback"

        class CallBackDeclaration(KnownCallback):
            types = (prototype._restype_,) + prototype._argtypes_

        CallBackDeclaration.__name__ = callback_name
        add_callback_to_module(CallBackDeclaration())

    class IATHook(object):
        """Look at my hook <3"""

        def __init__(self, IAT_entry, callback, types=None):
            if types is None:
                if not hasattr(callback, "_types_info"):
                    raise ValueError("Callback for IATHook has no type infomations")
                types = callback._types_info
            self.original_types = types
            self.callback_types = self.transform_arguments(self.original_types)
            self.entry = IAT_entry
            self.callback = callback
            ## No more circular ref -> but stub is destroyed -> segv :(
            self.stub = ctypes.WINFUNCTYPE(*self.callback_types)(self.hook_callback)
            # stub = ctypes.WINFUNCTYPE(*self.callback_types)(self.hook_callback)
            # self.stub_addr = ctypes.cast(stub, PVOID) # Same problem as keep stub... (GC..)
            self.stub_addr = ctypes.cast(self.stub, PVOID).value # Same problem as keep stub... (GC..)
            # IAT_entry.stub = stub
            self.realfunction = ctypes.WINFUNCTYPE(*types)(IAT_entry.nonhookvalue)
            self.is_enable = False

        def transform_arguments(self, types):
            res = []
            for type in types:
                if type in (ctypes.c_wchar_p, ctypes.c_char_p):
                    res.append(ctypes.c_void_p)
                else:
                    res.append(type)
            return res

        def enable(self):
            """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_addr
            self.is_enable = True
            self.entry.enabled = True

        def disable(self):
            """Disable the IAT hook"""
            with utils.VirtualProtected(self.entry.addr, ctypes.sizeof(PVOID), PAGE_EXECUTE_READWRITE):
                self.entry.value = self.entry.nonhookvalue
            self.is_enable = False
            self.entry.enabled = True

        def hook_callback(self, *args):
            adapted_args = []
            for value, type in zip(args, self.original_types[1:]):
                if type == ctypes.c_wchar_p:
                    adapted_args.append(ctypes.c_wchar_p(value))
                elif type == ctypes.c_char_p:
                    adapted_args.append(ctypes.c_char_p((value)))
                else:
                    adapted_args.append(value)

            def real_function(*args):
                if args == ():
                    args = adapted_args
                return self.realfunction(*args)
            return self.callback(*adapted_args, real_function=real_function)

    ## New simple hook API based on winproxy
    def setup_hook(target, hook, dll_to_hook):
        "TODO: Test and doc :D"
        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