From 8653c4c1e44fbbd512726bcc7648f20132b7e7ab Mon Sep 17 00:00:00 2001
From: hakril <project@hakril.net>
Date: Wed, 2 Nov 2016 22:35:18 +0100
Subject: [PATCH] No-closure pe_parse pass the test of WindowsTestCase

---
 windows/pe_parse.py        | 126 +++++++++++++++++--------------------
 windows/test/mytest.py     |  18 +++---
 windows/test/test_utils.py |  28 ++++++---
 3 files changed, 86 insertions(+), 86 deletions(-)

diff --git a/windows/pe_parse.py b/windows/pe_parse.py
index 63f725d..3158d04 100644
--- a/windows/pe_parse.py
+++ b/windows/pe_parse.py
@@ -430,29 +430,14 @@ class IMPORT_BY_NAME(ctypes.Structure):
         ("Name", BYTE)
     ]
 
-#class RVA(DWORD):
-#    @property
-#    def addr(self):
-#        return baseaddr + self.value
-#
-#    def __repr__(self):
-#        return "<DWORD {0} (RVA to '{1}')>".format(self.value, hex(self.addr))
-
 
 def get_string(target, addr):
     if target is None:
         return ctypes.c_char_p(addr).value
     return target.read_string(addr)
 
-#class StringRVa(RVA):
-#    def str(self):
-#        return get_string(self.addr).decode()
 
-
-class PESection((IMAGE_SECTION_HEADER)):
-    def __init__(self, baseaddr, target):
-        self.target = target
-        self.baseaddr = baseaddr
+class PESection(IMAGE_SECTION_HEADER):
 
     @property
     def name(self):
@@ -472,6 +457,14 @@ class PESection((IMAGE_SECTION_HEADER)):
         return "<PESection \"{0}\">".format(self.name)
 
 
