diff --git a/bindings/python/sample_x86.py b/bindings/python/sample_x86.py index 5c9c0d69..25d947c1 100755 --- a/bindings/python/sample_x86.py +++ b/bindings/python/sample_x86.py @@ -4,7 +4,7 @@ from __future__ import print_function from unicorn import * from unicorn.x86_const import * - +import pickle X86_CODE32 = b"\x41\x4a\x66\x0f\xef\xc1" # INC ecx; DEC edx; PXOR xmm0, xmm1 X86_CODE32_LOOP = b"\x41\x4a\xeb\xfe" # INC ecx; DEC edx; JMP self-loop @@ -453,11 +453,17 @@ def test_i386_context_save(): print(">>> Saving CPU context") saved_context = mu.context_save() + print(">>> Pickling CPU context") + pickled_saved_context = pickle.dumps(saved_context) + print(">>> Running emulation for the second time") mu.emu_start(address, address+1) print(">>> Emulation done. Below is the CPU context") print(">>> EAX = 0x%x" %(mu.reg_read(UC_X86_REG_EAX))) + print(">>> Unpickling CPU context") + saved_context = pickle.loads(pickled_saved_context) + print(">>> CPU context restored. Below is the CPU context") mu.context_restore(saved_context) print(">>> EAX = 0x%x" %(mu.reg_read(UC_X86_REG_EAX))) diff --git a/bindings/python/unicorn/unicorn.py b/bindings/python/unicorn/unicorn.py index 7b8f49be..b97129e4 100644 --- a/bindings/python/unicorn/unicorn.py +++ b/bindings/python/unicorn/unicorn.py @@ -604,7 +604,7 @@ class Uc(object): return context def context_update(self, context): - status = _uc.uc_context_save(self._uch, context) + status = _uc.uc_context_save(self._uch, context.context) if status != uc.UC_ERR_OK: raise UcError(status) @@ -628,16 +628,40 @@ class Uc(object): _uc.uc_free(regions) -class UcContext(ctypes.Structure): +class UcContext: def __init__(self, h): - self.context = uc_context() - - status = _uc.uc_context_alloc(h, ctypes.byref(self.context)) + self._context = uc_context() + self._size = _uc.uc_context_size(h) + self._to_free = True + status = _uc.uc_context_alloc(h, ctypes.byref(self._context)) if status != uc.UC_ERR_OK: raise UcError(status) + + @property + def context(self): + return self._context + + @property + def size(self): + return self._size + + # Make UcContext picklable + def __getstate__(self): + return (bytes(self), self.size) + + def __setstate__(self, state): + self._size = state[1] + self._context = ctypes.cast(ctypes.create_string_buffer(state[0], self._size), uc_context) + # __init__ won'e be invoked, so we are safe to set it here. + self._to_free = False + + def __bytes__(self): + return ctypes.string_at(self.context, self.size) def __del__(self): - _uc.uc_free(self.context) + # We need this property since we shouldn't free it if the object is constructed from pickled bytes. + if self._to_free: + _uc.uc_free(self._context) # print out debugging info diff --git a/uc.c b/uc.c index 3e2f7dc9..f27a6a11 100644 --- a/uc.c +++ b/uc.c @@ -1321,12 +1321,12 @@ UNICORN_EXPORT uc_err uc_context_alloc(uc_engine *uc, uc_context **context) { struct uc_context **_context = context; - size_t size = cpu_context_size(uc->arch, uc->mode); + size_t size = uc_context_size(uc); *_context = malloc(size); if (*_context) { (*_context)->jmp_env_size = sizeof(*uc->cpu->jmp_env); - (*_context)->context_size = size - sizeof(uc_context) - (*_context)->jmp_env_size; + (*_context)->context_size = cpu_context_size(uc->arch, uc->mode); return UC_ERR_OK; } else { return UC_ERR_NOMEM;