unicorn_obj.py

from __future__ import print_function
from unicorn import *
from unicorn.arm64_const import *
from unicorn.arm_const import *
from unicorn.m68k_const import *
from unicorn.mips_const import *
from unicorn.sparc_const import *
from unicorn.unicorn_const import *
from unicorn.x86_const import *

class UnicornMemory(object):
    def __init__(self, uc, addr, size):
        self.uc = uc
        self.addr = addr
        self.size = size
        uc.mem_map(addr, size)

    def __getitem__(self, i):
        if isinstance(i, slice):
            start = i.start
            stop = i.stop
            step = i.step
            if start is None: start = self.addr
            if stop is None: stop = self.addr + self.size
            if step is not None:
                return bytearray(''.join(str(self.uc.mem_read(addr, 1))
                        for addr in range(*i.indices(self.addr+self.size))))
            else:
                return self.uc.mem_read(start, stop - start)
        else:
            return self.uc.mem_read(i, 1)

    def __setitem__(self, i, val):
        if isinstance(i, slice):
            start = i.start
            stop = i.stop
            step = i.step
            if start is None: start = self.addr
            if stop is None: stop = self.addr + self.size
            if step is not None:
                for addr in range(*i.indices(self.addr+self.size)):
                    self.uc.mem_write(addr, val[(addr-start)/step])
            else:
                if i.stop: assert len(val) == stop - start
                return self.uc.mem_write(start, val)
        else:
            # This is a slight abuse -- really setting one
            # index should only set one byte -- but it's 
            # more convenient to be able to say 
            #   mem[i] = "hello"
            # to write "hello" at address i
            return self.uc.mem_write(i, val)

class UnicornCPU(object):
    def __init__(self, uc, arch):
        # Prevent infinite recursion
        super(UnicornCPU, self).__setattr__('uc', uc)
        super(UnicornCPU, self).__setattr__('arch', arch)

    def get_reg_name_val(self, name):
        reg_name = 'UC_' + self.arch + '_REG_' + name
        try:
            reg_name_val = globals()[reg_name]
        except KeyError:
            raise AttributeError(item)
        return reg_name_val

    def __getattr__(self, item):
        return self.uc.reg_read(self.get_reg_name_val(item))

    def __setattr__(self, item, value):
        return self.uc.reg_write(self.get_reg_name_val(item), value)

class UnicornSystem(object):
    def __init__(self, arch, mode, memory_start=0x1000000, memory_size=2*1024*1024):
        arch_val = getattr(unicorn, 'UC_ARCH_' + arch)
        mode_val = getattr(unicorn, 'UC_MODE_' + mode)
        self.arch = arch
        self.uc = Uc(arch_val, mode_val)
        self.mem = UnicornMemory(self.uc, memory_start, memory_size)
        self.cpu = UnicornCPU(self.uc, self.arch)

    def start(self, addr, stop):
        self.uc.emu_start(addr, stop)

if __name__ == "__main__":
    X86_CODE32 = b"\x41\x4a" # INC ecx; DEC edx
    sys = UnicornSystem("X86", "32")
    sys.mem[sys.mem.addr] = X86_CODE32
    sys.cpu.ECX = 0x1234
    sys.cpu.EDX = 0x7890
    sys.start(sys.mem.addr, sys.mem.addr + len(X86_CODE32))
    print("Emulation done. Below is the CPU context")
    print(">>> ECX = 0x%x" % sys.cpu.ECX)
    print(">>> EDX = 0x%x" % sys.cpu.EDX)