Hooking IAT in remote process with Python?

Question

I'm trying to hook a win32 function call (CreateFileW) inside of a notepad process to have the function do additional actions before returning what it should do. Ultimately, this will assist me in some RE translation efforts for another project I'm working on, but I'm just doing this as a proof of concept.

I'm the most comfortable in Python, so I've been trying to accomplish this task in this language.

There are a lot of small to large scale Python libraries out there on how to RPM/WPM, but hooking/detouring are not focus points with these projects.

Here is an example of what I've attempted to do with both the pymem and PythonForWindows libraries to accomplish this task:

import pymem
import sys
import os
from json import dumps

wdir = os.path.abspath('.')

log_path = os.path.join(wdir, 'out.log').replace("\\", "\\\\")
err_path = os.path.join(wdir, 'err.log').replace("\\", "\\\\")

shellcode = """
import sys
from os import chdir
from traceback import format_exc
sys.path=%s
chdir(sys.path[0])

def write_file(message):
    with open("%s", "a") as f:
        f.write(str(message))

try:
    import windows
    from windows.hooks import *

    @CreateFileWCallback
    def createfile_callback(lpFileName, dwDesiredAccess, dwShareMode, lpSecurityAttributes, dwCreationDisposition, dwFlagsAndAttributes, hTemplateFile, real_function):
        try:
            write_file("Trying to open {0}".format(lpFileName))
            return real_function()
        except:
            write_file(format_exc())

    peb = windows.current_process.peb.modules[0]
    imp = peb.pe.imports
    
    iat_create_file = [entry for entry in imp['kernel32.dll'] if entry.name == "CreateFileW"]
    
    result = iat_create_file[0].set_hook(createfile_callback)
except:
    write_file(format_exc())
""" % (
    dumps(sys.path).replace("\\", "\\\\"),
    'err.log'
)

pm = pymem.Pymem('notepad.exe')
pm.inject_python_interpreter()
pm.inject_python_shellcode(shellcode)

This code is supposed to hook the CreateFileW function, log the name of the file it's interacting with to a file and then return the rest of the function. My problem is that when this function is called within notepad, the program just hangs -- so not successful unfortunately.

I have a feeling that my code here:

pm.inject_python_shellcode(shellcode)

is spawning a thread and executing immediately, instead of only executing when it has been called - but I don't know of a different way to do this.

I'm totally open to different ways to go about this. I know ctypes is powerful for this work, but I'm unsure how to implement something like that in this method. I learn best by example code and have easily dumped dozens of hours searching for ways to do this without success. Thanks for any help you can provide!

Answer 1

I cross-posted this to another community that I'm a part of and ended up getting my answer (I should have done this sooner..):

It was likely that my code was getting garbage collected, hence the lock up. To keep the code there, I added:

while True:
  pass

to the end of my try block. The text is now being logged.