from unicorn import *
from unicorn.x86_const import *
from capstone import *
from capstone.x86_const import *

BASE_ADDRESS = 0x1000000
STACK_OFFSET = 0x200000

def setup(code_bin):
    """init capstone, return instance"""
    try:
        # Initialize emulator in X86-32bit mode
        mu = Uc(UC_ARCH_X86, UC_MODE_32)

        # map 2MB memory for this emulation
        mu.mem_map(BASE_ADDRESS, 8 * 1024 * 1024)

        # write data to memory
        mu.mem_write(BASE_ADDRESS, code_bin)

        # initialize register for stack 
        mu.reg_write(UC_X86_REG_ESP, BASE_ADDRESS + STACK_OFFSET)
        mu.reg_write(UC_X86_REG_EBP, BASE_ADDRESS + STACK_OFFSET)

    except UcError as e:
        print("ERROR SETUP:%s" % e)
        return None
    return mu


#TODO: calculuate the string size 
def get_string(offset, size=0x20):
    """read string from stack, example: lea     ecx, [ebp+var_44], enter 0x44 """
    return str(emu.mem_read(BASE_ADDRESS + STACK_OFFSET - offset, size)).replace("\x00","") 


def get_code():
    """ read bytes from idb"""
    start = SelStart()
    end = SelEnd()
    length =  end - start
    return "".join([byte for byte in GetManyBytes( SelStart(), length)])


data = get_code()
emu = setup(data)

if emu:
    try:
        emu.emu_start(BASE_ADDRESS, BASE_ADDRESS + len(data))
    except UcError as e:
        print("ERROR START: %s" % e)
    offset = AskLong(0, "Please enter stack offset")
    print get_string(offset)