WIP

server/funq_server/runner_win.py

@@ -34,22 +34,108 @@
 
 from funq_server.runner import RunnerInjector
 from ctypes import windll, wintypes, byref
+import ctypes
 import time
 
+kernel32 = ctypes.WinDLL('kernel32.dll', use_last_error=True)
+
+PROCESS_VM_OPERATION = 0x0008
+PROCESS_VM_WRITE = 0x0020
+PROCESS_CREATE_THREAD = 0x0002
+MEM_COMMIT = 0x1000
+MEM_RESERVE = 0x2000
+MEM_RELEASE = 0x8000
+PAGE_READWRITE = 0x0004
+INFINITE = -1
+
+SIZE_T = ctypes.c_size_t
+LPSIZE_T = ctypes.POINTER(SIZE_T)
+WCHAR_SIZE = ctypes.sizeof(wintypes.WCHAR)
+LPSECURITY_ATTRIBUTES = wintypes.LPVOID
+LPTHREAD_START_ROUTINE = wintypes.LPVOID
+
+class BOOL_CHECKED(ctypes._SimpleCData):
+    _type_ = "l"
+    def _check_retval_(retval):
+        if retval == 0:
+            raise ctypes.WinError(ctypes.get_last_error())
+        return retval
+
+class LPVOID_CHECKED(ctypes._SimpleCData):
+    _type_ = "P"
+    def _check_retval_(retval):
+        if retval is None:
+            raise ctypes.WinError(ctypes.get_last_error())
+        return retval
+
+HANDLE_CHECKED = LPVOID_CHECKED  # not file handles
+
+kernel32.OpenProcess.restype = HANDLE_CHECKED
+kernel32.OpenProcess.argtypes = (
+    wintypes.DWORD, # dwDesiredAccess
+    wintypes.BOOL,  # bInheritHandle
+    wintypes.DWORD) # dwProcessId
+
+kernel32.VirtualAllocEx.restype = LPVOID_CHECKED
+kernel32.VirtualAllocEx.argtypes = (
+    wintypes.HANDLE, # hProcess
+    wintypes.LPVOID, # lpAddress
+    SIZE_T,          # dwSize
+    wintypes.DWORD,  # flAllocationType
+    wintypes.DWORD)  # flProtect
+
+kernel32.VirtualFreeEx.argtypes = (
+    wintypes.HANDLE, # hProcess
+    wintypes.LPVOID, # lpAddress
+    SIZE_T,          # dwSize
+    wintypes.DWORD)  # dwFreeType
+
+kernel32.WriteProcessMemory.restype = BOOL_CHECKED
+kernel32.WriteProcessMemory.argtypes = (
+    wintypes.HANDLE,  # hProcess
+    wintypes.LPVOID,  # lpBaseAddress
+    wintypes.LPCVOID, # lpBuffer
+    SIZE_T,           # nSize
+    LPSIZE_T)         # lpNumberOfBytesWritten _Out_
+
+kernel32.CreateRemoteThread.restype = HANDLE_CHECKED
+kernel32.CreateRemoteThread.argtypes = (
+    wintypes.HANDLE,        # hProcess
+    LPSECURITY_ATTRIBUTES,  # lpThreadAttributes
+    SIZE_T,                 # dwStackSize
+    LPTHREAD_START_ROUTINE, # lpStartAddress
+    wintypes.LPVOID,        # lpParameter
+    wintypes.DWORD,         # dwCreationFlags
+    wintypes.LPDWORD)       # lpThreadId _Out_
+
+kernel32.WaitForSingleObject.argtypes = (
+    wintypes.HANDLE, # hHandle
+    wintypes.DWORD)  # dwMilliseconds
+
+kernel32.GetExitCodeThread.restype = BOOL_CHECKED
+kernel32.GetExitCodeThread.argtypes = (
+    wintypes.HANDLE,   # hThread
+    wintypes.LPDWORD)  # lpExitCode
+
+kernel32.CloseHandle.argtypes = (
+    wintypes.HANDLE,) # hObject
+
+
 # Useful resources regarding DLL injection:
 #
 # - https://docs.microsoft.com/en-us/windows/win32/dlls/dynamic-link-library-best-practices  # noqa: E501
 # - https://blog.nettitude.com/uk/dll-injection-part-two
+# - https://stackoverflow.com/questions/17392721/error-invalid-parameter-error-57-when-calling-createremotethread-with-python-3-2/17524073#17524073
 # - https://stackoverflow.com/questions/27332509/createremotethread-on-loadlibrary-and-get-the-hmodule-back  # noqa: E501
 # - https://github.com/numaru/injector
 
 # Constants from Windows API documentation.
-PROCESS_CREATE_THREAD = 0x0002
-PROCESS_VM_OPERATION = 0x0008
-PROCESS_VM_WRITE = 0x0020
-MEM_COMMIT = 0x00001000
-MEM_RESERVE = 0x00002000
-PAGE_READWRITE = 0x04
+#PROCESS_CREATE_THREAD = 0x0002
+#PROCESS_VM_OPERATION = 0x0008
+#PROCESS_VM_WRITE = 0x0020
+#MEM_COMMIT = 0x00001000
+#MEM_RESERVE = 0x00002000
+#PAGE_READWRITE = 0x04
 
 
 class WindowsRunnerInjector(RunnerInjector):
