Better handling of read_wstring for wstring at the end of readable memory

windows/test/mytest.py

@@ -112,6 +112,22 @@ class WindowsTestCase(unittest.TestCase):
                 calc.write_memory(k32.baseaddr, "XD")
             self.assertEqual(calc.read_memory(k32.baseaddr, 2), "XD")
 
+    def test_read_wstring(self):
+        test_string = "TEST_STRING"
+        with Calc32() as calc:
+            addr = calc.virtual_alloc(0x1000)
+            calc.write_memory(addr, "\x00".join(test_string + "\x00"))
+            self.assertEqual(calc.read_wstring(addr), test_string)
+
+
+    def test_read_wstring_end_page(self):
+        test_string = "TEST_STRING"
+        with Calc32() as calc:
+            # Setup string addr at end of page
+            addr = calc.virtual_alloc(0x1000) + 0x1000 - 26
+            calc.write_memory(addr, "\x00".join(test_string + "\x00"))
+            self.assertEqual(calc.read_wstring(addr), test_string)
+
     # Native execution
     def test_execute_to_32(self):
         with Calc32() as calc:
@@ -490,6 +506,7 @@ class WindowsTestCase(unittest.TestCase):
             t = calc.threads[0]
             self.assertNotEqual(t.teb_base, 0)
 
+
 class WindowsAPITestCase(unittest.TestCase):
     def test_createfileA_fail(self):
         with self.assertRaises(WindowsError) as ar:

windows/winobject/process.py

@@ -549,8 +549,21 @@ class Process(AutoHandle):
     def read_wstring(self, addr):
         """Read a windows UTF16 string at ``addr``"""
         res = []
-        for i in itertools.count():
-            x = self.read_memory(addr + (i * 0x100), 0x100)
+        read_size = 0x100
+        readden = 0
+        # I am trying to do something smart here..
+        while True:
+            try:
+                x = self.read_memory(addr + readden, read_size)
+            except winproxy.Kernel32Error as e:
+                if read_size == 2:
+                    raise
+                # handle read_wstring at end of page
+                # Of read failed: read only the half of size
+                # read_size must remain a multiple of 2
+                read_size = read_size / 2
+                continue
+            readden += read_size
             utf16_chars = ["".join(c) for c in zip(*[iter(x)] * 2)]
             if "\x00\x00" in utf16_chars:
                 res.extend(utf16_chars[:utf16_chars.index("\x00\x00")])