test_hook.py 文件源码

def test_self_iat_hook_multithread():
    """Test IAT hook in current process with multi thread trigger"""
    cp = windows.current_process
    # Might change this to XP compat ?
    kernelbase_mod = [m for m in cp.peb.modules if m.name == "kernelbase.dll"][0]
    LdrLoadDll = [n for n in kernelbase_mod.pe.imports['ntdll.dll'] if n.name == "LdrLoadDll"][0]

    calling_thread = set([])
    @windows.hooks.LdrLoadDllCallback
    def MyHook(*args, **kwargs):
        calling_thread.add(windows.current_thread.tid)
        return kwargs["real_function"]()

    x = LdrLoadDll.set_hook(MyHook)
    # Trigger from local thread
    ctypes.WinDLL("kernel32.dll")
    assert calling_thread == set([windows.current_thread.tid])
    # Trigger from another thread
    k32 = [m for m in cp.peb.modules if m.name == "kernel32.dll"][0]
    load_libraryA = k32.pe.exports["LoadLibraryA"]
    with cp.allocated_memory(0x1000) as addr:
        cp.write_memory(addr, "DLLNOTFOUND.NOT_A_REAL_DLL" + "\x00")
        t = cp.create_thread(load_libraryA, addr)
        t.wait()
    assert len(calling_thread) == 2
    x.disable()