diff --git a/windows/test/mytest.py b/windows/test/mytest.py index 353d4a7..530835c 100644 --- a/windows/test/mytest.py +++ b/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: diff --git a/windows/winobject/process.py b/windows/winobject/process.py index b361112..2853006 100644 --- a/windows/winobject/process.py +++ b/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")])