Add support for remote TEB + tests

Author: hakril

ctypes_generation/extended_structs/_LIST_ENTRY.py

@@ -0,0 +1,18 @@
+# From: ctypes_generation\extended_structs\_LIST_ENTRY.py
+# _LIST_ENTRY is a self referencing structure
+# Currently ctypes generation does not support extending self referencing structures
+# Ass the _fields_ assignement should happen after the extended structure definition
+# So we just redefine fully _LIST_ENTRY without inheriting the real one
+
+class _LIST_ENTRY(Structure):
+    def get_real_struct(self, targetcls, target_field):
+        # >>> gdef.LDR_DATA_TABLE_ENTRY.InMemoryOrderLinks
+        # <Field type=_LIST_ENTRY, ofs=16, size=16>
+        # This field object does not allow to retrieve the type..
+        # So we need to basse the target class AND the target field..
+        return targetcls.from_address(ctypes.addressof(self) - target_field.offset)
+
+_LIST_ENTRY._fields_ = [
+    ("Flink", POINTER(_LIST_ENTRY)),
+    ("Blink", POINTER(_LIST_ENTRY)),
+]

tests/test_process.py

@@ -28,7 +28,7 @@ def test_get_current_process_peb(self):
         return windows.current_process.peb
 
     def test_get_current_process_modules(self):
-        # Use module filename because this executable can be: 
+        # Use module filename because this executable can be:
         # 1. A PyInstaller exe
         # 2. A Windows App execution alias (Microsoft Store builds)
         assert os.path.basename(windows.current_process.peb.ProcessParameters[0].ImagePathName.str) in windows.current_process.peb.modules[0].name
@@ -474,19 +474,44 @@ def test_mapped_filename(self, proc32_64):
         with proc32_64.allocated_memory(0x1000) as addr:
             assert proc32_64.get_mapped_filename(addr) is None
 
+    def test_current_thread_teb():
+        teb = windows.current_thread.teb
+        assert ctypes.addressof(teb) == ctypes.addressof(windows.current_thread.teb.NtTib.Self[0])
+        assert ctypes.addressof(windows.current_process.peb) == ctypes.addressof(teb.ProcessEnvironmentBlock[0])
+        # Check type of teb.peb is the correct subclass (with modules & co)
+        assert teb.peb.modules
 
     def test_thread_teb_base(self, proc32_64):
         t = proc32_64.threads[0]
         assert t.teb_base != 0
 
+    def test_teb(self, proc32_64):
+        teb = proc32_64.threads[0].teb
+        if proc32_64.bitness == 32:
+            assert type(teb) == windows.winobject.process.RemoteTEB32
+        else:
+            assert type(teb) == windows.winobject.process.RemoteTEB64
+        assert teb.NtTib.Self.value == teb._base_addr
+        assert teb.ProcessEnvironmentBlock.value == teb.peb._base_addr
+        # Check type of teb.peb is the correct subclass (with modules & co)
+        assert teb.peb.modules
+
     @windows_64bit_only
     def test_thread_teb_syswow_base(self, proc32):
         t = proc32.threads[0]
         assert t.teb_base != 0
         assert t.teb_syswow_base != 0
         assert t.teb_base == t.teb_syswow_base + 0x2000
 
-
+    @windows_64bit_only
+    def test_thread_teb_syswow(self, proc32):
+        teb_syswow = proc32.threads[0].teb_syswow
+        assert type(teb_syswow) == windows.winobject.process.RemoteTEB64
+        assert type(teb_syswow.peb) == windows.winobject.process.RemotePEB64
+        assert teb_syswow.NtTib.Self.value == teb_syswow._base_addr
+        assert teb_syswow.ProcessEnvironmentBlock.value == teb_syswow.peb._base_addr
+        # Check type of teb.peb is the correct subclass (with modules & co)
+        assert teb.peb.modules
 
     def test_thread_owner_from_tid(self, proc32_64):
         thread = proc32_64.threads[0]

