def fix_InInitializationOrderModuleList(myjit, module_info):
    # first binary is ntdll
    # second binary is kernel32
    olist = []
    ntdll_e = None
    kernel_e = None
    for bname, (addr, e) in module_info.items():
        if bname[::2].lower() == "ntdll.dll":
            ntdll_e = (e, bname, addr)
            continue
        elif bname[::2].lower() == "kernel32.dll":
            kernel_e = (e, bname, addr)
            continue
        elif e == dummy_e:
            d_e = (e, bname, addr)
            continue
        elif e == main_pe:
            continue
        olist.append((e, bname, addr))
    if not ntdll_e or not kernel_e or not d_e:
        log.warn('No kernel ntdll, ldr data will be unconsistant')
    else:
        olist[0:0] = [ntdll_e]
        olist[1:1] = [kernel_e]

    olist.append(d_e)

    last_addr = 0
    for i in xrange(len(olist)):
        e, bname, addr = olist[i]
        p_e, p_bname, p_addr = olist[(i - 1) % len(olist)]
        n_e, n_bname, n_addr = olist[(i + 1) % len(olist)]
        myjit.vm.set_mem(
            addr + 0x10, pck32(n_addr + 0x10) + pck32(p_addr + 0x10))

def build_teb(jitter, teb_address):
    """
    Build TEB informations using following structure:

    +0x000 NtTib                     : _NT_TIB
    +0x01c EnvironmentPointer        : Ptr32 Void
    +0x020 ClientId                  : _CLIENT_ID
    +0x028 ActiveRpcHandle           : Ptr32 Void
    +0x02c ThreadLocalStoragePointer : Ptr32 Void
    +0x030 ProcessEnvironmentBlock   : Ptr32 _PEB
    +0x034 LastErrorValue            : Uint4B
    ...
    @jitter: jitter instance
    @teb_address: the TEB address
    """

    o = ""
    o += pck32(default_seh)
    o += (0x18 - len(o)) * "\x00"
    o += pck32(tib_address)

    o += (0x30 - len(o)) * "\x00"
    o += pck32(peb_address)
    o += pck32(0x11223344)

    jitter.vm.add_memory_page(teb_address, PAGE_READ | PAGE_WRITE, o)

def init_seh(myjit):
    global seh_count
    seh_count = 0
    build_teb(myjit, FS_0_AD)
    build_peb(myjit, peb_address)

    module_info = create_modules_chain(myjit, loaded_modules)
    fix_InLoadOrderModuleList(myjit, module_info)
    fix_InMemoryOrderModuleList(myjit, module_info)
    fix_InInitializationOrderModuleList(myjit, module_info)

    build_ldr_data(myjit, module_info)
    add_process_env(myjit)
    add_process_parameters(myjit)

    myjit.vm.add_memory_page(default_seh, PAGE_READ | PAGE_WRITE, pck32(
        0xffffffff) + pck32(0x41414141) + pck32(0x42424242))

    myjit.vm.add_memory_page(
        context_address, PAGE_READ | PAGE_WRITE, '\x00' * 0x2cc)
    myjit.vm.add_memory_page(
        exception_record_address, PAGE_READ | PAGE_WRITE, '\x00' * 200)

    myjit.vm.add_memory_page(
        FAKE_SEH_B_AD, PAGE_READ | PAGE_WRITE, 0x10000 * "\x00")

def build_peb(jitter, peb_address):
    """
    Build PEB informations using following structure:

    +0x000 InheritedAddressSpace    : UChar
    +0x001 ReadImageFileExecOptions : UChar
    +0x002 BeingDebugged            : UChar
    +0x003 SpareBool                : UChar
    +0x004 Mutant                   : Ptr32 Void
    +0x008 ImageBaseAddress         : Ptr32 Void
    +0x00c Ldr                      : Ptr32 _PEB_LDR_DATA
    +0x010 processparameter

    @jitter: jitter instance
    @peb_address: the PEB address
    """

    offset = peb_address + 8
    o = ""
    if main_pe:
        o += pck32(main_pe.NThdr.ImageBase)
    else:
        offset += 4
    o += pck32(peb_ldr_data_address)
    o += pck32(process_parameters_address)
    jitter.vm.add_memory_page(offset, PAGE_READ | PAGE_WRITE, o)