@@ -84,71 +170,49 @@ def start_subprocess(self):
             raise
 
     def _inject_dll(self, pid, dll_path):
-        # Get handle to kernel32.dll.
-        kernel32_handle = windll.kernel32.GetModuleHandleA(b"kernel32.dll")
-        if not kernel32_handle:
-            self._raise_windows_error("GetModuleHandleA()", kernel32_handle)
-
-        # Get handle to LoadLibraryA().
-        loadlibrary_address = windll.kernel32.GetProcAddress(
-            kernel32_handle, b"LoadLibraryA")
-        if not loadlibrary_address:
-            self._raise_windows_error("GetProcAddress()", loadlibrary_address)
-
-        # Get handle to the running process.
-        process_handle = windll.kernel32.OpenProcess(
-            PROCESS_CREATE_THREAD | PROCESS_VM_OPERATION | PROCESS_VM_WRITE,
-            0, pid)
-        if not process_handle:
-            self._raise_windows_error("OpenProcess()", process_handle)
-
-        # Allocate memory for the DLL path.
-        dll_path = dll_path.encode("ascii")
-        path_address = windll.kernel32.VirtualAllocEx(
-            process_handle, 0, len(dll_path), MEM_COMMIT | MEM_RESERVE,
-            PAGE_READWRITE)
-        if not path_address:
-            self._raise_windows_error("VirtualAllocEx()", path_address)
-
-        # Write DLL path into the allocated memory region.
-        success = windll.kernel32.WriteProcessMemory(
-            process_handle, path_address, dll_path, len(dll_path), None)
-        if not success:
-            self._raise_windows_error("WriteProcessMemory()", success)
-
-        # Create and start new thread in the process. The entry point of the
-        # new thread is LoadLibraryA() with our DLL path as argument.
-        thread_handle = windll.kernel32.CreateRemoteThread(
-            process_handle, 0, 0, loadlibrary_address, path_address, 0, None)
-        if not thread_handle:
-            self._raise_windows_error("CreateRemoteThread()", thread_handle)
-
-        # Release process handle since we no longer need it.
-        success = windll.kernel32.CloseHandle(process_handle)
-        if not success:
-            self._raise_windows_error("CloseHandle()", success)
-
-        # Wait (with 10s timeout) until the thread exited, i.e. our DLL
-        # injection either succeeded or failed.
-        error = windll.kernel32.WaitForSingleObject(thread_handle, 10000)
-        if error:
-            self._raise_windows_error("WaitForSingleObject()", error)
-
-        # Get the exit code of our thread, which corresponds to the return
-        # value of LoadLibraryA() so we can check if the DLL was loaded
-        # successfully or not.
-        libfunq_handle = wintypes.DWORD(0)
-        success = windll.kernel32.GetExitCodeThread(
-            thread_handle, byref(libfunq_handle))
-        if not success:
-            self._raise_windows_error("GetExitCodeThread()", success)
-        if not libfunq_handle:
-            self._raise_windows_error("LoadLibraryA()", libfunq_handle)
-
-        # Release thread handle since we no longer need it.
-        success = windll.kernel32.CloseHandle(thread_handle)
-        if not success:
-            self._raise_windows_error("CloseHandle()", success)
+        # dll_path = dll_path.encode("ascii")
+        size = (len(dll_path) + 1) * WCHAR_SIZE
+        process_handle = None
+        path_address = None
+        thread_handle = None
+        try:
+            # Get handle to the running process.
+            process_handle = kernel32.OpenProcess(
+                PROCESS_CREATE_THREAD | PROCESS_VM_OPERATION |
+                PROCESS_VM_WRITE, False, pid)
+
+            # Allocate memory for the DLL path.
+            path_address = kernel32.VirtualAllocEx(
+                process_handle, None, size, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE)
+
+            # Write DLL path into the allocated memory region.
+            kernel32.WriteProcessMemory(
+                process_handle, path_address, dll_path, size, None)
+
+            # Create and start new thread in the process. The entry point of the
+            # new thread is LoadLibraryW() with our DLL path as argument.
+            thread_handle = kernel32.CreateRemoteThread(
+                process_handle, None, 0, kernel32.LoadLibraryW, path_address, 0, None)
+
+            # Wait (with 10s timeout) until the thread exited, i.e. our DLL
+            # injection either succeeded or failed.
+            kernel32.WaitForSingleObject(thread_handle, 10000)
+
+            # Get the exit code of our thread, which corresponds to the return
+            # value of LoadLibraryW() so we can check if the DLL was loaded
+            # successfully or not.
+            libfunq_handle = wintypes.DWORD(0)
+            kernel32.GetExitCodeThread(thread_handle, byref(libfunq_handle))
+            if not libfunq_handle:
+                self._raise_windows_error("LoadLibraryW()", libfunq_handle)
+        finally:
+            if path_address is not None:
+                kernel32.VirtualFreeEx(process_handle, path_address, 0, MEM_RELEASE)
+            if thread_handle is not None:
+                kernel32.CloseHandle(thread_handle)
+            if process_handle is not None:
+                kernel32.CloseHandle(process_handle)
+
 
     def _raise_windows_error(self, function_name, return_value):
         """