windows/generated_def/winstructs.py

@@ -3896,6 +3896,27 @@ class _LIST_ENTRY(Structure):
     ("Blink", POINTER(_LIST_ENTRY)),
 ]
 
+# From: ctypes_generation\extended_structs\_LIST_ENTRY.py
+# _LIST_ENTRY is a self referencing structure
+# Currently ctypes generation does not support extending self referencing structures
+# Ass the _fields_ assignement should happen after the extended structure definition
+# So we just redefine fully _LIST_ENTRY without inheriting the real one
+
+class _LIST_ENTRY(Structure):
+    def get_real_struct(self, targetcls, target_field):
+        # >>> gdef.LDR_DATA_TABLE_ENTRY.InMemoryOrderLinks
+        # <Field type=_LIST_ENTRY, ofs=16, size=16>
+        # This field object does not allow to retrieve the type..
+        # So we need to basse the target class AND the target field..
+        return targetcls.from_address(ctypes.addressof(self) - target_field.offset)
+
+_LIST_ENTRY._fields_ = [
+    ("Flink", POINTER(_LIST_ENTRY)),
+    ("Blink", POINTER(_LIST_ENTRY)),
+]
+LIST_ENTRY = _LIST_ENTRY
+PLIST_ENTRY = POINTER(_LIST_ENTRY)
+PRLIST_ENTRY = POINTER(_LIST_ENTRY)
 class _LSA_UNICODE_STRING(Structure):
     _fields_ = [
         ("Length", USHORT),

windows/winobject/process.py

@@ -899,7 +899,9 @@ def teb_base(self):
 
     @property
     def teb(self):
-        return RemoteTEB(self.teb_base, target=self.owner)
+        if self.owner.bitness == 32:
+            return RemoteTEB32(self.teb_base, target=self.owner)
+        return RemoteTEB64(self.teb_base, target=self.owner)
 
     @property
     def teb_syswow_base(self):
@@ -914,7 +916,7 @@ def teb_syswow_base(self):
 
     @property
     def teb_syswow(self):
-        return TEB64.from_address(self.teb_syswow_base)
+        return RemoteTEB64.from_address(self.teb_syswow_base)
 
 
     def exit(self, code=0):
@@ -1222,8 +1224,6 @@ def exit(self, code=0):
         return winproxy.TerminateProcess(self.handle, code)
 
 
-
-
 def transform_ctypes_fields(struct, replacement):
     return [(name, replacement.get(name, type)) for name, type in struct._fields_]
 
@@ -1266,11 +1266,6 @@ def pe(self):
         return pe_parse.GetPEFile(self.baseaddr)
 
 
-class LIST_ENTRY_PTR(PVOID):
-    def TO_LDR_ENTRY(self):
-        return LDR_DATA_TABLE_ENTRY.from_address(self.value - sizeof(PVOID) * 2)
-
-
 class PEB(gdef.PEB):
     """The PEB (Process Environment Block) of the current process"""
 
@@ -1305,13 +1300,13 @@ def modules(self):
         :type: [:class:`LoadedModule`] -- List of loaded modules
 		"""
         res = []
-        list_entry_ptr = ctypes.cast(self.Ldr.contents.InMemoryOrderModuleList.Flink, LIST_ENTRY_PTR)
-        current_dll = list_entry_ptr.TO_LDR_ENTRY()
+        first_flink = self.Ldr.contents.InMemoryOrderModuleList.Flink[0]
+        current_dll = first_flink.get_real_struct(LoadedModule, LoadedModule.InMemoryOrderLinks)
         while current_dll.DllBase:
             res.append(current_dll)
-            list_entry_ptr = ctypes.cast(current_dll.InMemoryOrderLinks.Flink, LIST_ENTRY_PTR)
-            current_dll = list_entry_ptr.TO_LDR_ENTRY()
-        return [LoadedModule.from_address(addressof(LDR)) for LDR in res]
+            next_flink = current_dll.InMemoryOrderLinks.Flink[0]
+            current_dll = next_flink.get_real_struct(LoadedModule, LoadedModule.InMemoryOrderLinks)
+        return res
 
     @staticmethod
     def _extract_environment(env_block_addr, target):
@@ -1341,15 +1336,6 @@ def apisetmap(self):
             raise NotImplementedError("ApiSetMap does not exist prior to Windows 7")
         return apisetmap.get_api_set_map_for_current_process(self.ApiSetMap)
 
-# TEB enhanced, same bitness as PEB (current process)
-class TEB(gdef.TEB):
-    def peb(self):
-        return ctypes.cast(self.ProcessEnvironmentBlock, ctypes.POINTER(PEB))[0]
-
-class RemoteTEB(rctypes.RemoteStructure.from_structure(TEB)):
-    def peb(self):
-        return ctypes.cast(self.ProcessEnvironmentBlock, ctypes.POINTER(PEB))[0]
-
 # Memory stuff
 
 class EPSAPI_WORKING_SET_BLOCK_BASE(object):
@@ -1461,10 +1447,24 @@ def apisetmap(self):
         raise NotImplementedError("ApiSetMap for remote process not implemented yet")
 
 
+# TEB enhanced, same bitness as PEB (current process)
+class TEB(gdef.TEB):
+    @property
+    def peb(self):
+        return ctypes.cast(self.ProcessEnvironmentBlock, ctypes.POINTER(PEB))[0]
 
-
+# mote TEB enhanced, same bitness as PEB (current process)
+class RemoteTEB(rctypes.RemoteStructure.from_structure(TEB)):
+    @property
+    def peb(self):
+        ctypes_peb = self.ProcessEnvironmentBlock.value
+        return RemotePEB(ctypes_peb, self._target)
 
 if CurrentProcess().bitness == 32:
+    RemoteLoadedModule32 = RemoteLoadedModule
+    RemotePEB32 = RemotePEB
+    RemoteTEB32 = RemoteTEB
+
     class RemoteLoadedModule64(rctypes.transform_type_to_remote64bits(LoadedModule)):
         @property
         def pe(self):
@@ -1479,7 +1479,6 @@ class RemotePEB64(rctypes.transform_type_to_remote64bits(PEB)):
         def ptr_flink_to_remote_module(self, ptr_value):
             return RemoteLoadedModule64(ptr_value - ctypes.sizeof(rctypes.c_void_p64) * 2, self._target)
 
-
         @property
         def exe(self):
             """The executable of the process, as pointed by PEB.ImageBaseAddress
@@ -1512,7 +1511,17 @@ def environment(self):
 
         apisetmap = RemotePEB.apisetmap
 
+    class RemoteTEB64(rctypes.transform_type_to_remote64bits(TEB)):
+        @property
+        def peb(self):
+            ctypes_peb = self.ProcessEnvironmentBlock.value
+            return RemotePEB64(ctypes_peb, self._target)
+
+
 if CurrentProcess().bitness == 64:
+    RemoteLoadedModule64 = RemoteLoadedModule
+    RemotePEB64 = RemotePEB
+    RemoteTEB64 = RemoteTEB
 
     class RemoteLoadedModule32(rctypes.transform_type_to_remote32bits(LoadedModule)):
         @property
@@ -1558,4 +1567,10 @@ def environment(self):
             # TODO: Tests
             return self._extract_environment(self.ProcessParameters.contents.Environment, self._target)
 
-        apisetmap = RemotePEB.apisetmap
+        apisetmap = RemotePEB.apisetmap
+
+    class RemoteTEB32(rctypes.transform_type_to_remote32bits(TEB)):
+        @property
+        def peb(self):
+            ctypes_peb = self.ProcessEnvironmentBlock.value
+            return RemotePEB32(ctypes_peb, self._target)