From 5a29c59b3e6cfe6c1b7efb971987606ac7888741 Mon Sep 17 00:00:00 2001
From: Clement Rouault <clement.rouault@sogeti.com>
Date: Tue, 22 Dec 2015 16:17:30 +0100
Subject: [PATCH] execute_python is now safe + test

---
 generated_def/winfuncs.py   | 10 +++++-
 generated_def/winstructs.py |  1 -
 injection.py                | 43 +++++++++++++++++++++--
 test/__init__.py            |  3 +-
 test/mytest.py              | 23 +++++++++---
 utils/pythonutils.py        |  2 +-
 winobject.py                | 70 ++++++++++++++++++++++++++++++++++---
 7 files changed, 136 insertions(+), 16 deletions(-)

diff --git a/generated_def/winfuncs.py b/generated_def/winfuncs.py
index 9f21695..a5ffb64 100644
--- a/generated_def/winfuncs.py
+++ b/generated_def/winfuncs.py
@@ -3,7 +3,7 @@ from ctypes import *
 from ctypes.wintypes import *
 from .winstructs import *
 
-functions = ['ExitProcess', 'TerminateProcess', 'GetLastError', 'GetCurrentProcess', 'CreateFileA', 'CreateFileW', 'NtQuerySystemInformation', 'NtQueryInformationProcess', 'NtQueryVirtualMemory', 'NtCreateThreadEx', 'NtQueryInformationThread', 'VirtualAlloc', 'VirtualAllocEx', 'VirtualFree', 'VirtualFreeEx', 'VirtualProtect', 'VirtualQuery', 'VirtualQueryEx', 'GetModuleFileNameA', 'GetModuleFileNameW', 'CreateThread', 'CreateRemoteThread', 'VirtualProtect', 'CreateProcessA', 'CreateProcessW', 'GetThreadContext', 'NtGetContextThread', 'SetThreadContext', 'OpenThread', 'OpenProcess', 'CloseHandle', 'ReadProcessMemory', 'NtWow64ReadVirtualMemory64', 'WriteProcessMemory', 'CreateToolhelp32Snapshot', 'Thread32First', 'Thread32Next', 'Process32First', 'Process32Next', 'Process32FirstW', 'Process32NextW', 'GetProcAddress', 'LoadLibraryA', 'LoadLibraryW', 'OpenProcessToken', 'LookupPrivilegeValueA', 'LookupPrivilegeValueW', 'AdjustTokenPrivileges', 'FindResourceA', 'FindResourceW', 'SizeofResource', 'LoadResource', 'LockResource', 'GetVersionExA', 'GetVersionExW', 'GetVersion', 'GetCurrentThread', 'GetCurrentThreadId', 'GetCurrentProcessorNumber', 'AllocConsole', 'FreeConsole', 'GetStdHandle', 'SetStdHandle', 'SetThreadAffinityMask', 'WriteFile', 'GetExtendedTcpTable', 'GetExtendedUdpTable', 'SetTcpEntry', 'AddVectoredContinueHandler', 'AddVectoredExceptionHandler', 'TerminateThread', 'ExitThread', 'RemoveVectoredExceptionHandler', 'ResumeThread', 'SuspendThread', 'WaitForSingleObject', 'GetThreadId', 'LoadLibraryExA', 'LoadLibraryExW', 'SymInitialize', 'SymFromName', 'SymLoadModuleEx', 'SymSetOptions', 'SymGetTypeInfo', 'DeviceIoControl', 'GetTokenInformation', 'RegOpenKeyExA', 'RegOpenKeyExW', 'RegGetValueA', 'RegGetValueW', 'RegCloseKey', 'Wow64DisableWow64FsRedirection', 'Wow64RevertWow64FsRedirection', 'Wow64EnableWow64FsRedirection', 'Wow64GetThreadContext', 'SetConsoleCtrlHandler']
+functions = ['ExitProcess', 'TerminateProcess', 'GetLastError', 'GetCurrentProcess', 'CreateFileA', 'CreateFileW', 'NtQuerySystemInformation', 'NtQueryInformationProcess', 'NtQueryVirtualMemory', 'NtCreateThreadEx', 'NtQueryInformationThread', 'GetExitCodeThread', 'GetExitCodeProcess', 'VirtualAlloc', 'VirtualAllocEx', 'VirtualFree', 'VirtualFreeEx', 'VirtualProtect', 'VirtualQuery', 'VirtualQueryEx', 'GetModuleFileNameA', 'GetModuleFileNameW', 'CreateThread', 'CreateRemoteThread', 'VirtualProtect', 'CreateProcessA', 'CreateProcessW', 'GetThreadContext', 'NtGetContextThread', 'SetThreadContext', 'OpenThread', 'OpenProcess', 'CloseHandle', 'ReadProcessMemory', 'NtWow64ReadVirtualMemory64', 'WriteProcessMemory', 'CreateToolhelp32Snapshot', 'Thread32First', 'Thread32Next', 'Process32First', 'Process32Next', 'Process32FirstW', 'Process32NextW', 'GetProcAddress', 'LoadLibraryA', 'LoadLibraryW', 'OpenProcessToken', 'LookupPrivilegeValueA', 'LookupPrivilegeValueW', 'AdjustTokenPrivileges', 'FindResourceA', 'FindResourceW', 'SizeofResource', 'LoadResource', 'LockResource', 'GetVersionExA', 'GetVersionExW', 'GetVersion', 'GetCurrentThread', 'GetCurrentThreadId', 'GetCurrentProcessorNumber', 'AllocConsole', 'FreeConsole', 'GetStdHandle', 'SetStdHandle', 'SetThreadAffinityMask', 'WriteFile', 'GetExtendedTcpTable', 'GetExtendedUdpTable', 'SetTcpEntry', 'AddVectoredContinueHandler', 'AddVectoredExceptionHandler', 'TerminateThread', 'ExitThread', 'RemoveVectoredExceptionHandler', 'ResumeThread', 'SuspendThread', 'WaitForSingleObject', 'GetThreadId', 'LoadLibraryExA', 'LoadLibraryExW', 'SymInitialize', 'SymFromName', 'SymLoadModuleEx', 'SymSetOptions', 'SymGetTypeInfo', 'DeviceIoControl', 'GetTokenInformation', 'RegOpenKeyExA', 'RegOpenKeyExW', 'RegGetValueA', 'RegGetValueW', 'RegCloseKey', 'Wow64DisableWow64FsRedirection', 'Wow64RevertWow64FsRedirection', 'Wow64EnableWow64FsRedirection', 'Wow64GetThreadContext', 'SetConsoleCtrlHandler']
 
 # ExitProcess(uExitCode):
 ExitProcessPrototype = WINFUNCTYPE(VOID, UINT)