def fix_InInitializationOrderModuleList(jitter, modules_info):
    """Fix InInitializationOrderModuleList double link list. First module is the
    ntdll, then kernel32. dummy is last pe.

    @jitter: the jitter instance
    @modules_info: the LoadedModules instance

    """

    log.debug("Fix InInitializationOrderModuleList")
    main_pe = modules_info.name2module.get(main_pe_name, None)
    kernel32_pe = modules_info.name2module.get("kernel32.dll", None)
    ntdll_pe = modules_info.name2module.get("ntdll.dll", None)
    dummy_pe = modules_info.name2module.get("", None)
    special_modules = [main_pe, kernel32_pe, ntdll_pe, dummy_pe]
    if not all(special_modules):
        log.warn('No main pe, ldr data will be unconsistant')
        loaded_modules = modules_info.modules
    else:
        loaded_modules = [module for module in modules_info.modules
                          if module not in special_modules]
        loaded_modules[0:0] = [ntdll_pe]
        loaded_modules[1:1] = [kernel32_pe]
        loaded_modules.append(dummy_pe)

    for i, module in enumerate(loaded_modules):
        cur_module_entry = modules_info.module2entry[module]
        prev_module = loaded_modules[(i - 1) % len(loaded_modules)]
        next_module = loaded_modules[(i + 1) % len(loaded_modules)]
        prev_module_entry = modules_info.module2entry[prev_module]
        next_module_entry = modules_info.module2entry[next_module]
        jitter.vm.set_mem(cur_module_entry + 0x10,
                          (pck32(next_module_entry + 0x10) +
                           pck32(prev_module_entry + 0x10)))

def init_seh(jitter):
    """
    Build the modules entries and create double links
    @jitter: jitter instance
    """

    global seh_count
    seh_count = 0
    build_teb(jitter, FS_0_AD)
    build_peb(jitter, peb_address)

    modules_info = create_modules_chain(jitter, name2module)
    fix_InLoadOrderModuleList(jitter, modules_info)
    fix_InMemoryOrderModuleList(jitter, modules_info)
    fix_InInitializationOrderModuleList(jitter, modules_info)

    build_ldr_data(jitter, modules_info)
    add_process_env(jitter)
    add_process_parameters(jitter)

    jitter.vm.add_memory_page(default_seh, PAGE_READ | PAGE_WRITE, pck32(
        0xffffffff) + pck32(0x41414141) + pck32(0x42424242))

    jitter.vm.add_memory_page(
        context_address, PAGE_READ | PAGE_WRITE, '\x00' * 0x2cc)
    jitter.vm.add_memory_page(
        exception_record_address, PAGE_READ | PAGE_WRITE, '\x00' * 200)

    jitter.vm.add_memory_page(
        FAKE_SEH_B_AD, PAGE_READ | PAGE_WRITE, 0x10000 * "\x00")

def add_process_parameters(myjit):
    o = ""
    o += pck32(0x1000)  # size
    o += "E" * (0x48 - len(o))
    o += pck32(process_environment_address)
    myjit.vm.add_memory_page(process_parameters_address,
                             PAGE_READ | PAGE_WRITE,
                             o)

def add_process_parameters(jitter):
    """
    Build a process parameters structure
    @jitter: jitter instance
    """

    o = ""
    o += pck32(0x1000)  # size
    o += "E" * (0x48 - len(o))
    o += pck32(process_environment_address)
    jitter.vm.add_memory_page(process_parameters_address,
                              PAGE_READ | PAGE_WRITE,
                              o)

def regs2ctxt(jitter):
    """
    Build x86_32 cpu context for exception handling
    @jitter: jitload instance
    """

    ctxt = []
    # ContextFlags
    ctxt += [pck32(0x0)]
    # DRX
    ctxt += [pck32(0x0)] * 6
    # Float context
    ctxt += ['\x00' * 112]
    # Segment selectors
    ctxt += [pck32(reg) for reg in (jitter.cpu.GS, jitter.cpu.FS,
                                    jitter.cpu.ES, jitter.cpu.DS)]
    # Gpregs
    ctxt += [pck32(reg) for reg in (jitter.cpu.EDI, jitter.cpu.ESI,
                                    jitter.cpu.EBX, jitter.cpu.EDX,
                                    jitter.cpu.ECX, jitter.cpu.EAX,
                                    jitter.cpu.EBP, jitter.cpu.EIP)]
    # CS
    ctxt += [pck32(jitter.cpu.CS)]
    # Eflags
    # XXX TODO real eflag
    ctxt += [pck32(0x0)]
    # ESP
    ctxt += [pck32(jitter.cpu.ESP)]
    # SS
    ctxt += [pck32(jitter.cpu.SS)]
    return "".join(ctxt)

