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 """
    read_str = str(emu.mem_read(BASE_ADDRESS + STACK_OFFSET - offset, size))
    temp = read_str[:read_str.find("\x00\x00")]
    read_str = temp.replace("\x00","")
    return read_str


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


start,end, data = get_code()
if start:
    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")
        comment = get_string(offset)
        print "0x%x, %s" % (PrevHead(end), comment)
        MakeComm(PrevHead(end), comment)