GeezerVison is showing github.com — PythonMemoryModule/pythonmemorymodule/windows/hooks.py at main · freeide/PythonMemoryModule

Expand file tree

hooks.py

Copy path

More file actions

Latest commit

History

130 lines (104 loc) · 4.47 KB

hooks.py

Copy path

File metadata and controls

130 lines (104 loc) · 4.47 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

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