@@ -49,6 +49,14 @@ NtCreateThreadExParams = ((1, 'ThreadHandle'), (1, 'DesiredAccess'), (1, 'Object
 NtQueryInformationThreadPrototype = WINFUNCTYPE(NTSTATUS, HANDLE, THREAD_INFORMATION_CLASS, PVOID, ULONG, PULONG)
 NtQueryInformationThreadParams = ((1, 'ThreadHandle'), (1, 'ThreadInformationClass'), (1, 'ThreadInformation'), (1, 'ThreadInformationLength'), (1, 'ReturnLength'))
 
+# GetExitCodeThread(hThread, lpExitCode):
+GetExitCodeThreadPrototype = WINFUNCTYPE(BOOL, HANDLE, LPDWORD)
+GetExitCodeThreadParams = ((1, 'hThread'), (1, 'lpExitCode'))
+
+# GetExitCodeProcess(hProcess, lpExitCode):
+GetExitCodeProcessPrototype = WINFUNCTYPE(BOOL, HANDLE, LPDWORD)
+GetExitCodeProcessParams = ((1, 'hProcess'), (1, 'lpExitCode'))
+
 # VirtualAlloc(lpAddress, dwSize, flAllocationType, flProtect):
 VirtualAllocPrototype = WINFUNCTYPE(LPVOID, LPVOID, SIZE_T, DWORD, DWORD)
 VirtualAllocParams = ((1, 'lpAddress'), (1, 'dwSize'), (1, 'flAllocationType'), (1, 'flProtect'))
diff --git a/generated_def/winstructs.py b/generated_def/winstructs.py
index 6fecad8..2c3735a 100644
--- a/generated_def/winstructs.py
+++ b/generated_def/winstructs.py
@@ -858,7 +858,6 @@ PM128A = POINTER(_M128A)
 
 # Struct _CONTEXT64 definitions
 class _CONTEXT64(Structure):
-        _pack_ = 8
         _fields_ = [
         ("P1Home", DWORD64),
         ("P2Home", DWORD64),
diff --git a/injection.py b/injection.py
index d2f2ff0..c96f589 100644
--- a/injection.py
+++ b/injection.py
@@ -1,3 +1,6 @@
+import struct
+import ctypes
+
 import windows
 import windows.utils as utils
 
@@ -55,7 +58,7 @@ def generate_python_exec_shellcode_64(target, PYDLL_addr, PyInit, PyRun, PYCODE_
 
     code = x64.MultipleInstr()
     # Do stack alignement
-    code += x64.Push('RAX')
+    code += x64.Push('RCX')
     # Load python27.dll
     code += x64.Mov('RCX', PYDLL_addr)
     code += x64.Mov('RAX', LoadLibraryA)
@@ -89,7 +92,7 @@ def generate_python_exec_shellcode_64(target, PYDLL_addr, PyInit, PyRun, PYCODE_
     code += x64.Call('RAX')
     code += Clean_space_for_call
     # Remove stack alignement
-    code += x64.Pop('RAX')
+    code += x64.Pop('RCX')
     code += x64.Ret()
     return code.get_code()
 
@@ -130,3 +133,39 @@ def inject_python_command(process, code_injected, PYDLL="python27.dll\x00"):
 def execute_python_code(process, code):
     shellcode_remote_addr = inject_python_command(process, code)
     return process.create_thread(shellcode_remote_addr, 0)
+
+retrieve_exc = r"""
+import traceback
+import sys
+addr = {0}
+txt = "".join(traceback.format_exception(sys.last_type, sys.last_value, sys.last_traceback))
+import ctypes
+
+size = ctypes.c_uint.from_address(addr)
+size.value = len(txt)
+buff = (ctypes.c_char * len(txt)).from_address(addr + ctypes.sizeof(ctypes.c_uint))
+buff[:] = txt
+"""
+
+def retrieve_last_exception_data(process):
+    # TODO : FREE THIS
+    mem = process.virtual_alloc(0x1000)
+    execute_python_code(process, retrieve_exc.format(mem))
+    size = struct.unpack("<I", process.read_memory(mem, ctypes.sizeof(ctypes.c_uint)))[0]
+    data = process.read_memory(mem + ctypes.sizeof(ctypes.c_uint), size)
+    return data
+
+class RemotePythonError(Exception):
+    pass
+
+def safe_execute_python(process, code):
+    t = execute_python_code(process, code)
+    t.wait() # Wait terminaison of the thread
+    if t.exit_code == 0:
+        return True
+    if t.exit_code != 0xffffffff:
+        raise ValueError("Unknown exit code {0}".format(hex(t.exit_code)))
+    data = retrieve_last_exception_data(process)
+    raise RemotePythonError(data)
+
+
diff --git a/test/__init__.py b/test/__init__.py
index bea6043..cd33ee6 100644
--- a/test/__init__.py
+++ b/test/__init__.py
@@ -1,4 +1,3 @@
-from mytest import WindowsTestCase, pop_calc_32, pop_calc_64
-
+from mytest import WindowsTestCase, pop_calc_32, pop_calc_64, Calc32, Calc64
 
 __all__ = ["WindowsTestCase"]
diff --git a/test/mytest.py b/test/mytest.py
index 58f91c1..26cf669 100644
--- a/test/mytest.py
+++ b/test/mytest.py
@@ -129,7 +129,7 @@ class WindowsTestCase(unittest.TestCase):
         with Calc64() as calc:
             data = calc.virtual_alloc(0x1000)
             calc.execute_python('import ctypes; ctypes.c_uint.from_address({0}).value = 0x42424242'.format(data))
-            time.sleep(0.1)
+            #time.sleep(0.1)
             dword = struct.unpack("<I", calc.read_memory(data, 4))[0]
             self.assertEqual(dword, 0x42424242)
 
@@ -137,7 +137,7 @@ class WindowsTestCase(unittest.TestCase):
         with Calc32() as calc:
             data = calc.virtual_alloc(0x1000)
             calc.execute_python('import ctypes; ctypes.c_uint.from_address({0}).value = 0x42424242'.format(data))
-            time.sleep(0.1)
+            #time.sleep(0.1)
             dword = struct.unpack("<I", calc.read_memory(data, 4))[0]
             self.assertEqual(dword, 0x42424242)
 
@@ -169,7 +169,7 @@ class WindowsTestCase(unittest.TestCase):
                                 ctypes.c_uint.from_address({1}).value = GetCurrentProcessId
                                 """.format(os.getcwd(), data)
             calc.execute_python(textwrap.dedent(remote_python_code))
-            time.sleep(0.5)
+            #time.sleep(0.5)
             dword = struct.unpack("<I", calc.read_memory(data, 4))[0]
             self.assertEqual(dword, get_current_proc_id)
 
@@ -190,7 +190,7 @@ class WindowsTestCase(unittest.TestCase):
                                 ctypes.c_ulonglong.from_address({1}).value = GetCurrentProcessId
                                 """.format(os.getcwd(), data)
             calc.execute_python(textwrap.dedent(remote_python_code))
-            time.sleep(0.5)
+            #time.sleep(0.5)
             dword = struct.unpack("<Q", calc.read_memory(data, 8))[0]
             self.assertEqual(dword, get_current_proc_id)
 
@@ -228,6 +228,21 @@ class WindowsTestCase(unittest.TestCase):
             _winreg.OpenKey(*open_args)
         self.assertEqual(ar.exception.winerror, 0x11223344)
 
+    def test_thread_exit_value_32(self):
+        with Calc32() as calc:
+            res = calc.execute_python("import time;time.sleep(0.1); 2")
+            self.assertEqual(res, True)
+            with self.assertRaises(windows.injection.RemotePythonError) as ar:
+                t = calc.execute_python("import time;time.sleep(0.1); raise ValueError('BYE')")
+
+    def test_thread_exit_value_64(self):
+        with Calc64() as calc:
+            res = calc.execute_python("import time;time.sleep(0.1); 2")
+            self.assertEqual(res, True)
+            with self.assertRaises(windows.injection.RemotePythonError) as ar:
+                t = calc.execute_python("import time;time.sleep(0.1); raise ValueError('BYE')")
+
+
 
 if __name__ == '__main__':
     alltests = unittest.TestSuite()
diff --git a/utils/pythonutils.py b/utils/pythonutils.py
index f97ee50..730480c 100644
--- a/utils/pythonutils.py
+++ b/utils/pythonutils.py
@@ -11,7 +11,7 @@ def fixedpropety(f):
         except AttributeError:
             setattr(self, cache_name, f(self))
             return getattr(self, cache_name)
-    return property(prop)
+    return property(prop, doc=f.__doc__)
 
 
 def swallow_ctypes_copy(ctypes_object):
diff --git a/winobject.py b/winobject.py
index cd6628a..de0b858 100644
--- a/winobject.py
+++ b/winobject.py
@@ -25,7 +25,7 @@ import windows.pe_parse as pe_parse
 
 
 class AutoHandle(object):
-    """An abstract class that allow easy handle creation/destruction"""
+    """An abstract class that allow easy handle creation/destruction/wait"""
     def _get_handle(self):
         raise NotImplementedError("{0} is abstract".format(type(self).__name__))
 
@@ -44,6 +44,9 @@ class AutoHandle(object):
         self._handle = self._get_handle()
         return self._handle
 
+    def wait(self, timeout=INFINITE):
+        return winproxy.WaitForSingleObject(self.handle, timeout)
+
     def __del__(self):
         if hasattr(self, "_handle") and self._handle:
             winproxy.CloseHandle(self._handle)
@@ -156,6 +159,16 @@ class WinThread(THREADENTRY32, AutoHandle):
     def _get_handle(self):
         return winproxy.OpenThread(dwThreadId=self.tid)
 
+    @property
+    def is_exit(self):
+        return self.exit_code != STILL_ACTIVE
+
+    @property
+    def exit_code(self):
+        res = DWORD()
+        winproxy.GetExitCodeThread(self.handle, byref(res))
+        return res.value
+
     def __repr__(self):
         owner = self.owner
         if owner is None:
@@ -168,9 +181,32 @@ class WinThread(THREADENTRY32, AutoHandle):
     def _from_handle(handle):
         tid = winproxy.GetThreadId(handle)
         try:
-            return [t for t in System().threads if t.tid == tid][0]
+            # Really useful ?
+            thread = [t for t in System().threads if t.tid == tid][0]
+            # set AutoHandle _handle
+            thread._handle = handle
+            return thread
         except IndexError:
-            return (tid, handle)
+            return DeadThread(handle, tid)
+
+class DeadThread(AutoHandle):
+    """A simple object arround an already dead thread"""
+    def __init__(self, handle, tid=None):
+        if tid is None:
+            tid = winproxy.GetThreadId(handle)
+        self.tid = tid
+        # set AutoHandle _handle
+        self._handle = handle
+
+    @property
+    def is_exit(self):
+        return self.exit_code != STILL_ACTIVE
+
+    @property
+    def exit_code(self):
+        res = DWORD()
+        winproxy.GetExitCodeThread(self.handle, byref(res))
+        return res.value
 
 
 class Process(AutoHandle):
@@ -207,6 +243,16 @@ class Process(AutoHandle):
     def virtual_alloc(self, size):
         raise NotImplementedError("virtual_alloc")
 
+    @property
+    def exit_code(self):
+        res = DWORD()
+        winproxy.GetExitCodeProcess(self.handle, byref(res))
+        return res.value
+
+    @property
+    def is_exit(self):
+        return self.exit_code == STILL_ACTIVE
+
     def execute(self, code):
         """Execute some raw code in the context of the process"""
         x = self.virtual_alloc(len(code))
@@ -267,6 +313,10 @@ class CurrentThread(AutoHandle):
         """Exit the thread"""
         return winproxy.ExitThread(code)
 
+    def wait(self):
+        """Raise ValueError to prevent deadlock :D"""
+        raise ValueError("wait() on current thread")
+
 
 class CurrentProcess(Process):
     """The current process"""
@@ -348,7 +398,7 @@ class CurrentProcess(Process):
         return True
 
     def read_memory(self, addr, size):
-        """Read size from adddr"""
+        """Read size from addr"""
         dbgprint('Read CurrentProcess Memory', 'READMEM')
         buffer = (c_char * size).from_address(addr)
         return buffer[:]
@@ -366,6 +416,10 @@ class CurrentProcess(Process):
         """Exit the process"""
         return winproxy.ExitProcess(code)
 
+    def wait(self):
+        """Raise ValueError to prevent deadlock :D"""
+        raise ValueError("wait() on current thread")
+
 
 class WinProcess(PROCESSENTRY32, Process):
     """A Process on the system"""
@@ -453,7 +507,9 @@ class WinProcess(PROCESSENTRY32, Process):
     def create_thread(self, addr, param):
         """Create a remote thread"""
         if windows.current_process.bitness == 32 and self.bitness == 64:
-            return windows.syswow64.NtCreateThreadEx_32_to_64(ProcessHandle=self.handle, lpStartAddress=addr, lpParameter=param)
+            thread_handle = HANDLE()
+            windows.syswow64.NtCreateThreadEx_32_to_64(ThreadHandle=byref(thread_handle) ,ProcessHandle=self.handle, lpStartAddress=addr, lpParameter=param)
+            return WinThread._from_handle(thread_handle.value)
         return WinThread._from_handle(winproxy.CreateRemoteThread(hProcess=self.handle, lpStartAddress=addr, lpParameter=param))
 
     def load_library(self, dll_path):
@@ -464,6 +520,10 @@ class WinProcess(PROCESSENTRY32, Process):
         return self.create_thread(LoadLibrary, x)
 
     def execute_python(self, pycode):
+        """Execute Python code into the remote process"""
+        return injection.safe_execute_python(self, pycode)
+
+    def execute_python_unsafe(self, pycode):
         """Execute Python code into the remote process"""
         return injection.execute_python_code(self, pycode)