Merge branch 'feat/reg_save_restore' of https://github.com/rhelmot/unicorn into rhelmot-feat/reg_save_restore

bindings/python/sample_x86.py

@@ -312,6 +312,56 @@ def test_i386_inout():
         print("ERROR: %s" % e)
 
 
+def test_i386_reg_save():
+    print("Save/restore registers in opaque blob")
+    address = 0
+    code = '\x40'  # inc eax
+    try:
+        # Initialize emulator
+        mu = Uc(UC_ARCH_X86, UC_MODE_32)
+
+        # map 8KB memory for this emulation
+        mu.mem_map(address, 8 * 1024, UC_PROT_ALL)
+
+        # write machine code to be emulated to memory
+        mu.mem_write(address, code)
+
+        print(">>> set eax to 1")
+        mu.reg_write(UC_X86_REG_EAX, 1)
+
+        print(">>> execute 'inc eax'")
+        mu.emu_start(address, address+1)
+
+        print(">>> save the register state")
+        saved_regs = mu.regstate_save()
+
+        print(">>> execute 'inc eax'")
+        mu.emu_start(address, address+1)
+
+        print(">>> assert eax == 3")
+        assert mu.reg_read(UC_X86_REG_EAX) == 3
+
+        print(">>> restore the register state")
+        mu.regstate_restore(saved_regs)
+
+        print(">>> assert eax == 2")
+        assert mu.reg_read(UC_X86_REG_EAX) == 2
+
+        print(">>> execute 'inc eax'")
+        mu.emu_start(address, address+1)
+
+        print(">>> assert eax == 3")
+        assert mu.reg_read(UC_X86_REG_EAX) == 3
+
+        print(">>> restore the register state")
+        mu.regstate_restore(saved_regs)
+
+        print(">>> assert eax == 2")
+        assert mu.reg_read(UC_X86_REG_EAX) == 2
+
+    except UcError as e:
+        print("ERROR: %s" % e)
+
 def test_x86_64():
     print("Emulate x86_64 code")
     try:
@@ -483,6 +533,8 @@ if __name__ == '__main__':
     print("=" * 20)
     test_i386_inout()
     print("=" * 20)
+    test_i386_reg_save()
+    print("=" * 20)
     test_x86_64()
     print("=" * 20)
     test_x86_64_syscall()

bindings/python/unicorn/unicorn.py

@@ -100,6 +100,9 @@ _setup_prototype(_uc, "uc_mem_map_ptr", ucerr, uc_engine, ctypes.c_uint64, ctype
 _setup_prototype(_uc, "uc_mem_unmap", ucerr, uc_engine, ctypes.c_uint64, ctypes.c_size_t)
 _setup_prototype(_uc, "uc_mem_protect", ucerr, uc_engine, ctypes.c_uint64, ctypes.c_size_t, ctypes.c_uint32)
 _setup_prototype(_uc, "uc_query", ucerr, uc_engine, ctypes.c_uint32, ctypes.POINTER(ctypes.c_size_t))
+_setup_prototype(_uc, "uc_regstate_save", ctypes.c_voidp, uc_engine, ctypes.c_voidp)
+_setup_prototype(_uc, "uc_regstate_restore", None, uc_engine, ctypes.c_voidp)
+_setup_prototype(_uc, "free", None, ctypes.c_voidp)
 
 # uc_hook_add is special due to variable number of arguments
 _uc.uc_hook_add = _uc.uc_hook_add
@@ -440,6 +443,27 @@ class Uc(object):
             raise UcError(status)
         h = 0
 
+    def regstate_save(self, store=None):
+        if store is None:
+            ptr = ctypes.cast(0, ctypes.c_voidp)
+            return _ActivePointer(_uc.uc_regstate_save(self._uch, ptr))
+        elif type(store) is _ActivePointer:
+            _uc.uc_regstate_save(self._uch, store.pointer)
+            return store
+        else:
+            raise TypeError("Bad register store %s" % repr(store))
+
+    def regstate_restore(self, store):
+        if type(store) is not _ActivePointer:
+            raise TYpeError("Bad register store %s" % repr(store))
+        _uc.uc_regstate_restore(self._uch, store.pointer)
+
+class _ActivePointer(object):
+    def __init__(self, pointer):
+        self.pointer = pointer
+
+    def __del__(self):
+        _uc.free(self.pointer)
 
 # print out debugging info
 def debug():