def set_link_list_entry(jitter, loaded_modules, modules_info, offset):
    for i, module in enumerate(loaded_modules):
        cur_module_entry = modules_info.module2entry[module]
        prev_module = loaded_modules[(i - 1) % len(loaded_modules)]
        next_module = loaded_modules[(i + 1) % len(loaded_modules)]
        prev_module_entry = modules_info.module2entry[prev_module]
        next_module_entry = modules_info.module2entry[next_module]
        if i == 0:
            prev_module_entry = peb_ldr_data_address + 0xC
        if i == len(loaded_modules) - 1:
            next_module_entry = peb_ldr_data_address + 0xC
        jitter.vm.set_mem(cur_module_entry + offset,
                          (pck32(next_module_entry + offset) +
                           pck32(prev_module_entry + offset)))

def return_from_seh(myjit):
    "Handle return after a call to fake seh handler"

    # Get current context
    context_address = upck32(myjit.vm.get_mem(myjit.cpu.ESP + 0x8, 4))
    log.info('Context address: %x', context_address)
    myjit.cpu.ESP = upck32(myjit.vm.get_mem(context_address + 0xc4, 4))
    log.info('New esp: %x', myjit.cpu.ESP)

    # Rebuild SEH
    old_seh = upck32(myjit.vm.get_mem(tib_address, 4))
    new_seh = upck32(myjit.vm.get_mem(old_seh, 4))
    log.info('Old seh: %x New seh: %x', old_seh, new_seh)
    myjit.vm.set_mem(tib_address, pck32(new_seh))

    dump_seh(myjit)

    if myjit.cpu.EAX == 0x0:
        # ExceptionContinueExecution
        ctxt_ptr = context_address
        log.info('Seh continues Context: %x', ctxt_ptr)

        # Get registers changes
        ctxt_str = myjit.vm.get_mem(ctxt_ptr, 0x2cc)
        ctxt2regs(ctxt_str, myjit)
        myjit.pc = myjit.cpu.EIP
        log.info('Context::Eip: %x', myjit.pc)

    elif myjit.cpu.EAX == -1:
        raise NotImplementedError("-> seh try to go to the next handler")

    elif myjit.cpu.EAX == 1:
        # ExceptionContinueSearch
        raise NotImplementedError("-> seh, gameover")

def return_from_seh(jitter):
    """Handle the return from an exception handler
    @jitter: jitter instance"""

    # Get current context
    context_address = upck32(jitter.vm.get_mem(jitter.cpu.ESP + 0x8, 4))
    log.info('Context address: %x', context_address)
    jitter.cpu.ESP = upck32(jitter.vm.get_mem(context_address + 0xc4, 4))
    log.info('New esp: %x', jitter.cpu.ESP)

    # Rebuild SEH
    old_seh = upck32(jitter.vm.get_mem(tib_address, 4))
    new_seh = upck32(jitter.vm.get_mem(old_seh, 4))
    log.info('Old seh: %x New seh: %x', old_seh, new_seh)
    jitter.vm.set_mem(tib_address, pck32(new_seh))

    dump_seh(jitter)

    if jitter.cpu.EAX == 0x0:
        # ExceptionContinueExecution
        ctxt_ptr = context_address
        log.info('Seh continues Context: %x', ctxt_ptr)

        # Get registers changes
        ctxt_str = jitter.vm.get_mem(ctxt_ptr, 0x2cc)
        ctxt2regs(ctxt_str, jitter)
        jitter.pc = jitter.cpu.EIP
        log.info('Context::Eip: %x', jitter.pc)

    elif jitter.cpu.EAX == -1:
        raise NotImplementedError("-> seh try to go to the next handler")

    elif jitter.cpu.EAX == 1:
        # ExceptionContinueSearch
        raise NotImplementedError("-> seh, gameover")