+    @classmethod
+    def create(cls, pefile, addr):
+        self = pefile.transformers.create_structure_at(cls, addr)
+        self.baseaddr = pefile.baseaddr
+        self.target = pefile.target
+        return self
+
+
 class IATEntry(ctypes.Structure):
     """Represent an entry in the IAT of a module
     Can be used to get resolved value and setup hook
@@ -527,9 +520,7 @@ class IATEntry(ctypes.Structure):
         return True
 
 
-class IMAGE_IMPORT_DESCRIPTOR(ctypes.Structure):
-    _fields_ = transform_ctypes_fields(IMAGE_IMPORT_DESCRIPTOR, {})
-
+class IMAGE_IMPORT_DESCRIPTOR(IMAGE_IMPORT_DESCRIPTOR): # TODO: use explicite name winstructs.IMAGE_IMPORT_DESCRIPTOR
     def get_INT(self):
         if not self.OriginalFirstThunk:
             return None
@@ -561,6 +552,40 @@ class IMAGE_IMPORT_DESCRIPTOR(ctypes.Structure):
             iat_entry = self.transformers.create_structure_at(THUNK_DATA, iat_addr)
         return res
 
+    @classmethod
+    def create(cls, pefile, addr):
+        self = pefile.transformers.create_structure_at(cls, addr)
+        self.baseaddr = pefile.baseaddr
+        self.transformers = pefile.transformers
+        self.IMAGE_ORDINAL_FLAG = pefile.IMAGE_ORDINAL_FLAG
+        self.target = pefile.target
+        return self
+
+class IMAGE_EXPORT_DIRECTORY(IMAGE_EXPORT_DIRECTORY): # TODO: use explicite name winstructs._IMAGE_EXPORT_DIRECTORY
+    def get_exports(self):
+        NameOrdinals = self.transformers.create_structure_at((WORD * self.NumberOfNames), self.AddressOfNameOrdinals + self.baseaddr)
+        NameOrdinals = list(NameOrdinals)
+        Functions = self.transformers.create_structure_at((DWORD * self.NumberOfFunctions), self.AddressOfFunctions + self.baseaddr)
+        Names = self.transformers.create_structure_at((DWORD * self.NumberOfNames), self.AddressOfNames + self.baseaddr)
+        res = []
+        for nb, func in enumerate(Functions):
+            func += self.baseaddr
+            if nb in NameOrdinals:
+                name = get_string(self.target, Names[NameOrdinals.index(nb)] + self.baseaddr)
+            else:
+                name = None
+            res.append((nb, func, name))
+        return res
+
+    @classmethod
+    def create(cls, pefile, addr):
+        self = pefile.transformers.create_structure_at(cls, addr)
+        self.transformers = pefile.transformers
+        self.target = pefile.target
+        self.baseaddr = pefile.baseaddr
+        return self
+
+
 class IMAGE_DOS_HEADER(ctypes.Structure):
     _fields_ = [
         ("e_magic", CHAR * 2),
@@ -586,10 +611,6 @@ class IMAGE_DOS_HEADER(ctypes.Structure):
 
 class PEFile(object):
     """Represent a PE loaded in a process (current or remote)"""
-     #def __del__(self):
-     #    print("YOLODEL")
-     #    import pprint
-     #    pprint.pprint(vars(self))
 
     def __init__(self, target, baseaddr, targetedbitness, transformers):
         self.target = target
@@ -638,23 +659,12 @@ class PEFile(object):
         if import_datadir.VirtualAddress == 0:
             return []
         import_descriptor_addr = self.baseaddr + import_datadir.VirtualAddress
-        current_import_descriptor = self.transformers.create_structure_at(IMAGE_IMPORT_DESCRIPTOR, import_descriptor_addr)
-        # YOLO FIX: todo better
-        current_import_descriptor.baseaddr = self.baseaddr
-        current_import_descriptor.transformers = self.transformers
-        current_import_descriptor.IMAGE_ORDINAL_FLAG = self.IMAGE_ORDINAL_FLAG
-        current_import_descriptor.target = self.target
-
+        current_import_descriptor =  IMAGE_IMPORT_DESCRIPTOR.create(self, import_descriptor_addr)
         res = []
         while current_import_descriptor.FirstThunk:
             res.append(current_import_descriptor)
             import_descriptor_addr += ctypes.sizeof(IMAGE_IMPORT_DESCRIPTOR)
-            current_import_descriptor = self.transformers.create_structure_at(IMAGE_IMPORT_DESCRIPTOR, import_descriptor_addr)
-            # YOLO FIX: todo better
-            current_import_descriptor.baseaddr = self.baseaddr
-            current_import_descriptor.transformers = self.transformers
-            current_import_descriptor.IMAGE_ORDINAL_FLAG = self.IMAGE_ORDINAL_FLAG
-            current_import_descriptor.target = self.target
+            current_import_descriptor =  IMAGE_IMPORT_DESCRIPTOR.create(self, import_descriptor_addr)
         return res
 
     def get_EXPORT_DIRECTORY(self):
@@ -662,11 +672,7 @@ class PEFile(object):
         if export_directory_rva == 0:
             return None
         export_directory_addr = self.baseaddr + export_directory_rva
-        exp_dir = self.transformers.create_structure_at(self._IMAGE_EXPORT_DIRECTORY, export_directory_addr)
-        # YOLO AGAIN.. FIXME
-        exp_dir.transformers = self.transformers
-        exp_dir.target = self.target
-        exp_dir.baseaddr = self.baseaddr
+        exp_dir = IMAGE_EXPORT_DIRECTORY.create(self, export_directory_addr)
         return exp_dir
 
     @utils.fixedpropety
@@ -679,8 +685,10 @@ class PEFile(object):
         else:
             opt_header_addr = self.get_NT_HEADER().OptionalHeader._base_addr
         base_section = opt_header_addr + SizeOfOptionalHeader
-        sections_array = self.transformers.create_structure_at((self.PESection * nb_section), base_section)
-        return list(sections_array)
+        #pe_section_type = IMAGE_SECTION_HEADER
+        return [PESection.create(self, base_section + (sizeof(IMAGE_SECTION_HEADER) * i)) for i in range(nb_section)]
+        #sections_array = self.transformers.create_structure_at((self.PESection * nb_section), base_section)
+        #return list(sections_array)
 
     @utils.fixedpropety
     def exports(self):
@@ -702,10 +710,12 @@ class PEFile(object):
     @utils.fixedpropety
     def export_name(self):
         """The Name attribute of the ``EXPORT_DIRECTORY``"""
-        try:
-            return self.get_EXPORT_DIRECTORY().Name.str
-        except AttributeError:
+        exp_dir = self.get_EXPORT_DIRECTORY()
+        if exp_dir is None:
             return None
+        if not exp_dir.Name:
+            return None
+        return get_string(self.target, self.baseaddr + exp_dir.Name)
 
     # TODO: get imports by parsing other modules exports if no INT
     @utils.fixedpropety
@@ -726,26 +736,4 @@ class PEFile(object):
                     iat_entry.name = str(name.decode()) if name else ""
             name = get_string(self.target, self.baseaddr + import_descriptor.Name)
             res.setdefault(name.lower(), []).extend(IAT)
-        return res
-
-    # Will be usable as `self.IMAGE_IMPORT_DESCRIPTOR`
-
-
-    # Will be usable as `self._IMAGE_EXPORT_DIRECTORY`
-    class _IMAGE_EXPORT_DIRECTORY(ctypes.Structure):
-        _fields_ = transform_ctypes_fields(IMAGE_EXPORT_DIRECTORY, {})
-
-        def get_exports(self):
-            NameOrdinals = self.transformers.create_structure_at((WORD * self.NumberOfNames), self.AddressOfNameOrdinals + self.baseaddr)
-            NameOrdinals = list(NameOrdinals)
-            Functions = self.transformers.create_structure_at((DWORD * self.NumberOfFunctions), self.AddressOfFunctions + self.baseaddr)
-            Names = self.transformers.create_structure_at((DWORD * self.NumberOfNames), self.AddressOfNames + self.baseaddr)
-            res = []
-            for nb, func in enumerate(Functions):
-                func += self.baseaddr
-                if nb in NameOrdinals:
-                    name = get_string(self.target, Names[NameOrdinals.index(nb)] + self.baseaddr)
-                else:
-                    name = None
-                res.append((nb, func, name))
-            return res
+        return res
\ No newline at end of file
diff --git a/windows/test/mytest.py b/windows/test/mytest.py
index 8354600..217ce22 100644
--- a/windows/test/mytest.py
+++ b/windows/test/mytest.py
@@ -198,10 +198,11 @@ class WindowsTestCase(unittest.TestCase):
             dword = struct.unpack("<I", calc.read_memory(data, 4))[0]
             self.assertEqual(dword, 0x42424242)
             # Check calc32 is still suspended:
-                # 1 thread
-                # suspend count == 1
-            self.assertEqual(len(calc.threads), 1)
+                # 1 thread | except windows 10 that pop threads
+                # main thread suspend count == 1
             self.assertEqual(calc.threads[0].suspend(), 1)
+            if not is_windows_10:
+                self.assertEqual(len(calc.threads), 1)
 
     @windows_64bit_only
     def test_execute_python_to_64_suspended(self):
@@ -211,10 +212,11 @@ class WindowsTestCase(unittest.TestCase):
             dword = struct.unpack("<I", calc.read_memory(data, 4))[0]
             self.assertEqual(dword, 0x42424242)
             # Check calc32 is still suspended:
-                # 1 thread
-                # suspend count == 1
-            self.assertEqual(len(calc.threads), 1)
+                # 1 thread | except windows 10 that pop threads
+                # main thread suspend count == 1
             self.assertEqual(calc.threads[0].suspend(), 1)
+            if not is_windows_10:
+                self.assertEqual(len(calc.threads), 1)
 
 
     def test_parse_remote_32_peb(self):
@@ -223,12 +225,12 @@ class WindowsTestCase(unittest.TestCase):
             # Yeah a don't know but on 32bits system the parsing might begin before
             # InMemoryOrderModuleList is setup..
             import time; time.sleep(0.1)
-            self.assertEqual(calc.peb.modules[0].name, "calc.exe")
+            self.assertEqual(calc.peb.modules[0].name, test_binary_name)
 
     @windows_64bit_only
     def test_parse_remote_64_peb(self):
         with Calc64() as calc:
-            self.assertEqual(calc.peb.modules[0].name, "calc.exe")
+            self.assertEqual(calc.peb.modules[0].name, test_binary_name)
 
     def test_parse_remote_32_pe(self):
         with Calc32() as calc:
diff --git a/windows/test/test_utils.py b/windows/test/test_utils.py
index da75deb..8da0ffd 100644
--- a/windows/test/test_utils.py
+++ b/windows/test/test_utils.py
@@ -6,6 +6,7 @@ import windows.debug
 import windows.native_exec.simple_x86 as x86
 import windows.native_exec.simple_x64 as x64
 import windows.native_exec.nativeutils as nativeutils
+from windows.generated_def import CREATE_NEW_CONSOLE
 
 
 is_process_32_bits = windows.current_process.bitness == 32
@@ -14,30 +15,39 @@ is_process_64_bits = windows.current_process.bitness == 64
 is_windows_32_bits = windows.system.bitness == 32
 is_windows_64_bits = windows.system.bitness == 64
 
+is_windows_10 = (windows.system.version[0] == 10)
+
 windows_32bit_only = unittest.skipIf(not is_windows_32_bits, "Test for 32bits Kernel only")
 windows_64bit_only = unittest.skipIf(not is_windows_64_bits, "Test for 64bits Kernel only")
 
 process_32bit_only = unittest.skipIf(not is_process_32_bits, "Test for 32bits process only")
 process_64bit_only = unittest.skipIf(not is_process_64_bits, "Test for 64bits process only")
 
+if windows.system.version[0] < 10:
+    test_binary_name = "calc.exe"
+else:
+    test_binary_name = "cmd.exe"
+
+DEFAULT_CREATION_FLAGS = CREATE_NEW_CONSOLE
+#DEFAULT_CREATION_FLAGS = 0
 
 if is_windows_32_bits:
-    def pop_calc_32(dwCreationFlags=0):
-        return windows.utils.create_process(r"C:\Windows\system32\calc.exe", dwCreationFlags=dwCreationFlags, show_windows=True)
+    def pop_calc_32(dwCreationFlags=DEFAULT_CREATION_FLAGS):
+        return windows.utils.create_process(r"C:\Windows\system32\{0}".format(test_binary_name), dwCreationFlags=dwCreationFlags, show_windows=True)
 
-    def pop_calc_64(dwCreationFlags=0):
+    def pop_calc_64(dwCreationFlags=DEFAULT_CREATION_FLAGS):
         raise WindowsError("Cannot create calc64 in 32bits system")
 else:
-    def pop_calc_32(dwCreationFlags=0):
-        return windows.utils.create_process(r"C:\Windows\syswow64\calc.exe", dwCreationFlags=dwCreationFlags, show_windows=True)
+    def pop_calc_32(dwCreationFlags=DEFAULT_CREATION_FLAGS):
+        return windows.utils.create_process(r"C:\Windows\syswow64\{0}".format(test_binary_name), dwCreationFlags=dwCreationFlags, show_windows=True)
 
     if is_process_32_bits:
-        def pop_calc_64(dwCreationFlags=0):
+        def pop_calc_64(dwCreationFlags=DEFAULT_CREATION_FLAGS):
             with windows.utils.DisableWow64FsRedirection():
-                return windows.utils.create_process(r"C:\Windows\system32\calc.exe", dwCreationFlags=dwCreationFlags, show_windows=True)
+                return windows.utils.create_process(r"C:\Windows\system32\{0}".format(test_binary_name), dwCreationFlags=dwCreationFlags, show_windows=True)
     else:
-        def pop_calc_64(dwCreationFlags=0):
-            return windows.utils.create_process(r"C:\Windows\system32\calc.exe", dwCreationFlags=dwCreationFlags, show_windows=True)
+        def pop_calc_64(dwCreationFlags=DEFAULT_CREATION_FLAGS):
+            return windows.utils.create_process(r"C:\Windows\system32\{0}".format(test_binary_name), dwCreationFlags=dwCreationFlags, show_windows=True)
 
 
 @contextmanager