def init_seh(myjit):
    global seh_count
    seh_count = 0
    # myjit.vm.add_memory_page(tib_address, PAGE_READ | PAGE_WRITE,
    # p(default_seh) + p(0) * 11 + p(peb_address))
    myjit.vm.add_memory_page(
        FS_0_AD, PAGE_READ | PAGE_WRITE, build_fake_teb())
    # myjit.vm.add_memory_page(peb_address, PAGE_READ | PAGE_WRITE, p(0) *
    # 3 + p(peb_ldr_data_address))
    myjit.vm.add_memory_page(
        peb_address, PAGE_READ | PAGE_WRITE, build_fake_peb())
    # myjit.vm.add_memory_page(peb_ldr_data_address, PAGE_READ |
    # PAGE_WRITE, p(0) * 3 + p(in_load_order_module_list_address) + p(0) *
    # 0x20)

    """
    ldr_data += "\x00"*(InInitializationOrderModuleList_offset - len(ldr_data))
    ldr_data += build_fake_InInitializationOrderModuleList(loaded_modules)
    ldr_data += "\x00"*(InLoadOrderModuleList_offset - len(ldr_data))
    ldr_data += build_fake_InLoadOrderModuleList(loaded_modules)
    """
    myjit.vm.add_memory_page(
        LDR_AD, PAGE_READ | PAGE_WRITE, "\x00" * MAX_MODULES * 0x1000)
    module_info = create_modules_chain(myjit, loaded_modules)
    fix_InLoadOrderModuleList(myjit, module_info)
    fix_InMemoryOrderModuleList(myjit, module_info)
    fix_InInitializationOrderModuleList(myjit, module_info)

    ldr_data = build_fake_ldr_data(module_info)
    myjit.vm.set_mem(LDR_AD, ldr_data)
    add_process_env(myjit)
    add_process_parameters(myjit)

    # myjit.vm.add_memory_page(in_load_order_module_list_address,
    #     PAGE_READ | PAGE_WRITE, p(0) * 40)
    # myjit.vm.add_memory_page(in_load_order_module_list_address,
    #     PAGE_READ | PAGE_WRITE, build_fake_inordermodule(loaded_modules))
    myjit.vm.add_memory_page(default_seh, PAGE_READ | PAGE_WRITE, pck32(
        0xffffffff) + pck32(0x41414141) + pck32(0x42424242))

    myjit.vm.add_memory_page(
        context_address, PAGE_READ | PAGE_WRITE, '\x00' * 0x2cc)
    myjit.vm.add_memory_page(
        exception_record_address, PAGE_READ | PAGE_WRITE, '\x00' * 200)

    myjit.vm.add_memory_page(
        FAKE_SEH_B_AD, PAGE_READ | PAGE_WRITE, 0x10000 * "\x00")

 def test_init(self):
     init_regs(self)
     self.buf = ""
     for reg_name in reversed(["EAX", "ECX",
                               "EDX", "EBX",
                               "ESP", "EBP",
                               "ESI", "EDI"]):
         self.buf += pck32(getattr(self.myjit.cpu, reg_name))

def fix_InMemoryOrderModuleList(myjit, module_info):
    log.debug("Fix InMemoryOrderModuleList")
    # first binary is PE
    # last is dumm_e
    olist = []
    m_e = None
    d_e = None
    for m in [main_pe_name, ""] + loaded_modules:

        if isinstance(m, tuple):
            fname, e = m
        else:
            fname, e = m, None

        if "/" in fname:
            fname = fname[fname.rfind("/") + 1:]
        bname_str = fname
        bname = '\x00'.join(bname_str) + '\x00'
        if not bname.lower() in module_info:
            log.warn('Module not found, ldr data will be unconsistant')
            continue
        addr, e = module_info[bname.lower()]
        log.debug(bname_str)
        if e == main_pe:
            m_e = (e, bname, addr)
            continue
        elif e == dummy_e:
            d_e = (e, bname, addr)
            continue
        olist.append((e, bname, addr))
    if not m_e or not d_e:
        log.warn('No main pe, ldr data will be unconsistant')
    else:
        olist[0:0] = [m_e]
    olist.append(d_e)

    last_addr = 0

    for i in xrange(len(olist)):
        e, bname, addr = olist[i]
        p_e, p_bname, p_addr = olist[(i - 1) % len(olist)]
        n_e, n_bname, n_addr = olist[(i + 1) % len(olist)]
        myjit.vm.set_mem(
            addr + 0x8, pck32(n_addr + 0x8) + pck32(p_addr + 0x8))

def regs2ctxt(regs):
    ctxt = ""
    ctxt += '\x00\x00\x00\x00'  # ContextFlags
    ctxt += '\x00\x00\x00\x00' * 6  # drX
    ctxt += '\x00' * 112  # float context
    ctxt += '\x00\x00\x00\x00' + '\x3b\x00\x00\x00' + \
        '\x23\x00\x00\x00' + '\x23\x00\x00\x00'  # segment selectors
    ctxt += pck32(regs['EDI']) + pck32(regs['ESI']) + pck32(regs['EBX']) + \
        pck32(regs['EDX']) + pck32(regs['ECX']) + pck32(regs['EAX']) + \
        pck32(regs['EBP']) + pck32(regs['EIP'])  # gpregs
    ctxt += '\x23\x00\x00\x00'  # cs
    ctxt += '\x00\x00\x00\x00'  # eflags
    ctxt += pck32(regs['ESP'])  # esp
    ctxt += '\x23\x00\x00\x00'  # ss segment selector
    return ctxt

def build_ldr_data(jitter, modules_info):
    """
    Build Loader informations using following structure:

    +0x000 Length                          : Uint4B
    +0x004 Initialized                     : UChar
    +0x008 SsHandle                        : Ptr32 Void
    +0x00c InLoadOrderModuleList           : _LIST_ENTRY
    +0x014 InMemoryOrderModuleList         : _LIST_ENTRY
    +0x01C InInitializationOrderModuleList         : _LIST_ENTRY
    # dummy dll base
    +0x024 DllBase : Ptr32 Void

    @jitter: jitter instance
    @modules_info: LoadedModules instance

    """
    # ldr offset pad
    offset = 0xC
    addr = LDR_AD + peb_ldr_data_offset
    ldrdata = PEB_LDR_DATA(jitter.vm, addr)

    main_pe = modules_info.name2module.get(main_pe_name, None)
    ntdll_pe = modules_info.name2module.get("ntdll.dll", None)


    size = 0
    if main_pe:
        size += ListEntry.sizeof() * 2
        main_addr_entry = modules_info.module2entry[main_pe]
    if ntdll_pe:
        size += ListEntry.sizeof()
        ntdll_addr_entry = modules_info.module2entry[ntdll_pe]

    jitter.vm.add_memory_page(addr + offset, PAGE_READ | PAGE_WRITE,
                              "\x00" * size,
                              "Loader struct")  # (ldrdata.get_size() - offset))

    if main_pe:
        ldrdata.InLoadOrderModuleList.flink = main_addr_entry
        ldrdata.InLoadOrderModuleList.blink = 0

        ldrdata.InMemoryOrderModuleList.flink = main_addr_entry + \
            LdrDataEntry.get_type().get_offset("InMemoryOrderLinks")
        ldrdata.InMemoryOrderModuleList.blink = 0

    if ntdll_pe:
        ldrdata.InInitializationOrderModuleList.flink = ntdll_addr_entry + \
            LdrDataEntry.get_type().get_offset("InInitializationOrderLinks")
        ldrdata.InInitializationOrderModuleList.blink = 0

    # Add dummy dll base
    jitter.vm.add_memory_page(peb_ldr_data_address + 0x24,
                              PAGE_READ | PAGE_WRITE, pck32(0),
                              "Loader struct dummy dllbase")

def build_fake_teb():
    """
    +0x000 NtTib                     : _NT_TIB
    +0x01c EnvironmentPointer        : Ptr32 Void
    +0x020 ClientId                  : _CLIENT_ID
    +0x028 ActiveRpcHandle           : Ptr32 Void
    +0x02c ThreadLocalStoragePointer : Ptr32 Void
    +0x030 ProcessEnvironmentBlock   : Ptr32 _PEB
    +0x034 LastErrorValue            : Uint4B
    ...
    """
    o = ""
    o += pck32(default_seh)
    o += (0x18 - len(o)) * "\x00"
    o += pck32(tib_address)

    o += (0x30 - len(o)) * "\x00"
    o += pck32(peb_address)
    o += pck32(0x11223344)

    return o

def build_teb(myjit, teb_address):
    """
    +0x000 NtTib                     : _NT_TIB
    +0x01c EnvironmentPointer        : Ptr32 Void
    +0x020 ClientId                  : _CLIENT_ID
    +0x028 ActiveRpcHandle           : Ptr32 Void
    +0x02c ThreadLocalStoragePointer : Ptr32 Void
    +0x030 ProcessEnvironmentBlock   : Ptr32 _PEB
    +0x034 LastErrorValue            : Uint4B
    ...
    """
    o = ""
    o += pck32(default_seh)
    o += (0x18 - len(o)) * "\x00"
    o += pck32(tib_address)

    o += (0x30 - len(o)) * "\x00"
    o += pck32(peb_address)
    o += pck32(0x11223344)

    myjit.vm.add_memory_page(teb_address, PAGE_READ | PAGE_WRITE, o)

def build_peb(myjit, peb_address):
    """
    +0x000 InheritedAddressSpace    : UChar
    +0x001 ReadImageFileExecOptions : UChar
    +0x002 BeingDebugged            : UChar
    +0x003 SpareBool                : UChar
    +0x004 Mutant                   : Ptr32 Void
    +0x008 ImageBaseAddress         : Ptr32 Void
    +0x00c Ldr                      : Ptr32 _PEB_LDR_DATA
    +0x010 processparameter
    """

    offset = peb_address + 8
    o = ""
    if main_pe:
        o += pck32(main_pe.NThdr.ImageBase)
    else:
        offset += 4
    o += pck32(peb_ldr_data_address)
    o += pck32(process_parameters_address)
    myjit.vm.add_memory_page(offset, PAGE_READ | PAGE_WRITE, o)