# plugins

repository: DockTree/mefq
file: qiling/extensions/idaplugin/qilingida.py
wanted_name: Qiling Emulator
created_at: 2023-10-22 02:13:08+00:00
pushed_at: 2023-12-14 10:43:25+00:00
url:
https://github.com/DockTree/mefq/blob/46efdd83ebd5db0b6331f800ff1bedc87b9cbf88/qiling/extensions/idaplugin/qilingida.py

content:
# Cross Platform and Multi Architecture Advanced Binary Emulation Framework Plugin For IDA
# Learn how to use? Please visit https://docs.qiling.io/en/latest/ida/

UseAsScript = True

import sys
import collections
import time
import struct
import re
import logging
import shlex
from enum import Enum
from elftools.elf.elffile import ELFFile
from json import load

# IDA Python SDK
from idaapi import *
from idc import *
from idautils import *
import idc
import ida_ua
import ida_idaapi
import ida_funcs
import ida_nalt
import idautils
import ida_xref
import ida_kernwin
import ida_ida
import ida_bytes
import ida_segment
import ida_name
import ida_gdl
import ida_frame
import ida_idp
import ida_auto
import ida_netnode
import ida_hexrays
import ida_range
# PyQt
from PyQt5 import QtCore, QtWidgets
from PyQt5.QtWidgets import (QPushButton, QHBoxLayout)

# Qiling
from qiling import Qiling
from qiling.const import *
from qiling.arch.x86_const import reg_map_32 as x86_reg_map_32
from qiling.arch.x86_const import reg_map_64 as x86_reg_map_64
from qiling.arch.x86_const import reg_map_misc as x86_reg_map_misc
from qiling.arch.x86_const import reg_map_st as x86_reg_map_st
from qiling.arch.arm_const import reg_map as arm_reg_map
from qiling.arch.arm64_const import reg_map as arm64_reg_map
from qiling.arch.mips_const import reg_map as mips_reg_map
from qiling import __version__ as QLVERSION
from qiling.os.filestruct import ql_file
from keystone import *


QilingHomePage = 'https://www.qiling.io'
QilingStableVersionURL = 'https://raw.githubusercontent.com/qilingframework/qiling/master/qiling/__version__.py'
logging.basicConfig(level=logging.INFO, format='[%(levelname)s][%(module)s:%(lineno)d] %(message)s')

class Colors(Enum):
    Blue = 0xE8864A
    Pink = 0xC0C0FB
    White = 0xFFFFFF
    Black = 0x000000
    Green = 0xd3ead9
    Gray = 0xd9d9d9
    Beige = 0xCCF2FF

class IDA:
    def __init__(self):
        pass

    @staticmethod
    def get_function(addr):
        return ida_funcs.get_func(addr)

    @staticmethod
    def get_function_start(addr):
        return IDA.get_function(addr).start_ea

    @staticmethod
    def get_function_end(addr):
        return IDA.get_function(addr).end_ea

    @staticmethod
    def get_function_framesize(addr):
        return IDA.get_function(addr).frsize

    @staticmethod
    def get_function_name(addr):
        return ida_funcs.get_func_name(addr)

    @staticmethod
    def get_functions():
        return [IDA.get_function(func) for func in idautils.Functions()]

    @staticmethod
    def set_color(addr, what, color):
        return idc.set_color(addr, what, color)

    @staticmethod
    def color_block(bb, color):
        for i in range(bb.start_ea, bb.end_ea):
            IDA.set_color(i, idc.CIC_ITEM, color)

    # note:
    # corresponds to IDA graph view
    # a good example to iterate the graph
    # https://github.com/idapython/src/blob/bc9b51b1c70083815a574a57b7a783698de3698d/examples/core/dump_flowchart.py
    # arg can be a function or a (start, end) tuple or an address in the function
    @staticmethod
    def get_flowchart(arg):
        if type(arg) is int:
            func = IDA.get_function(arg)
            if func is None:
                return None
            return ida_gdl.FlowChart(func)
        return ida_gdl.FlowChart(arg)

    @staticmethod
    def get_block(addr):
        flowchart = IDA.get_flowchart(addr)
        for bb in flowchart:
            if bb.start_ea <= addr and addr < bb.end_ea:
                return bb
        return None

    @staticmethod
    def block_is_terminating(bb):
        # fcb_ret: has a retn instruction in the end
        # fcb_noret: in most cases, exit() is called
        # fcb_indjump: jmp $eax
        if (bb.type == ida_gdl.fcb_ret or bb.type == ida_gdl.fcb_noret or
                (bb.type == ida_gdl.fcb_indjump and len(list(bb.succs())) == 0)):
            return True
        for b in bb.succs():
            if b.type == ida_gdl.fcb_extern:
                return True
        return False

    @staticmethod
    def get_starting_block(addr):
        flowchart = IDA.get_flowchart(addr)
        if flowchart is None:
            return None
        func = IDA.get_function(addr)
        for bb in flowchart:
            if bb.start_ea == func.start_ea:
                return bb
        return None

    @staticmethod
    def get_terminating_blocks(addr):
        flowchart = IDA.get_flowchart(addr)
        return [bb for bb in flowchart if IDA.block_is_terminating(bb)]

    @staticmethod
    def get_prev_head(addr, minea=0):
        return ida_bytes.prev_head(addr, minea)

    @staticmethod
    def get_segments():
        r = []
        seg = ida_segment.get_first_seg()
        while seg is not None:
            r.append(seg)
            seg = ida_segment.get_next_seg(seg.start_ea)
        return r

    @staticmethod
    def get_segment_name(s):
        return ida_segment.get_segm_name(s)

    @staticmethod
    def get_segment_by_name(name):
        return ida_segment.get_segm_by_name(name)

    @staticmethod
    def __addr_in_seg(addr):
        segs = IDA.get_segments()
        for seg in segs:
            if addr < seg.end_ea and addr >= seg.start_ea:
                return seg
        return None

    # note: accept name and address in the segment
    @staticmethod
    def get_segment(arg):
        if type(arg) is int:
            return IDA.__addr_in_seg(arg)
        else: # str
            return IDA.get_segment_by_name(arg)

    @staticmethod
    def get_segment_start(arg):
        seg = IDA.get_segment(arg)
        if seg is not None:
            return seg.start_ea
        return None

    @staticmethod
    def get_segment_end(arg):
        seg = IDA.get_segment(arg)
        if seg is not None:
            return seg.end_ea
        return None

    @staticmethod
    def get_segment_perm(arg):
        seg = IDA.get_segment(arg)
        if seg is not None:
            return seg.perm # RWX e.g. 0b101 = R + X
        return None

    @staticmethod
    def get_segment_type(arg):
        seg = IDA.get_segment(arg)
        if seg is not None:
            return seg.type # 0x1 SEG_DATA 0x2 SEG_CODE See doc for details
        return None

    @staticmethod
    def get_instruction(addr):
        r = ida_ua.print_insn_mnem(addr)
        if r == "":
            return None
        return r

    # immidiate value
    @staticmethod
    def get_operand(addr, n):
        return (idc.get_operand_type(addr, n), idc.get_operand_value(addr, n))

    # eax, ecx, etc
    @staticmethod
    def print_operand(addr, n):
        return idc.print_operand(addr, n)

    @staticmethod
    def get_instruction_size(addr):
        return ida_bytes.get_item_size(addr)

    @staticmethod
    def get_instructions_count(begin, end):
        p = begin
        cnt = 0
        while p < end:
            sz = IDA.get_instruction_size(p)
            cnt += 1
            p += sz
        return cnt

    @staticmethod
    def get_name(addr):
        return ida_name.get_name(addr)

    @staticmethod
    def get_name_address(name, addr=0):
        return ida_name.get_name_ea(addr, name)

    @staticmethod
    def get_bytes(addr, l):
        return ida_bytes.get_bytes(addr, l)

    @staticmethod
    def get_byte(addr):
        return ida_bytes.get_byte(addr)

    @staticmethod
    def get_word(addr):
        return ida_bytes.get_word(addr)

    @staticmethod
    def get_dword(addr):
        return ida_bytes.get_dword(addr)

    @staticmethod
    def get_qword(addr):
        return ida_bytes.get_qword(addr)

    @staticmethod
    def get_xrefsto(addr, flags=ida_xref.XREF_ALL):
        return [ref.frm for ref in idautils.XrefsTo(addr, flags)]

    @staticmethod
    def get_xrefsfrom(addr, flags=ida_xref.XREF_ALL):
        return [ref.frm for ref in idautils.XrefsFrom(addr, flags)]

    @staticmethod
    def get_input_file_path():
        return ida_nalt.get_input_file_path()

    @staticmethod
    def get_info_structure():
        return ida_idaapi.get_inf_structure()

    @staticmethod
    def get_main_address():
        return IDA.get_info_structure().main

    @staticmethod
    def get_max_address():
        return IDA.get_info_structure().max_ea

    @staticmethod
    def get_min_address():
        return IDA.get_info_structure().min_ea

    @staticmethod
    def is_big_endian():
        return IDA.get_info_structure().is_be()

    @staticmethod
    def is_little_endian():
        return not IDA.is_big_endian()

    @staticmethod
    def get_filetype():
        info = IDA.get_info_structure()
        ftype = info.filetype
        if ftype == ida_ida.f_MACHO:
            return "macho"
        elif ftype == ida_ida.f_PE or ftype == ida_ida.f_EXE or ftype == ida_ida.f_EXE_old: # is this correct?
            return "pe"
        elif ftype == ida_ida.f_ELF:
            return "elf"
        else:
            return None

    @staticmethod
    def get_ql_arch_string():
        info = IDA.get_info_structure()
        proc = info.procname.lower()
        result = None
        if proc == "metapc":
            result = "x86"
            if info.is_64bit():
                result = "x8664"
        elif "mips" in proc:
            result = "mips"
        elif "arm" in proc:
            result = "arm32"
            if info.is_64bit():
                result = "arm64"
        # That's all we support :(
        return result

    @staticmethod
    def get_current_address():
        return ida_kernwin.get_screen_ea()

    # return (?, start, end)
    @staticmethod
    def get_last_selection():
        return ida_kernwin.read_range_selection(None)

    # Use with skipcalls
    # note that the address is the end of target instruction
    # e.g.:
    # 0x1 push eax
    # 0x4 mov eax, 0
    # call get_frame_sp_delta(0x4) and get -4.
    @staticmethod
    def get_frame_sp_delta(addr):
        return ida_frame.get_sp_delta(IDA.get_function(addr), addr)

    @staticmethod
    def patch_bytes(addr, bs):
        return ida_bytes.patch_bytes(addr, bs)

    @staticmethod
    def fill_bytes(start, end, bs = b'\x90'):
        return ida_bytes.patch_bytes(start, bs*(end-start))

    @staticmethod
    def nop_selection():
        _, start, end = IDA.get_last_selection()
        return IDA.fill_bytes(start, end)

    @staticmethod
    def fill_block(bb, bs=b'\x90'):
        return IDA.fill_bytes(bb.start_ea, bb.end_ea, bs)

    @staticmethod
    def assemble(ea, cs, ip, use32, line):
        return ida_idp.assemble(ea, cs, ip, use32, line)

    @staticmethod
    def create_data(ea, dataflag, size, tid=ida_netnode.BADNODE):
        return ida_bytes.create_data(ea, dataflag, size, tid)

    @staticmethod
    def create_bytes_array(start, end):
        return IDA.create_data(start, ida_bytes.byte_flag(), end-start)

    @staticmethod
    def create_byte(ea, length, force=False):
        return ida_bytes.create_byte(ea, length, force)

    @staticmethod
    def perform_analysis(start, end, final_pass=True):
        return ida_auto.plan_and_wait(start, end)

    @staticmethod
    def get_item_head(ea):
        return ida_bytes.get_item_head(ea)

    @staticmethod
    def get_item_end(ea):
        return ida_bytes.get_item_end(ea)

    @staticmethod
    def get_item_size(ea):
        return ida_bytes.get_item_size(ea)

    @staticmethod
    def get_item(ea):
        return (IDA.get_item_head(ea), IDA.get_item_end(ea))

    @staticmethod
    def is_colored_item(ea):
        return ida_nalt.is_colored_item(ea)

    # NOTE: The [start, end) range should include all control flows except long calls.
    @staticmethod
    def get_micro_code_mba(start, end, decomp_flags=ida_hexrays.DECOMP_WARNINGS, maturity=7):
        mbrgs = ida_hexrays.mba_ranges_t()
        rg = ida_range.range_t(start, end)
        mbrgs.ranges.push_back(rg)
        fl = ida_hexrays.hexrays_failure_t()
        mba = ida_hexrays.gen_microcode(mbrgs, fl, None, decomp_flags, maturity)
        if mba is None:
            logging.error(f"Fail to get mba because: {fl}")
        return mba

    @staticmethod
    def micro_code_from_mbb(mbb):
        cur = mbb.head
        while cur is not None:
            yield cur
            cur = cur.next
        return

### View Class

class QlEmuRegView(simplecustviewer_t):
    def __init__(self, ql_emu_plugin):
        super(QlEmuRegView, self).__init__()
        self.hooks = None
        self.ql_emu_plugin = ql_emu_plugin

    def Create(self):
        title = "QL Register View"
        if not simplecustviewer_t.Create(self, title):
            return False

        self.menu_update = 1

        class Hooks(UI_Hooks):
            class PopupActionHandler(action_handler_t):
                def __init__(self, subview, menu_id):
                    action_handler_t.__init__(self)
                    self.subview = subview
                    self.menu_id = menu_id

                def activate(self, ctx):
                    self.subview.OnPopupMenu(self.menu_id)

                def update(self, ctx):
                    return AST_ENABLE_ALWAYS

            def __init__(self, form):
                UI_Hooks.__init__(self)
                self.form = form

            def finish_populating_widget_popup(self, widget, popup):
                if self.form.title == get_widget_title(widget):
                    attach_dynamic_action_to_popup(widget, popup, action_desc_t(None, "Edit Register", self.PopupActionHandler(self.form, self.form.menu_update),  None, None, -1))

        if self.hooks is None:
            self.hooks = Hooks(self)
            self.hooks.hook()

        return True

    def SetReg(self, addr, ql:Qiling):
        #clear
        self.ClearLines()

        view_title = COLSTR("Reg value at { ", SCOLOR_AUTOCMT)
        view_title += COLSTR("IDA Address:0x%X | QL Address:0x%X" % (addr - self.ql_emu_plugin.qlemu.baseaddr + get_imagebase(), addr), SCOLOR_DREF)
        # TODO: Add disass should be better
        view_title += COLSTR(" }", SCOLOR_AUTOCMT)
        self.AddLine(view_title)
        self.AddLine("")

        reglist = QlEmuMisc.get_reg_map(ql)
        line = ""
        cols = 3
        reglist = [reglist[i:i+cols] for i in range(0,len(reglist),cols)]
        for regs in reglist:
            for reg in regs:
                line += COLSTR(" %4s: " % str(reg), SCOLOR_REG)
                regvalue = ql.arch.regs.read(reg)
                if arch in [QL_ARCH.X8664, QL_ARCH.ARM64]:
                    value_format = "0x%.16X"
                else:
                    value_format = "0x%.8X"
                line += COLSTR(str(value_format % regvalue), SCOLOR_NUMBER)
                # TODO: ljust will looks better
            self.AddLine(line)
            line = ''
        self.AddLine(line)
        self.Refresh()

    def OnPopupMenu(self, menu_id):
        if menu_id == self.menu_update:
            self.ql_emu_plugin.ql_chang_reg()

    def OnClose(self):
        if self.hooks:
            self.hooks.unhook()
            self.hooks = None
        self.ql_emu_plugin.ql_close_reg_view()

class QlEmuStackView(simplecustviewer_t):
    def __init__(self, ql_emu_plugin):
        super(QlEmuStackView, self).__init__()
        self.ql_emu_plugin = ql_emu_plugin

    def Create(self):
        title = "QL Stack View"
        if not simplecustviewer_t.Create(self, title):
            return False
        return True

    def SetStack(self, ql:Qiling):
        self.ClearLines()
        if ql is None:
            return

        sp = ql.arch.regs.arch_sp
        self.AddLine('')
        self.AddLine(COLSTR('  Stack at 0x%X' % sp, SCOLOR_AUTOCMT))
        self.AddLine('')

        reg_bit_size = ql.arch.bits
        reg_byte_size = reg_bit_size // 8
        value_format = '% .16X' if reg_bit_size == 64 else '% .8X'

        for i in range(-30, 30):
            clr = SCOLOR_DREF if i < 0 else SCOLOR_INSN
            cur_addr = (sp + i * reg_byte_size)
            line = ('  ' + value_format + ': ') % cur_addr
            try:
                value = ql.mem.read(cur_addr, reg_byte_size)
                value, = struct.unpack('Q' if reg_bit_size == 64 else 'I', value)
                line += value_format % value
            except Exception:
                line += '?' * reg_byte_size * 2

            self.AddLine(COLSTR(line, clr))

    def OnClose(self):
        self.ql_emu_plugin.ql_close_stack_view()

class QlEmuMemView(simplecustviewer_t):
    def __init__(self, ql_emu_plugin, addr, size):
        super(QlEmuMemView, self).__init__()
        self.ql_emu_plugin = ql_emu_plugin
        self.viewid = addr
        self.addr = addr
        self.size = size
        self.lastContent = []

    def Create(self, title):
        if not simplecustviewer_t.Create(self, title):
            return False
        return True

    def SetMem(self, ql:Qiling):
        self.ClearLines()

        if ql is None:
            return

        try:
            memory = ql.mem.read(self.addr, self.size)
        except:
            return

        size = len(memory)

        view_title = COLSTR("  Memory at [ ", SCOLOR_AUTOCMT)
        view_title += COLSTR("0x%X: %d byte(s)" % (self.addr, size), SCOLOR_DREF)
        view_title += COLSTR(" ]", SCOLOR_AUTOCMT)
        self.AddLine(str(view_title))
        self.AddLine("")
        self.AddLine(COLSTR("                0  1  2  3  4  5  6  7  8  9  A  B  C  D  E  F", SCOLOR_AUTOCMT))

        startAddress = self.addr
        line = ""
        chars = ""
        get_char = lambda byte: chr(byte) if 0x20 <= byte <= 0x7E else '.'

        if size != 0:
            for x in range(size):
                if x%16==0:
                    line += COLSTR(" %.12X: " % startAddress, SCOLOR_AUTOCMT)
                if len(self.lastContent) == len(memory):
                    if memory[x] != self.lastContent[x]:
                        line += COLSTR(str("%.2X " % memory[x]), SCOLOR_VOIDOP)
                        chars += COLSTR(get_char(memory[x]), SCOLOR_VOIDOP)
                    else:
                        line += COLSTR(str("%.2X " % memory[x]), SCOLOR_NUMBER)
                        chars += COLSTR(get_char(memory[x]), SCOLOR_NUMBER)
                else:
                    line += COLSTR(str("%.2X " % memory[x]), SCOLOR_NUMBER)
                    chars += COLSTR(get_char(memory[x]), SCOLOR_NUMBER)

                if (x+1)%16==0:
                    line += "  " + chars
                    self.AddLine(line)
                    startAddress += 16
                    line = ""
                    chars = ""

            # add padding
            tail = 16 - size%16
            if tail != 0:
                for x in range(tail): line += "   "
                line += "  " + chars
                self.AddLine(line)

        self.Refresh()
        self.lastContent = memory

    def OnClose(self):
        self.ql_emu_plugin.ql_close_mem_view(self.viewid)

### Dialog Class
class QlEmuMemDialog(Form):
    def __init__(self):
        Form.__init__(self, r"""STARTITEM {id:mem_addr}
BUTTON YES* Add
BUTTON CANCEL Cancel
Show Memory Range
Specify start address and size of new memory range.
<##Address\::{mem_addr}> <##Size\::{mem_size}>
<##Comment\::{mem_cmnt}>
""", {
        'mem_addr': Form.NumericInput(swidth=20, tp=Form.FT_HEX),
        'mem_size': Form.NumericInput(swidth=10, tp=Form.FT_DEC),
        'mem_cmnt': Form.StringInput(swidth=41)
    })

class QlEmuSetupDialog(Form):
    def __init__(self):
        Form.__init__(self, r"""STARTITEM {id:path_name}
BUTTON YES* Start
BUTTON CANCEL Cancel
Setup Qiling
<#Select Rootfs to open#Rootfs path\:        :{path_name}>
<#Custom script path   #Custom script path\: :{script_name}>
<#Custom parameter   #Custom parameter\: :{parameters}>
<#Custom env   #Custom env\: :{env_var}>
""", {
            'path_name': Form.DirInput(swidth=50),
            'script_name': Form.FileInput(swidth=50, open=True),
            'parameters': Form.StringInput(swidth=50),
            'env_var': Form.FileInput(swidth=70, open=True),
        })

class QlEmuSaveDialog(Form):
    def __init__(self):
        Form.__init__(self, r"""STARTITEM {id:path_name}
BUTTON YES* Save
BUTTON CANCEL Cancel
Save Path
<#Save to#Path\::{path_name}>
""", {
        'path_name': Form.FileInput(swidth=50, save=True),
    })

class QlEmuLoadDialog(Form):
    def __init__(self):
        Form.__init__(self, r"""STARTITEM {id:file_name}
BUTTON YES* Load
BUTTON CANCEL Cancel
Load File
<#Load From#File\::{file_name}>
""", {
        'file_name': Form.FileInput(swidth=50, open=True)
    })

class QlEmuAboutDialog(Form):
    def __init__(self, version):
        super(QlEmuAboutDialog, self).__init__(
            r"""STARTITEM 0
BUTTON YES* Open Qiling Website
Qiling:: About
            {FormChangeCb}
            Qiling IDA plugin v%s with Qiling Framework v%s.
            Author:
            Ziqiao Kong,
            Chenxu Wu,
            Qiling Team.
            Qiling is released under the GPL v2.
            Find more info at https://www.qiling.io
            """ %(version, QLVERSION), {
            'FormChangeCb': self.FormChangeCb(self.OnFormChange),
            })

        self.Compile()

    # callback to be executed when any form control changed
    def OnFormChange(self, fid):
        if fid == -2:   # Goto homepage
            import webbrowser
            # open Keypatch homepage in a new tab, if possible
            webbrowser.open(QilingHomePage, new = 2)

        return 1

class QlEmuUpdateDialog(Form):
    def __init__(self, version, message):
        super(QlEmuUpdateDialog, self).__init__(
            r"""STARTITEM 0
BUTTON YES* Open Qiling Website
Qiling:: Check for update
            {FormChangeCb}
            Your Qiling is v%s
            %s
            """ %(version, message), {
            'FormChangeCb': self.FormChangeCb(self.OnFormChange),
            })
        self.Compile()

    # callback to be executed when any form control changed
    def OnFormChange(self, fid):
        if fid == -2:   # Goto homepage
            import webbrowser
            # open Keypatch homepage in a new tab, if possible
            webbrowser.open(QilingHomePage, new = 2)

        return 1

class QlEmuRegEditDialog(Form):
    def __init__(self, regName):
        Form.__init__(self, r"""STARTITEM {id:reg_val}
BUTTON YES* Save
BUTTON CANCEL Cancel
Register Value
{reg_label}
<##:{reg_val}>
""", {
        'reg_label': Form.StringLabel("Edit [ " + regName + " ] value"),
        'reg_val': Form.NumericInput(tp=Form.FT_HEX, swidth=20)
        })

class QlEmuRegDialog(Choose):
    def __init__(self, reglist, flags=0, width=None, height=None, embedded=False):
        Choose.__init__(
            self, "QL Register Edit",
            [ ["Register", 10 | Choose.CHCOL_PLAIN],
              ["Value", 30] ])
        self.popup_names = ["", "", "Edit Value", ""]
        self.items = reglist

    def show(self):
        return self.Show(True) >= 0

    def OnEditLine(self, n):
        edit_dlg = QlEmuRegEditDialog(self.items[n][0])
        edit_dlg.Compile()
        edit_dlg.reg_val.value = self.items[n][1]
        ok = edit_dlg.Execute()
        if ok == 1:
            newvalue = edit_dlg.reg_val.value
            self.items[n][1] = int("%X" % newvalue, 16)
        self.Refresh()

    def OnGetLine(self, n):
        if self.items[n][2] == 32:
            return [ self.items[n][0], "0x%08X" % self.items[n][1] ]
        if self.items[n][2] == 64:
            return [ self.items[n][0], "0x%16X" % self.items[n][1] ]

    def OnGetSize(self):
        return len(self.items)

    def OnClose(self):
        pass

### Misc
class QlEmuMisc:
    MenuItem = collections.namedtuple("MenuItem", ["action", "handler", "title", "tooltip", "shortcut", "popup"])
    class menu_action_handler(action_handler_t):
        def __init__(self, handler, action):
            action_handler_t.__init__(self)
            self.action_handler = handler
            self.action_type = action

        def activate(self, ctx):
            if ctx.widget_type == BWN_DISASM:
                self.action_handler.ql_handle_menu_action(self.action_type)
            return 1

        # This action is always available.
        def update(self, ctx):
            return AST_ENABLE_ALWAYS

    @staticmethod
    def get_reg_map(ql:Qiling):
        tables = {
            QL_ARCH.X86     : list({**x86_reg_map_32, **x86_reg_map_misc, **x86_reg_map_st}.keys()),
            QL_ARCH.X8664   : list({**x86_reg_map_64, **x86_reg_map_misc, **x86_reg_map_st}.keys()),
            QL_ARCH.ARM     : list({**arm_reg_map}.keys()),
            QL_ARCH.ARM64   : list({**arm64_reg_map}.keys()),
            QL_ARCH.MIPS    : list({**mips_reg_map}.keys()),
        }

        return tables.get(ql.arch.type, [])

    @staticmethod
    def url_download(url):
        try:
            from urllib2 import Request, urlopen, URLError, HTTPError
        except:
            from urllib.request import Request, urlopen
            from urllib.error import URLError, HTTPError

        # create the url and the request
        req = Request(url)

        # Open the url
        try:
            # download this URL
            f = urlopen(req)
            content = f.read()
            return (0, content)

        # handle errors
        except HTTPError as e:
            # print "HTTP Error:", e.code , url
            # fail to download this file
            return (1, None)
        except URLError as e:
            # print "URL Error:", e.reason , url
            # fail to download this file
            return (1, None)
        except Exception as e:
            # fail to save the downloaded file
            # print("Error:", e)
            return (2, None)

    class QLStdIO(ql_file):
        def __init__(self, path, fd):
            super().__init__(path, fd)
            self.__fd = fd

        def write(self, write_buf):
            super().write(write_buf)
            msg(write_buf.decode('utf-8'))

        def flush(self):
            pass

        def isatty(self):
            return False

### Qiling wrapper

class QlEmuQiling:
    def __init__(self):
        self.path = None
        self.rootfs = None
        self.ql: Qiling = None
        self.status = None
        self.exit_addr = None
        self.baseaddr = None
        self.env = {}

    def start(self, *args, **kwargs):
        self.ql = Qiling(argv=self.path, rootfs=self.rootfs, verbose=QL_VERBOSE.DEBUG, env=self.env, log_plain=True, *args, **kwargs)

        if sys.platform != 'win32':
            self.ql.os.stdin = QlEmuMisc.QLStdIO('stdin', sys.__stdin__.fileno())
            self.ql.os.stdout = QlEmuMisc.QLStdIO('stdout', sys.__stdout__.fileno())
            self.ql.os.stderr = QlEmuMisc.QLStdIO('stderr', sys.__stderr__.fileno())

        self.exit_addr = self.ql.os.exit_point
        if self.ql.os.type == QL_OS.LINUX:
            f = open(self.ql.path, 'rb')
            elffile = ELFFile(f)
            elf_header = elffile.header
            if elf_header['e_type'] == 'ET_EXEC':
                self.baseaddr = self.ql.os.elf_mem_start
            elif elf_header['e_type'] == 'ET_DYN':
                if self.ql.arch.bits == 32:
                    self.baseaddr = int(self.ql.os.profile.get("OS32", "load_address"), 16)
                elif self.ql.arch.bits == 64:
                    self.baseaddr = int(self.ql.os.profile.get("OS64", "load_address"), 16)
        else:
            self.baseaddr = 0x0

    def run(self, begin=None, end=None):
        self.ql.run(begin, end)

    def set_reg(self):
        reglist = QlEmuMisc.get_reg_map(self.ql)
        regs = [ [ row, int(self.ql.arch.regs.read(row)), self.ql.arch.bits ] for row in reglist ]
        regs_len = len(regs)
        RegDig = QlEmuRegDialog(regs)
        if RegDig.show():
            for idx, val in enumerate(RegDig.items[0:regs_len-1]):
                self.ql.arch.regs.write(reglist[idx], val[1])
            return True
        else:
            return False

    def save(self):
        savedlg = QlEmuSaveDialog()
        savedlg.Compile()

        if savedlg.Execute() != 1:
            return False

        savepath = savedlg.path_name.value

        self.ql.save(reg=True, mem=True,fd=True, cpu_context=True, snapshot=savepath)
        logging.info('Save to ' + savepath)
        return True

    def load(self):
        loaddlg = QlEmuLoadDialog()
        loaddlg.Compile()

        if loaddlg.Execute() != 1:
            return False

        loadname = loaddlg.file_name.value

        self.ql.restore(snapshot=loadname)
        logging.info('Restore from ' + loadname)
        return True

    def remove_ql(self):
        if self.ql is not None:
            del self.ql
            self.ql = None

    def ql_addr_from_ida(self, addr):
        return addr - get_imagebase() + self.baseaddr

    def ida_addr_from_ql_addr(self, addr):
        return addr - self.baseaddr + get_imagebase()

### Plugin

class QlEmuPlugin(plugin_t, UI_Hooks):
    ### Ida Plugin Data

    popup_menu_hook = None

    flags = PLUGIN_HIDE
    comment = ""

    help = "Qiling Emulator"
    wanted_name = "Qiling Emulator"
    wanted_hotkey = ""

    ### View Data

    qlemuregview = None
    qlemustackview = None
    qlemumemview = {}

    def __init__(self):
        super(QlEmuPlugin, self).__init__()
        self.plugin_name = "Qiling Emulator"
        self.qlemu = None
        self.ql = None
        self.stepflag = True
        self.stephook = None
        self.qlinit = False
        self.lastaddr = None
        self.is_change_addr = -1
        self.userobj = None
        self.customscriptpath = None
        self.bb_mapping = {}
        self.ks = None

    ### Main Framework

    def init(self):
        # init data
        logging.info('---------------------------------------------------------------------------------------')
        logging.info('Qiling Emulator Plugin For IDA, by Qiling Team. Version {0}, 2020'.format(QLVERSION))
        logging.info('Based on Qiling v{0}'.format(QLVERSION))
        logging.info('Find more information about Qiling at https://qiling.io')
        logging.info('---------------------------------------------------------------------------------------')
        self.qlemu = QlEmuQiling()
        self.ql_hook_ui_actions()
        return PLUGIN_KEEP

    def run(self, arg = 0):
        logging.info(f"Registering actions.")
        self.ql_register_menu_actions()
        self.ql_attach_main_menu_actions()

    def ready_to_run(self):
        logging.info(f"UI is ready, register our menu actions.")
        self.run()

    def term(self):
        self.qlemu.remove_ql()
        self.ql_unhook_ui_actions()
        self.ql_detach_main_menu_actions()
        self.ql_unregister_menu_actions()

    ### Actions

    def ql_start(self):
        if self.qlemu is None:
            self.qlemu = QlEmuQiling()
        if self.ql_set_rootfs():
            logging.info(f'Rootfs: {self.qlemu.rootfs}')
            logging.info(f"Custom user script: {self.customscriptpath}")
            logging.info(f"Custom env: {self.qlemu.env}")
            show_wait_box("Qiling is processing ...")
            try:
                self.qlemu.start()
                self.qlinit = True
                self.lastaddr = None
            finally:
                hide_wait_box()
                logging.info("Qiling is initialized successfully.")
        if self.customscriptpath is not None:
            self.ql_load_user_script()
            self.userobj.custom_prepare(self.qlemu.ql)

        ida_ida.ql_plugin = self

    def ql_load_user_script(self):
        if self.qlinit :
            self.ql_get_user_script(is_reload=True, is_start=True)
        else:
            logging.error('Qiling should be setup firstly.')

    def ql_reload_user_script(self):
        if self.qlinit:
            self.ql_get_user_script(is_reload=True)
        else:
            logging.error('Qiling should be setup firstly.')

    def ql_continue(self):
        if self.qlinit:
            userhook = None
            pathhook = self.qlemu.ql.hook_code(self.ql_path_hook)
            if self.userobj is not None:
                userhook = self.userobj.custom_continue(self.qlemu.ql)
            if self.qlemu.status is not None:
                self.qlemu.ql.restore(self.qlemu.status)
                show_wait_box("Qiling is processing ...")
                try:
                    self.qlemu.run(begin=self.qlemu.ql.arch.regs.arch_pc, end=self.qlemu.exit_addr)
                finally:
                    hide_wait_box()
            else:
                show_wait_box("Qiling is processing ...")
                try:
                    self.qlemu.run()
                finally:
                    hide_wait_box()
            self.qlemu.ql.hook_del(pathhook)
            if userhook and userhook is not None:
                for hook in userhook:
                    self.qlemu.ql.hook_del(hook)
            self.ql_update_views(self.qlemu.ql.arch.regs.arch_pc, self.qlemu.ql)
        else:
            logging.error('Qiling should be setup firstly.')

    def _color_path(self, color):
        def _cb(ql, addr, size):
            real_addr = addr - self.qlemu.baseaddr + get_imagebase()
            set_color(real_addr, CIC_ITEM, color)
        return _cb

    def ql_run_selection(self):
        if self.qlinit:
            _, start, end = IDA.get_last_selection()
            # Drop all previous status
            self.qlemu.status = None
            if self.userobj is not None:
                userhook = self.userobj.custom_execute_selection(self.qlemu.ql)
            colorhook = self.qlemu.ql.hook_code(self._color_path(Colors.Pink.value))
            show_wait_box("Qiling is processing ...")
            try:
                self.qlemu.run(begin=start, end=end)
            finally:
                hide_wait_box()
            self.qlemu.ql.hook_del(colorhook)
            if userhook and userhook is not None:
                for hook in userhook:
                    self.qlemu.ql.hook_del(hook)
            self.qlemu.status = self.qlemu.ql.save()
            self.ql_update_views(self.qlemu.ql.arch.regs.arch_pc, self.qlemu.ql)
        else:
            logging.error('Qiling should be setup firstly.')

    def ql_set_pc(self):
        if self.qlinit:
            ea = IDA.get_current_address()
            self.qlemu.ql.arch.regs.arch_pc = ea
            logging.info(f"QIling PC set to {hex(ea)}")
            self.qlemu.status = self.qlemu.ql.save()
            self.ql_update_views(self.qlemu.ql.arch.regs.arch_pc, self.qlemu.ql)
        else:
            logging.error('Qiling should be setup firstly.')

    def ql_run_to_here(self):
        if self.qlinit:
            curr_addr = get_screen_ea()
            untillhook = self.qlemu.ql.hook_code(self.ql_untill_hook)
            if self.qlemu.status is not None:
                self.qlemu.ql.restore(self.qlemu.status)
                show_wait_box("Qiling is processing ...")
                try:
                    self.qlemu.run(begin=self.qlemu.ql.arch.regs.arch_pc, end=curr_addr+self.qlemu.baseaddr-get_imagebase())
                finally:
                    hide_wait_box()
            else:
                show_wait_box("Qiling is processing ...")
                try:
                    self.qlemu.run(end=curr_addr+self.qlemu.baseaddr-get_imagebase())
                finally:
                    hide_wait_box()

            set_color(curr_addr, CIC_ITEM, 0x00B3CBFF)
            self.qlemu.ql.hook_del(untillhook)
            self.qlemu.status = self.qlemu.ql.save()
            self.ql_update_views(self.qlemu.ql.arch.regs.arch_pc, self.qlemu.ql)
        else:
            logging.error('Qiling should be setup firstly.')

    def ql_step(self):
        if self.qlinit:
            userhook = None
            self.stepflag = True
            self.qlemu.ql.restore(saved_states=self.qlemu.status)
            self.stephook = self.qlemu.ql.hook_code(callback=self.ql_step_hook)
            if self.userobj is not None:
                userhook = self.userobj.custom_step(self.qlemu.ql)
            self.qlemu.run(begin=self.qlemu.ql.arch.regs.arch_pc, end=self.qlemu.exit_addr)
            if userhook and userhook is not None:
                for hook in userhook:
                    self.qlemu.ql.hook_del(hook)
            self.ql_update_views(self.qlemu.ql.arch.regs.arch_pc, self.qlemu.ql)
        else:
            logging.error('Qiling should be setup firstly.')

    def ql_save(self):
        if self.qlinit:
            if self.qlemu.save() != True:
                logging.error('Fail to save the snapshot.')
        else:
            logging.error('Qiling should be setup firstly.')

    def ql_load(self):
        if self.qlinit:
            if self.qlemu.load() != True:
                logging.error('Fail to load the snapshot.')
        else:
            logging.error('Qiling should be setup firstly.')

    def ql_chang_reg(self):
        if self.qlinit:
            self.qlemu.set_reg()
            self.ql_update_views(self.qlemu.ql.arch.regs.arch_pc, self.qlemu.ql)
            self.qlemu.status = self.qlemu.ql.save()
        else:
            logging.error('Qiling should be setup firstly.')

    def ql_reset(self):
        if self.qlinit:
            self.ql_close()
            self.qlemu = QlEmuQiling()
            self.ql_start()
        else:
            logging.error('Qiling should be setup firstly.')

    def ql_close(self):
        if self.qlinit:
            heads = Heads(get_segm_start(get_screen_ea()), get_segm_end(get_screen_ea()))
            for i in heads:
                set_color(i, CIC_ITEM, 0xFFFFFF)
            self.qlemu.remove_ql()
            del self.qlemu
            self.qlemu = None
            self.qlinit = False
            logging.info('Qiling is deleted.')
        else:
            logging.error('Qiling is not started.')

    def ql_show_reg_view(self):
        if self.qlinit:
            if self.qlemuregview is None:
                self.qlemuregview = QlEmuRegView(self)
                QlEmuRegView(self)
                self.qlemuregview.Create()
                self.qlemuregview.SetReg(self.qlemu.ql.arch.regs.arch_pc, self.qlemu.ql)
                self.qlemuregview.Show()
                self.qlemuregview.Refresh()
        else:
            logging.error('Qiling should be setup firstly.')

    def ql_show_stack_view(self):
        if self.qlinit:
            if self.qlemustackview is None:
                self.qlemustackview = QlEmuStackView(self)
                self.qlemustackview.Create()
                self.qlemustackview.SetStack(self.qlemu.ql)
                self.qlemustackview.Show()
                self.qlemustackview.Refresh()
        else:
            logging.error('Qiling should be setup firstly.')

    def ql_show_mem_view(self, addr=get_screen_ea(), size=0x10):
        if self.qlinit:
            memdialog = QlEmuMemDialog()
            memdialog.Compile()
            memdialog.mem_addr.value = addr
            memdialog.mem_size.value = size
            ok = memdialog.Execute()
            if ok == 1:
                mem_addr = memdialog.mem_addr.value - self.qlemu.baseaddr + get_imagebase()
                mem_size = memdialog.mem_size.value
                mem_cmnt = memdialog.mem_cmnt.value

                if mem_addr not in self.qlemumemview:
                    if not self.qlemu.ql.mem.is_mapped(mem_addr, mem_size):
                        ok = ask_yn(1, "Memory [%X:%X] is not mapped!\nDo you want to map it?\n   YES - Load Binary\n   NO - Fill page with zeroes\n   Cancel - Close dialog" % (mem_addr, mem_addr + mem_size))
                        if ok == 0:
                            self.qlemu.ql.mem.map(mem_addr, mem_size)
                            self.qlemu.ql.mem.write(mem_addr, b"\x00"*mem_size)
                        elif ok == 1:
                            # TODO: map_binary
                            return
                        else:
                            return
                    self.qlemumemview[mem_addr] = QlEmuMemView(self, mem_addr, mem_size)
                    if mem_cmnt == []:
                        self.qlemumemview[mem_addr].Create("QL Memory")
                    else:
                        self.qlemumemview[mem_addr].Create("QL Memory [ " + mem_cmnt + " ]")
                    self.qlemumemview[mem_addr].SetMem(self.qlemu.ql)
                self.qlemumemview[mem_addr].Show()
                self.qlemumemview[mem_addr].Refresh()
        else:
            logging.error('Qiling should be setup firstly.')

    def ql_unload_plugin(self):
        heads = Heads(get_segm_start(get_screen_ea()), get_segm_end(get_screen_ea()))
        for i in heads:
            set_color(i, CIC_ITEM, 0xFFFFFF)
        self.ql_close()
        self.ql_detach_main_menu_actions()
        self.ql_unregister_menu_actions()
        logging.info('Unload plugin successfully!')

    def ql_menu_null(self):
        pass

    def ql_about(self):
        self.aboutdlg = QlEmuAboutDialog(QLVERSION)
        self.aboutdlg.Execute()
        self.aboutdlg.Free()

    def ql_check_update(self):
        (r, content) = QlEmuMisc.url_download(QilingStableVersionURL)
        content = content.decode("utf-8")
        if r == 0:
            try:
                version_stable = re.findall(r"\"([\d\.]+)\"", content)[0]
            except (TypeError, IndexError):
                warning("ERROR: Failed to find the Qiling version string from response.")
                logging.warning("Failed to find the Qiling version string from response.")

            # compare with the current version
            if version_stable == QLVERSION:
                self.updatedlg = QlEmuUpdateDialog(QLVERSION, "Good, you are already on the latest stable version!")
                self.updatedlg.Execute()
                self.updatedlg.Free()
            else:
                self.updatedlg = QlEmuUpdateDialog(QLVERSION, "Download latest stable version {0} from https://github.com/qilingframework/qiling/blob/master/qiling/extensions/idaplugin".format(version_stable))
                self.updatedlg.Execute()
                self.updatedlg.Free()
        else:
            # fail to download
            warning("ERROR: Failed to connect to Github. Try again later.")
            logging.warning("Failed to connect to Github when checking for the latest update. Try again later.")

    def _remove_from_bb_lists(self, bbid):
        if bbid in self.real_blocks:
            self.real_blocks.remove(bbid)
        elif bbid in self.fake_blocks:
            self.fake_blocks.remove(bbid)
        elif bbid in self.retn_blocks:
            self.retn_blocks.remove(bbid)

    def ql_mark_real(self):
        if len(self.bb_mapping) == 0:
            logging.error(f"Please perform auto analysis before marking blocks manually!")
            return
        cur_addr = IDA.get_current_address()
        cur_block = IDA.get_block(cur_addr)
        self._remove_from_bb_lists(cur_block.id)
        self.real_blocks.append(cur_block.id)
        IDA.color_block(cur_block, Colors.Green.value)

    def ql_mark_fake(self):
        if len(self.bb_mapping) == 0:
            logging.error(f"Please perform auto analysis before marking blocks manually!")
            return
        cur_addr = IDA.get_current_address()
        cur_block = IDA.get_block(cur_addr)
        self._remove_from_bb_lists(cur_block.id)
        self.fake_blocks.append(cur_block.id)
        IDA.color_block(cur_block, Colors.Gray.value)

    def ql_mark_retn(self):
        if len(self.bb_mapping) == 0:
            logging.error(f"Please perform auto analysis before marking blocks manually!")
            return
        cur_addr = IDA.get_current_address()
        cur_block = IDA.get_block(cur_addr)
        self._remove_from_bb_lists(cur_block.id)
        self.retn_blocks.append(cur_block.id)
        IDA.color_block(cur_block, Colors.Pink.value)

    def __in_bb(self, addr, bb):
        return addr < bb.end_ea and addr >= bb.start_ea

    # Identify if the given microcode block is a jmp_mbb.
    # Core patterns:
    #     ...
    #     ...
    #     goto @dispatcher
    def __is_jmp_mbb(self, mbb):
        ins_list = list(IDA.micro_code_from_mbb(mbb))
        if len(ins_list) == 0:
            logging.warning(f"Get an empty mbb at {hex(mbb.start)}?!")
            return False
        last_ins = ins_list[-1]
        if last_ins.opcode != ida_hexrays.m_goto:
            logging.warning(f"jmp_mbb at {hex(mbb.start)} the opcode of last instruction {last_ins._print()} isn't goto")
            return False
        if last_ins.l.t != ida_hexrays.mop_b:
            logging.warning(f"jmp_mbb at {hex(mbb.start)} the l of last instruction {last_ins._print()} doesn't have a microcode block reference")
            return False
        goto_mbb = self.mbbs[last_ins.l.b]
        mbb_start = goto_mbb.start
        dispatcher_bb = self.bb_mapping[self.dispatcher]
        pre_dispatcher_bb = self.bb_mapping[self.pre_dispatcher]
        if self.__in_bb(mbb_start, dispatcher_bb) or self.__in_bb(mbb_start, pre_dispatcher_bb):
            return True
        logging.warning(f"The address {hex(mbb_start)} where jmp_mbb goes isn't pre_dispatcher or dispatcher block!")
        return False

    # Identify if the given microcode block is a next_mbb.
    # Core patterns:
    #     mov #imm, reg
    def __is_next_mbb(self, mbb):
        ins_list = list(IDA.micro_code_from_mbb(mbb))
        if len(ins_list) == 0:
            logging.warning(f"Get an empty mbb at {hex(mbb.start)}?!")
            return False
        first_ins = ins_list[0]
        if first_ins.opcode != ida_hexrays.m_mov:
            logging.warning(f"next_mbb at {hex(mbb.start)} the opcode of first instruction {first_ins._print()} isn't mov")
            return False
        if first_ins.l.t != ida_hexrays.mop_n:
            logging.warning(f"next_mbb at {hex(mbb.start)} the l of first instruction {first_ins._print()} isn't an immediate number")
            return False
        if first_ins.d.t != ida_hexrays.mop_r:
            logging.warning(f"next_mbb at {hex(mbb.start)} the d of first instruction {first_ins._print()} isn't a reg")
            return False
        if len(ins_list) == 1:
            logging.info(f"A block with only one instruction which is `mov #imm, reg` at {hex(mbb.start)}.")
        return True

    def _get_jmp_ins(self, ida_addr, insns):
        # This cloud really happen! See issue #804. TODO: Investigate or re-design insns structure or replace it with ESIL.
        # So we have to fallback to legacy path.
        if ida_addr not in insns:
            return (None, None)
        ins_list = insns[ida_addr]
        result = []
        for bbid, ins in ins_list:
            if ida_hexrays.is_mcode_jcond(ins.opcode):
                result.append((bbid, ins))
        if len(result) > 1:
            logging.warning(f"More than one conditional jmp detected at {hex(ida_addr)}!")
        elif len(result) == 0:
            logging.warning(f"No conditional jmp found at {hex(ida_addr)}!")
            return (None, None)
        return result[0]

    # cmov/it eq:
    #   < ... condiontal jump > --> < mov imm, reg > (next_bb) --> < dispatcher >
    #                           |-> < dispatcher > (jmp_mbb)
    # NOTE:
    #   LLVM IR -> ASM -> IDA IR
    #   I guess this pattern is stable enough for us to rely on.
    def _force_execution_with_microcode(self, ql, ida_addr):
        _, insns, mbbs = self._prepare_microcodes(maturity=7)
        bbid, ins = self._get_jmp_ins(ida_addr, insns)
        if ins is None:
            return False
        # According to comments in hexrays.hpp, it may be a mop_v. I guess that shouldn't exist
        # so we add a sanity check here.
        if ins.d.t != ida_hexrays.mop_b:
            logging.warning(f"Sanity check: microcode {ins._print()} doesn't refer a block!")
        jmp_mbb = mbbs[ins.d.b]
        next_mbb = mbbs[bbid].nextb
        if not (self.__is_next_mbb(next_mbb) and self.__is_jmp_mbb(jmp_mbb)):
            # Switch the branch and try again?
            logging.info("Switch the jmp_bb and next_bb and try again...")
            if self.__is_jmp_mbb(next_mbb) and self.__is_next_mbb(jmp_mbb):
                jmp_mbb, next_mbb = next_mbb, jmp_mbb
            else:
                logging.error(f"Fail to identify microcode blocks at {hex(ida_addr)}")
                return False
        ins_list = list(IDA.micro_code_from_mbb(next_mbb))
        first_ins = ins_list[0]
        imm = first_ins.l.nnn.value
        reg_name = ida_hexrays.get_mreg_name(first_ins.d.r, ql.arch.pointersize)
        logging.info(f"Froce set {reg_name} to {hex(imm)}")
        ql.arch.regs.__setattr__(reg_name, imm)
        return True

    def _ida_address_after_branch(self, ida_addr):
        _, ins = self._get_jmp_ins(ida_addr, self.insns)
        if ins is None:
            return None
        return self.mbbs[ins.d.b].start

    # legacy approach for x86_64 and arm64
    # TODO & NOTE: We still need this fallback since the microcode pattern we match now sometimes may change. A better
    #              approach is that we fully drop pattern machting and emulates the microcode while recording the correspoding
    #              branch condition.
    #              For example:
    #                CSEL            W8, W8, W9, LT
    #                =>
    #                'LT' -> W9->W8 (the value of W9 can be get from emulation)
    #                'GT' -> W8->W8
    #                We assume that the comparison (like cmp W10, 0) is the same between real assembly and microcode and the only
    #                difference is the condition of the jump instruction. So the 'LT' and 'GT' here represent condition both in microcode
    #                and original assembly.
    #              For current implementation we don't record 'LT' or 'GT' which makes it hard to patch code after force execution.
    def _force_execution_by_parsing_assembly(self, ql, ida_addr):
        if "x86" in IDA.get_ql_arch_string(): # cmovlg eax, ebx
            reg1 = IDA.print_operand(ida_addr, 0).lower()
            reg2 = IDA.print_operand(ida_addr, 1).lower()
            reg2_val = ql.arch.regs.__getattr__(reg2)
            logging.info(f"Force set {reg1} to {hex(reg2_val)}")
            ql.arch.regs.__setattr__(reg1, reg2_val)
            return True
        elif "arm" in IDA.get_ql_arch_string():
            instr = IDA.get_instruction(ida_addr).lower()
            logging.info(f"Going to force execute: {instr}")
            if instr.startswith("it"): # itt eq\n moveqw low\n movteq high\n
                ida_addr = ida_addr + IDA.get_instruction_size(ida_addr)
                low = IDA.get_operand(ida_addr, 1)
                ida_addr = ida_addr + IDA.get_instruction_size(ida_addr)
                high = IDA.get_operand(ida_addr, 1)
                reg = IDA.print_operand(ida_addr, 0).lower()
                val = (high << 16) + low
                logging.info(f"Force set {reg} to {hex(val)}")
                ql.arch.regs.__setattr__(reg, val)
                return True
            elif "csel" in instr: # csel dst, src1, src2, cond
                dst = IDA.print_operand(ida_addr, 0).lower()
                src = IDA.print_operand(ida_addr, 2).lower()
                src_val = ql.arch.regs.__getattr__(src)
                logging.info(f"Force set {dst} to {hex(src_val)}")
                ql.arch.regs.__setattr__(dst, src_val)
                return True
        return False

    def _has_call_insn(self, ida_addr):
        if ida_addr not in self.insns:
            return False
        ins_list = self.insns[ida_addr]
        for _, ins in ins_list:
            if ida_hexrays.is_mcode_call(ins.opcode):
                return True
        return False

    def _guide_hook(self, ql, addr, size):
        start_bb_id = self.hook_data['startbb']
        ida_addr = self.deflatqlemu.ida_addr_from_ql_addr(addr)
        func = self.hook_data['func']
        if ida_addr < func.start_ea or ida_addr >= func.end_ea:
            logging.error(f"Address {hex(ida_addr)} out of function boundaries!")
            ql.emu_stop()
            self.hook_data['result'] = False
            return
        cur_bb = IDA.get_block(ida_addr)
        if "force" in self.hook_data and ida_addr in self.hook_data['force']:
            if self.hook_data['force'][ida_addr]:
                logging.info(f"Force execution at {hex(ida_addr)}")
                result = self._force_execution_with_microcode(ql, ida_addr)
                if not result:
                    logging.warning(f"Fail to force execution by microcode at {hex(ida_addr)}, trying legacy approach")
                    result = self._force_execution_by_parsing_assembly(ql, ida_addr)
                    if not result:
                        logging.error(f"Fail to force execution by legacy approach at {hex(ida_addr)}, stop now...")
                        self.hook_data['result'] = False
                        ql.emu_stop()
                        return
                    self.hook_data['force_legacy'] = True
            else:
                pass
            if not self.hook_data['force_legacy']:
                next_ida_addr = self._ida_address_after_branch(ida_addr)
            else:
                next_ida_addr = ida_addr + IDA.get_instruction_size(ida_addr)
            logging.info(f"Goto {hex(next_ida_addr)} after branch...")
            ql.arch.regs.arch_pc = self.deflatqlemu.ql_addr_from_ida(next_ida_addr) + self.append
            ida_addr = next_ida_addr
        # TODO: Maybe we can detect whether the program will access unmapped
        #       here so that we won't map the memory.
        if self._has_call_insn(ida_addr):
            ql.arch.regs.arch_pc += IDA.get_instruction_size(ida_addr) + self.append
            return
        if start_bb_id == cur_bb.id:
            return
        if cur_bb.id in self.real_blocks or cur_bb.id in self.retn_blocks:
            if cur_bb.id not in self.paths[start_bb_id]:
                self.paths[start_bb_id].append(cur_bb.id)
            ql.emu_stop()

    def _skip_unmapped_rw(self, ql: Qiling, type, addr, size, value):
        map_addr = ql.mem.align(addr)
        map_size = ql.mem.align_up(size)

        if not ql.mem.is_mapped(map_addr, map_size):
            logging.warning(f"Invalid memory R/W, trying to map {hex(map_size)} at {hex(map_addr)}")

            ql.mem.map(map_addr, map_size)
            ql.mem.write(map_addr, b'\x00' * map_size)

        return True

    def _find_branch_in_real_block(self, bb):
        #paddr = bb.start_ea
        #while paddr < bb.end_ea:
        #    ins = IDA.get_instruction(paddr)
        #    sz = IDA.get_instruction_size(paddr)
        #    if ins.lower().startswith("cmov"):
        #        return paddr
        #    paddr += sz

        # microcode implementation
        #  Is it possible for the address of a microcode instruction not equal
        #  to any start address of a real instruction? IDK, only Hexrays knows.
        #  So we iterate all addresses so that we won't miss any posssible branchs.
        for addr in range(bb.start_ea, bb.end_ea):
            if addr in self.insns:
                for _, insn in self.insns[addr]:
                    opcode = insn.opcode
                    if ida_hexrays.is_mcode_jcond(opcode):
                        return addr
        return None

    def _log_paths_str(self):
        for bbid, succs in self.paths.items():
            if len(succs) == 1:
                logging.info(f"{self._block_str(bbid)} -> {self._block_str(succs[0])}")
            elif len(succs) == 2:
                logging.info(f"{self._block_str(bbid)} --(force jump)--> {self._block_str(succs[0])}")
                logging.info(f"|----(skip jump)----> {self._block_str(succs[1])}")
            elif len(succs) > 2:
                logging.warning(f"succs: {succs} found from {self._block_str(bbid)}!")

    # Is this correct?
    def _thumb_detect(self, ida_addr):
        return IDA.get_instruction_size(ida_addr) == 2

    def _log_verbose(self, ql, addr, size):
        logging.debug(f"addr: {hex(addr)} ida_addr: {hex(self.deflatqlemu.ida_addr_from_ql_addr(addr))}")
        registers = [ k for k in ql.arch.regs.register_mapping.keys() if type(k) is str ]
        for idx in range(0, len(registers), 3):
            regs = registers[idx:idx+3]
            s = "\t".join(map(lambda v: f"{v:4}: {ql.arch.regs.__getattr__(v):016x}", regs))
            logging.debug(s)

    # Q: Why we need emulation to help us find real control flow considering there are some
    #    switch-case patterns in mircocode which can be analysed statically?
    # A: Emulation makes the de-obf much more robust and general and can work under less assumptions, like
    #    the optimization level, the complexity of original control flow or even the verson of the
    #    Hexrays decompiler because we try to avoid matching fixed microcode patterns. In fact, static
    #    analysis is always possible since the brain of a reverse engineer is the best decompiler
    #    and emulator, isn't it? :D
    def _search_path(self):
        self.paths = {bbid: [] for bbid in self.bb_mapping.keys()}
        reals = [self.first_block, *self.real_blocks]
        self.deflatqlemu = QlEmuQiling()
        self.deflatqlemu.path = self.qlemu.path
        self.deflatqlemu.rootfs = self.qlemu.rootfs
        first_block = self.bb_mapping[self.first_block]
        if IDA.get_ql_arch_string() == "arm32":
            if self._thumb_detect(first_block.start_ea):
                logging.info(f"Thumb detected, enable it.")
                self.deflatqlemu.start(archtype=QL_ARCH.ARM, thumb=True)
                self.deflatqlemu.ql.arch.regs.cpsr |= 0x20
                self.append = 1
        else:
            self.deflatqlemu.start()
            self.append = 0
        ql = self.deflatqlemu.ql
        if logging.root.level <= logging.DEBUG:
            ql.hook_code(self._log_verbose)
        self.hook_data = None
        ql.hook_mem_read_invalid(self._skip_unmapped_rw)
        ql.hook_mem_write_invalid(self._skip_unmapped_rw)
        ql.hook_mem_unmapped(self._skip_unmapped_rw)
        # set up stack before we really run.
        ql.run(begin=self.deflatqlemu.ql_addr_from_ida(first_block.start_ea) + self.append, end=self.deflatqlemu.ql_addr_from_ida(first_block.end_ea), count=0xFFF)
        # okay, we can set up our core hook now.
        ql.hook_code(self._guide_hook)
        for bbid in reals:
            logging.debug(f"Search control flow for block: {self._block_str(bbid)}")
            bb = self.bb_mapping[bbid]
            braddr = self._find_branch_in_real_block(bb)
            self.hook_data = {
                "startbb": bbid,
                "func": IDA.get_function(first_block.start_ea),
                "result": True,
                "force_legacy": False
            }
            ql_bb_start_ea = self.deflatqlemu.ql_addr_from_ida(bb.start_ea) + self.append
            ctx = ql.save()
            # Skip force execution in the first block.
            # `end=0` is a workaround for ql remembering last exit_point.
            if braddr is None or bb.id == self.first_block:
                ql.run(begin=ql_bb_start_ea, end=0, count=0xFFF)
            else:
                self.hook_data['force'] = {braddr: True}
                ctx2 = ql.save()
                ql.run(begin=ql_bb_start_ea, end=0, count=0xFFF)
                ql.restore(ctx2)
                if not self.hook_data['result']:
                    return False
                self.hook_data['force'] = {braddr: False}
                ql.run(begin=ql_bb_start_ea, end=0, count=0xFFF)
            ql.restore(ctx)
            if not self.hook_data['result']:
                return False
        self._log_paths_str()
        return True


    # IDA doesn't support arm assembling, so it's a good chance to replace IDA
    # assembler implmentation with keystone.
    def _initialize_keystone(self):
        if self.ks is None:
            self.ks = self.deflatqlemu.ql.arch.assembler


    def _asm(self, *args, **kwargs):
        self._initialize_keystone()
        return self.ks.asm(*args, **kwargs)

    def _assemble(self, instr, addr):
        logging.debug(f"Keystone: Assemble {instr} at {hex(addr)}")
        bs, count = self._asm(instr, addr)
        IDA.patch_bytes(addr, bytes(bs))
        IDA.perform_analysis(addr, addr + len(bs))
        return bs, count

    # Patching microcode is TOO complex.
    # I would rahter write another 1e10 llvm passes than a single hexrays decompiler pass.
    def _arch_jmp_instruction(self, addr):
        arch = IDA.get_ql_arch_string()
        op = None
        if "x86" in arch:
            op = "jmp"
        elif "arm" in arch:
            op = "B"
        elif "mips" in arch:
            op = "j"
        return f"{op} {addr}"

    # See comments above.
    def _arch_cond_jmp_instruction(self, cond, addr):
        arch = IDA.get_ql_arch_string()
        op = None
        if "x86" in arch:
            op = f"j{cond}"
        elif "arm" in arch:
            op = f"b{cond}"
        elif "mips" in arch:
            op = f"j{cond}"
        return f"{op} {addr}"

    # See comments above.
    def _arch_parse_cond_from_addr(self, braddr):
        arch = IDA.get_ql_arch_string()
        instr = IDA.get_instruction(braddr).lower()
        if "x86" in arch: # cmovge
            return instr[4:]
        elif "arm" in arch:
            if instr.startswith("it"): # itt eq
                tks = instr.split(" ")
                if len(tks) != 2:
                    logging.error(f"Can't get condition from {instr}")
                    return None
                return tks[-1]
            elif "csel" in instr:
                return IDA.print_operand(braddr, 3)
        # TODO: mips
        return None

    # Really FORCE ida to analyse again.
    def _force_analysis(self, start, end):
        ida_bytes.del_items(start, 0, end-start)
        ida_auto.plan_and_wait(start, end)
        return

    # The only way to make IDA generate right assembly is to undefine bytes patched just now
    # and call IDA to analyse it again. Nice job again as always, Hexrays!
    def _patch_bytes_with_force_analysis(self, start, bs):
        IDA.patch_bytes(start, bs)
        return self._force_analysis(start, start+len(bs))

    def _arch_branch_patch(self, braddr, bbid):
        bb = self.bb_mapping[bbid]
        force_addr = self.bb_mapping[self.paths[bbid][0]].start_ea
        normal_addr = self.bb_mapping[self.paths[bbid][1]].start_ea
        # Temporary dirty fix.
        # See comments for _force_execution_by_parsing_assembly.
        if "arm64" == IDA.get_ql_arch_string():
            force_addr, normal_addr = normal_addr, force_addr
        # Parse condition before patching nop.
        cond = self._arch_parse_cond_from_addr(braddr)
        buffer = [0] * (bb.end_ea - braddr)
        instr_to_assemble = self._arch_cond_jmp_instruction(cond, f"{hex(force_addr)}h")
        logging.info(f"Assemble {instr_to_assemble} at {hex(force_addr)}")
        bs1, _ = self._asm(instr_to_assemble, braddr)
        buffer[:len(bs1)] = bs1
        next_instr_address = braddr + len(bs1)
        instr_to_assemble = self._arch_jmp_instruction(f"{hex(normal_addr)}h")
        logging.info(f"Assemble {instr_to_assemble} at {hex(normal_addr)}")
        bs2, _ = self._asm(instr_to_assemble, next_instr_address)
        buffer[len(bs1):len(bs1) + len(bs2)] = bs2
        logging.info(f"Patch real block with branch from {hex(braddr)} to {hex(bb.end_ea)}")
        self._patch_bytes_with_force_analysis(braddr, bytes(buffer))


    def _patch_codes(self):
        if len(self.paths[self.first_block]) != 1:
            logging.error(f"Found wrong ways in first block: {self._block_str(self.bb_mapping[self.first_block])}, should be 1 path but get {len(self.paths[self.first_block])}, exit.")
            return
        logging.info("NOP dispatcher block")
        dispatcher_bb = self.bb_mapping[self.dispatcher]
        # Some notes:
        #    Patching b'\x00' instead of 'nop' can help IDA decompile a better result. Don't know why...
        #    Besides
        buffer = [0] * (dispatcher_bb.end_ea - dispatcher_bb.start_ea)
        first_jmp_addr = dispatcher_bb.start_ea
        instr_to_assemble = self._arch_jmp_instruction(f"{hex(self.bb_mapping[self.paths[self.first_block][0]].start_ea)}h")
        logging.info(f"Assemble {instr_to_assemble} at {hex(first_jmp_addr)}")
        bs, _ = self._asm(instr_to_assemble, first_jmp_addr)
        buffer[:len(bs)] = bs
        logging.info(f"Patch first jump at {hex(first_jmp_addr)}")
        self._patch_bytes_with_force_analysis(first_jmp_addr, bytes(buffer))
        for bbid in self.real_blocks:
            logging.debug(f"Patching real block: {self._block_str(bbid)}")
            bb = self.bb_mapping[bbid]
            braddr = self._find_branch_in_real_block(bb)
            if braddr is None:
                last_instr_address = IDA.get_prev_head(bb.end_ea)
                buffer = [0x90] * (bb.end_ea - last_instr_address)
                if len(self.paths[bbid]) != 1:
                    logging.warning(f"Found wrong ways in block: {self._block_str(bb)}, should be 1 path but get {len(self.paths[bbid])}")
                    continue
                instr_to_assemble = self._arch_jmp_instruction(f"{hex(self.bb_mapping[self.paths[bbid][0]].start_ea)}h")
                logging.info(f"Assemble {instr_to_assemble} at {hex(last_instr_address)}")
                bs, _ = self._asm(instr_to_assemble, last_instr_address)
                buffer[:len(bs)] = bs
                logging.info(f"Patch real block from {hex(last_instr_address)} to {hex(bb.end_ea)}")
                self._patch_bytes_with_force_analysis(last_instr_address, bytes(buffer))
            else:
                if len(self.paths[bbid]) != 2:
                    logging.warning(f"Found wrong ways in block: {self._block_str(bb)}, should be 2 paths but get {len(self.paths[bbid])}")
                    continue
                self._arch_branch_patch(braddr, bbid)
        for bbid in self.fake_blocks:
            bb = self.bb_mapping[bbid]
            logging.info(f"Patch NOP for block: {self._block_str(bb)}")
            self._patch_bytes_with_force_analysis(bb.start_ea, b"\x00"*(bb.end_ea-bb.start_ea))
        logging.info(f"Patch NOP for pre_dispatcher.")
        bb = self.bb_mapping[self.pre_dispatcher]
        self._patch_bytes_with_force_analysis(bb.start_ea, b"\x00"*(bb.end_ea-bb.start_ea))

    def _prepare_microcodes(self, decomp_flags=ida_hexrays.DECOMP_WARNINGS | ida_hexrays.DECOMP_NO_WAIT, maturity=7):
        dispatcher_bb = self.bb_mapping[self.dispatcher]
        target_function = IDA.get_function(dispatcher_bb.start_ea)
        # Reduce optimization to make pattern more stable.
        logging.info(f"Generate microcode from {hex(target_function.start_ea)} to {hex(target_function.end_ea)}")
        mba = IDA.get_micro_code_mba(target_function.start_ea, target_function.end_ea, decomp_flags, maturity)
        insns = {}
        mbbs = {}
        for i in range(mba.qty):
            mbb = mba.get_mblock(i)
            mbbs[i] = mbb
            cur_insn = mbb.head
            while cur_insn != None:
                insn_ea = cur_insn.ea
                if insn_ea not in insns:
                    insns[insn_ea] = []
                insns[insn_ea].append((i, cur_insn))
                cur_insn = cur_insn.next
        return mba, insns, mbbs

    def ql_deflat(self):
        if len(self.bb_mapping) == 0:
            self.ql_parse_blocks_for_deobf()
        if not self.qlinit:
            logging.info("Qiling should be setup firstly!")
            return
        self.mba, self.insns, self.mbbs = self._prepare_microcodes(maturity=3)
        logging.debug("Microcode generation done. Going to search path.")
        if not self._search_path():
            logging.info(f"Fail to search path. Please fire an issue to us at https://github.com/qilingframework/qiling with relevant logs!")
            return
        logging.debug("Real control flows search done. Going to patch codes.")
        self._patch_codes()
        logging.debug("Codes patched. Let's tell IDA to analyse the whole function again.")
        IDA.perform_analysis(self.deflat_func.start_ea, self.deflat_func.end_ea)
        del self.deflatqlemu
        self.deflatqlemu = None
        self.ks = None

    def _block_str(self, bb):
        if type(bb) is int:
            bb = self.bb_mapping[bb]
        return f"Block id: {bb.id}, start_address: {bb.start_ea:x}, end_address: {bb.end_ea:x}, type: {bb.type}"

    def ql_parse_blocks_for_deobf(self):
        cur_addr = IDA.get_current_address()
        flowchart = IDA.get_flowchart(cur_addr)
        self.deflat_func = IDA.get_function(cur_addr)
        self.bb_mapping = {bb.id:bb for bb in flowchart}
        if flowchart is None:
            return
        bb_count = {}
        for bb in flowchart:
            for succ in bb.succs():
                if succ.id not in bb_count:
                    bb_count[succ.id] = 0
                bb_count[succ.id] += 1
        max_ref_bb_id = None
        max_ref = 0
        for bb_id, ref in bb_count.items():
            if ref > max_ref:
                max_ref = ref
                max_ref_bb_id = bb_id
        self.pre_dispatcher = max_ref_bb_id
        try:
            self.dispatcher = list(self.bb_mapping[self.pre_dispatcher].succs())[0].id
            self.first_block = flowchart[0].id
        except IndexError:
            logging.error("Fail to get dispatcher and first_block.")
            return
        self.real_blocks = []
        self.fake_blocks = []
        self.retn_blocks = []
        for bb in flowchart:
            if self.pre_dispatcher in [b.id for b in bb.succs()] and IDA.get_instructions_count(bb.start_ea, bb.end_ea) > 1:
                self.real_blocks.append(bb.id)
            elif IDA.block_is_terminating(bb):
                self.retn_blocks.append(bb.id)
            elif bb.id != self.first_block and bb.id != self.pre_dispatcher and bb.id != self.dispatcher:
                self.fake_blocks.append(bb.id)
        for bbid in self.real_blocks:
            IDA.color_block(self.bb_mapping[bbid], Colors.Green.value)
        for bbid in self.fake_blocks:
            IDA.color_block(self.bb_mapping[bbid], Colors.Gray.value)
        for bbid in self.retn_blocks:
            IDA.color_block(self.bb_mapping[bbid], Colors.Pink.value)
        IDA.color_block(self.bb_mapping[self.dispatcher], Colors.Blue.value)
        IDA.color_block(self.bb_mapping[self.pre_dispatcher], Colors.Blue.value)
        IDA.color_block(self.bb_mapping[self.first_block], Colors.Beige.value)
        logging.info(f"First block: {self._block_str(self.first_block)}")
        logging.info(f"Dispatcher: {self._block_str(self.dispatcher)}")
        logging.info(f"Pre dispatcher: {self._block_str(self.pre_dispatcher)}")
        logging.info(f"Real blocks:")
        for s in map(self._block_str, self.real_blocks): logging.info(s)
        logging.info(f"Fake blocks:")
        for s in map(self._block_str, self.fake_blocks): logging.info(s)
        logging.info(f"Return blocks:")
        for s in map(self._block_str, self.retn_blocks): logging.info(s)
        logging.info(f"Auto analysis finished, please check whether the result is correct.")
        logging.info(f"You may change the property of each block manually if necessary.")

    # jb addr
    # jnb addr
    # <=>
    # jmp addr / nop
    # nop
    def _junk_useless_jcc(self, start, end):
        bs = IDA.get_bytes(start, end-start)
        patterns = []
        # Volume 2 Table A-2
        for opc in range(0x70, 0x7F, 2):
            # jb -> jnb
            r_opc = opc + 1
            patterns.append( re.compile(re.escape(bytes([opc])) + rb'.' + re.escape(bytes([r_opc])) + rb'.') )
            patterns.append( re.compile(re.escape(bytes([r_opc])) + rb'.' + re.escape(bytes([opc])) + rb'.') )

        for pattern in patterns:
            logging.debug(f"Start with pattern: {pattern}")
            tmpbs = bs
            result = re.search(pattern, tmpbs)
            while result is not None:
                l = result.span()[0]
                r = result.span()[1]
                l_offset = tmpbs[l+1]
                r_offset = tmpbs[l+3]
                tmpbs = tmpbs[r:]
                result = re.search(pattern, tmpbs)
                if l_offset == r_offset + 2:
                    logging.info(f"Get a junk jcc at [{hex(start+l)}, {hex(start+r)}] with offset {hex(l_offset)}.")
                    # If it jumps down, check if we can fill the codes with nops safely.
                    if l_offset <= 0x7F:
                        can_fill_with_nops = True
                        for addr in range(start+l, start + l +  2 + l_offset):
                            if len(IDA.get_xrefsto(addr, flags=ida_xref.XREF_FAR)) != 0:
                                logging.info(f"Find multiple Xrefs at {hex(addr)}, patch a jmp...")
                                can_fill_with_nops = False
                                break
                        if can_fill_with_nops:
                            logging.info(f"Fill NOPs from {hex(start+l)} to {hex(start + l + 2 + l_offset)}.")
                            IDA.fill_bytes(start+l, start + l + 2 + l_offset)
                            continue
                    # Or we simply patch a jmp.
                    logging.info(f"Patch a jmp at {start+l:x}.")
                    IDA.fill_bytes(start+l, start+r)
                    IDA.patch_bytes(start+l, b"\xeb" + bytes([l_offset]))


    # Remove junk code by fixed patterns.
    # If you find new patterns, please fire an issue or PR!
    # TODO: Find some samples to add more patterns.
    def ql_remove_junk_code_by_patterns(self):
        _, start, end = IDA.get_last_selection()
        self._junk_useless_jcc(start, end)


    def ql_nop_items_without_colors(self):
        _, start, end = IDA.get_last_selection()
        cur_start, cur_end = IDA.get_item(start)
        while cur_start < end:
            if not IDA.is_colored_item(cur_start):
                logging.info(f"Item at [{hex(cur_start)}, {hex(cur_end)}) doesn't have a color, nop it.")
                IDA.fill_bytes(cur_start, cur_end)
            cur_start, cur_end = IDA.get_item(cur_end)

    ### Hook

    def ql_step_hook(self, ql, addr, size):
        self.stepflag = not self.stepflag
        addr = addr - self.qlemu.baseaddr + get_imagebase()
        if self.stepflag:
            set_color(addr, CIC_ITEM, 0x00FFD700)
            self.ql_update_views(self.qlemu.ql.arch.regs.arch_pc, ql)
            self.qlemu.status = ql.save()
            ql.os.stop()
            self.qlemu.ql.hook_del(self.stephook)
            jumpto(addr)

    def ql_path_hook(self, ql, addr, size):
        addr = addr - self.qlemu.baseaddr + get_imagebase()
        set_color(addr, CIC_ITEM, 0x007FFFAA)
        bp_count = get_bpt_qty()
        bp_list = []
        if bp_count > 0:
            for num in range(0, bp_count):
                bp_list.append(get_bpt_ea(num))

            if addr in bp_list and (addr != self.lastaddr or self.is_change_addr>1):
                self.qlemu.status = ql.save()
                ql.os.stop()
                self.lastaddr = addr
                self.is_change_addr = -1
                jumpto(addr)

            self.is_change_addr += 1


    def ql_untill_hook(self, ql, addr, size):
        addr = addr - self.qlemu.baseaddr + get_imagebase()
        set_color(addr, CIC_ITEM, 0x00B3CBFF)

    ### User Scripts

    def ql_get_user_script(self, is_reload=False, is_start=False):
        def get_user_scripts_obj(scriptpath:str, classname:str, is_reload:bool):
            try:
                import os, sys
                import importlib

                modulepath,filename = os.path.split(scriptpath)
                logging.info(modulepath)
                logging.info(filename)
                scriptname,_ = os.path.splitext(filename)
                logging.info(scriptname)
                sys.path.append(modulepath)
                module = importlib.import_module(scriptname)

                if is_reload:
                    importlib.reload(module)
                cls = getattr(module, classname)
                return cls()
            except Exception as e:
                logging.exception("")
                return None

        self.userobj = get_user_scripts_obj(self.customscriptpath, 'QILING_IDA', is_reload)
        if self.userobj is not None:
            if is_reload and not is_start:
                logging.info('Custom user script is reloaded.')
            else:
                logging.info('Custom user script is loaded successfully.')
        else:
            logging.info('Custom user script not found.')

    ### Dialog

    def ql_set_rootfs(self):
        setupdlg = QlEmuSetupDialog()
        setupdlg.Compile()

        if setupdlg.Execute() != 1:
            return False

        rootfspath = setupdlg.path_name.value
        customscript = setupdlg.script_name.value
        parameter = setupdlg.parameters.value
        env = setupdlg.env_var.value

        if customscript != '':
            self.customscriptpath = customscript

        para_array = shlex.split(parameter)
        self.qlemu.path = [get_input_file_path()] + para_array
        logging.info(self.qlemu.path)

        if env != '':
            try:
                env_f = open(env)
                self.qlemu.env = load(env_f)
                env_f.close()
            except Exception as e:
                logging.error("Error parsing the env file!!")
                logging.error(e.message)
                return False

        if self.qlemu is not None:
            self.qlemu.rootfs = rootfspath
            return True
        return False

    ### Menu

    menuitems = []

    def ql_register_new_action(self, act_name, act_text, act_handler, shortcut, tooltip, icon):
        new_action = action_desc_t(
            act_name,       # The action name. This acts like an ID and must be unique
            act_text,       # The action text.
            act_handler,    # The action handler.
            shortcut,       # Optional: the action shortcut
            tooltip,        # Optional: the action tooltip (available in menus/toolbar)
            icon)           # Optional: the action icon (shows when in menus/toolbars)
        register_action(new_action)

    def ql_handle_menu_action(self, action):
        [x.handler() for x in self.menuitems if x.action == action]

    def ql_register_menu_actions(self):
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":start",             self.ql_start,                 "Setup",                      "Setup",                     None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":reloaduserscripts", self.ql_reload_user_script,      "Reload User Scripts",        "Reload User Scripts",       None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem("-",                                     self.ql_menu_null,              "",                           None,                        None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":runtohere",         self.ql_run_to_here,             "Execute Till",               "Execute Till",              None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":runselection",      self.ql_run_selection,           "Execute Selection",         "Execute Selection",              None,                True))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":runfromhere",       self.ql_continue,              "Continue",                   "Continue",                  None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":setpc",             self.ql_set_pc,                "Set PC",                     "Set PC",              None,                True))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":step",              self.ql_step,                  "Step",                       "Step (CTRL+SHIFT+F9)",      "CTRL+SHIFT+F9",        True   ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":changreg",          self.ql_chang_reg,              "Edit Register",              "Edit Register",             None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem("-",                                     self.ql_menu_null,              "",                           None,                        None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":reset",             self.ql_reset,                 "Restart",                    "Restart Qiling",            None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":close",             self.ql_close,                 "Close",                      "Close Qiling",              None,                   False  ))
        self.menuitems.append(QlEmuMisc.MenuItem("-",                                     self.ql_menu_null,              "",                           None,                        None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":reg view",          self.ql_show_reg_view,           "View Register",              "View Register",             None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":stack view",        self.ql_show_stack_view,         "View Stack",                 "View Stack",                None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":memory view",       self.ql_show_mem_view,           "View Memory",                "View Memory",               None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem("-",                                     self.ql_menu_null,              "",                           None,                        None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":save",              self.ql_save,                  "Save Snapshot",              "Save Snapshot",             None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":load",              self.ql_load,                  "Load Snapshot",              "Load Snapshot",             None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem("-",                                     self.ql_menu_null,              "",                           None,                        None,                   True   ))
        if UseAsScript:
            self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":unload",            self.ql_unload_plugin,           "Unload Plugin",              "Unload Plugin",             None,                   False  ))
            self.menuitems.append(QlEmuMisc.MenuItem("-",                                     self.ql_menu_null,              "",                           None,                        None,                   False  ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":about",             self.ql_about,                 "About",                      "About",                     None,                   False  ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":checkupdate",       self.ql_check_update,           "Check Update",               "Check Update",              None,                   False  ))
        self.menuitems.append(QlEmuMisc.MenuItem("-",                                     self.ql_menu_null,              "",                           None,                        None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":parseblocks",       self.ql_parse_blocks_for_deobf,           "Auto Analysis For Deflat",               "Auto Analysis For Deflat",              None,                   True  ))
        # TODO: Dispatcher blocks should support assignment.
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":markreal",       self.ql_mark_real,           "Mark as Real Block",               "Mark as Real Block",              None,                   True  ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":markfake",       self.ql_mark_fake,           "Mark as Fake Block",               "Mark as Fake Block",              None,                   True  ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":markretn",       self.ql_mark_retn,           "Mark as Return Block",               "Mark as Return Block",              None,                   True  ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":deflat",       self.ql_deflat,           "Deflat",               "Deflat",              None,                   True  ))
        self.menuitems.append(QlEmuMisc.MenuItem("-",                                     self.ql_menu_null,              "",                           None,                        None,                   True   ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":removejunkcodebypatterns",       self.ql_remove_junk_code_by_patterns,           "Remove Junk Code by Patterns",               "Remove Junk Code by Patterns",              None,                   True  ))
        self.menuitems.append(QlEmuMisc.MenuItem(self.plugin_name + ":nopinstructions",       self.ql_nop_items_without_colors,           "Nop Items without Color",               "Nop Items without Color",              None,                   True  ))

        for item in self.menuitems:
            if item.action == "-":
                continue
            self.ql_register_new_action(item.action, item.title, QlEmuMisc.menu_action_handler(self, item.action), item.shortcut, item.tooltip,  -1)

    def ql_unregister_menu_actions(self):
        for item in self.menuitems:
            unregister_action(item.action)

    def ql_attach_main_menu_actions(self):
        for item in self.menuitems:
            attach_action_to_menu("Edit/Plugins/" + self.plugin_name + "/" + item.title, item.action, SETMENU_APP)

    def ql_detach_main_menu_actions(self):
        for item in self.menuitems:
            detach_action_from_menu("Edit/Plugins/" + self.plugin_name + "/" + item.title, item.action)

    ### POPUP MENU

    def ql_hook_ui_actions(self):
        self.popup_menu_hook = self
        self.popup_menu_hook.hook()

    def ql_unhook_ui_actions(self):
        if self.popup_menu_hook != None:
            self.popup_menu_hook.unhook()

    # IDA 7.x

    def finish_populating_widget_popup(self, widget, popup_handle):
        if get_widget_type(widget) == BWN_DISASM:
            for item in self.menuitems:
                if item.popup:
                    attach_action_to_popup(widget, popup_handle, item.action, self.plugin_name + "/")

    ### Close View

    def ql_close_reg_view(self):
        self.qlemuregview = None

    def ql_close_stack_view(self):
        self.qlemustackview = None

    def ql_close_mem_view(self, viewid):
        del self.qlemumemview[viewid]

    def ql_close_all_views(self):
        if self.qlemuregview is not None:
            self.qlemuregview.Close()
        if self.qlemustackview is not None:
            self.qlemustackview.Close()

        for viewid in self.qlemumemview:
            self.qlemumemview[viewid].Close()
            self.qlemumemview = None

    def ql_update_views(self, addr, ql):
        if self.qlemuregview is not None:
            self.qlemuregview.SetReg(addr, ql)

        if self.qlemustackview is not None:
            self.qlemustackview.SetStack(self.qlemu.ql)

        for id in self.qlemumemview:
            self.qlemumemview[id].SetMem(self.qlemu.ql)


def PLUGIN_ENTRY():
    qlEmu = QlEmuPlugin()
    return qlEmu

if UseAsScript:
    if __name__ == "__main__":
        qlEmu = QlEmuPlugin()
        qlEmu.init()
        qlEmu.run()


repository: merydwin/idascrtipt
file: Scripts/VirusTotal.py
wanted_name: VirusTotal report
comment: VirusTotal plugin for IDA
created_at: 2015-06-25 13:32:18+00:00
pushed_at: 2015-06-25 13:33:48+00:00
url:
https://github.com/merydwin/idascrtipt/blob/431e04847e55adbb1d263aa2aadc2d489d068f50/Scripts/VirusTotal.py

content:
# -----------------------------------------------------------------------
# VirusTotal IDA Plugin
# By Elias Bachaalany <elias at hex-rays.com>
# (c) Hex-Rays 2011
#
# Special thanks:
# - VirusTotal team
# - Bryce Boe for his VirusTotal Python code
#
import idaapi
import idc
from idaapi import Choose2, plugin_t
import BboeVt as vt
import webbrowser
import urllib
import os


PLUGIN_TEST = 0

# -----------------------------------------------------------------------
# Configuration file
VT_CFGFILE = os.path.join(idaapi.get_user_idadir(), "virustotal.cfg")

# -----------------------------------------------------------------------
# VirusTotal Icon in PNG format
VT_ICON = (
    "\x89\x50\x4E\x47\x0D\x0A\x1A\x0A\x00\x00\x00\x0D\x49\x48\x44\x52"
    "\x00\x00\x00\x10\x00\x00\x00\x10\x04\x03\x00\x00\x00\xED\xDD\xE2"
    "\x52\x00\x00\x00\x30\x50\x4C\x54\x45\x03\x8B\xD3\x5C\xB4\xE3\x9C"
    "\xD1\xED\xF7\xFB\xFD\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
    "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
    "\x00\x00\x00\x00\x00\x00\x00\x00\x00\xD3\xF2\x42\x61\x00\x00\x00"
    "\x4B\x49\x44\x41\x54\x78\x9C\x2D\xCA\xC1\x0D\x80\x30\x0C\x43\x51"
    "\x27\x2C\x50\x89\x05\x40\x2C\x40\xEB\xFD\x77\xC3\x76\xC9\xE9\xEB"
    "\xC5\x20\x5F\xE8\x1A\x0F\x97\xA3\xD0\xE4\x1D\xF9\x49\xD1\x59\x29"
    "\x4C\x43\x9B\xD0\x15\x01\xB5\x4A\x9C\xE4\x70\x14\x39\xB3\x31\xF8"
    "\x15\x70\x04\xF4\xDA\x20\x39\x02\x8A\x0D\xA8\x0F\x94\xA7\x09\x0E"
    "\xC5\x16\x2D\x54\x00\x00\x00\x00\x49\x45\x4E\x44\xAE\x42\x60\x82")


# -----------------------------------------------------------------------
class VirusTotalConfig(object):
    def __init__(self):
        self.Default()


    def Default(self):
        self.md5sum = GetInputMD5()
        self.infile = idaapi.dbg_get_input_path()
        if not self.infile:
            self.infile = ""

        # Persistent options
        self.apikey = ""
        self.options = 1 | 2


    def Read(self):
        """
        Read configuration from file
        """
        if not os.path.exists(VT_CFGFILE):
            return
        f = open(VT_CFGFILE, 'r')
        lines = f.readlines()
        for i, line in enumerate(lines):
            line = line.strip()
            if i == 0:
                self.apikey = line
            elif i == 1:
                self.options = int(line)
            else:
                break


    def Write(self):
        """
        Write configuration to file
        """
        lines = (self.apikey.strip(), str(self.options))
        try:
            f = open(VT_CFGFILE, 'w')
            f.write("\n".join(lines))
            f.close()
        except:
            pass


# -----------------------------------------------------------------------
def VtReport(apikey, filename=None, md5sum=None):
    if filename is None and md5sum is None:
        return (False, "No parameters passed!")

    # Check filename existance
    if filename is not None and not os.path.exists(filename):
        return (False, "Input file '%s' does not exist!" % filename)

    #print("fn=%s md5=%s" % (filename, md5sum))
    # Get file report from VirusTotal
    try:
        vt.set_apikey(apikey)
        result = vt.get_file_report(filename=filename, md5sum=md5sum)
    except Exception as e:
        return (False, "Exception:\n%s" % str(e))

    # Already analyzed?
    if result is not None:
        # Transform the results
        items = []
        for av, mwname in result.items():
            mwname = str(mwname) if mwname else "n/a"
            av = str(av)
            items.append([av, mwname])
        result = items

    return (True, result)


# -----------------------------------------------------------------------
class VirusTotalChooser(Choose2):
    """
    Chooser class to display results from VT
    """
    def __init__(self, title, items, icon, embedded=False):
        Choose2.__init__(self,
                         title,
                         [ ["Antivirus", 20], ["Result", 40] ],
                         embedded=embedded)
        self.items = items
        self.icon = icon


    def GetItems(self):
        return self.items


    def SetItems(self, items):
        self.items = [] if items is None else items


    def OnClose(self):
        pass


    def OnGetLine(self, n):
        return self.items[n]


    def OnGetSize(self):
        return len(self.items)


    def OnSelectLine(self, n):
        # Google search for the malware name and the antivirus name
        s = urllib.urlencode({"q" : " ".join(self.items[n])})
        webbrowser.open_new_tab("http://www.google.com/search?%s" % s)


# --------------------------------------------------------------------------
class VirusTotalForm(Form):
    def __init__(self, icon):
        self.EChooser = VirusTotalChooser("E1", [], icon, embedded=True)
        Form.__init__(self, r"""STARTITEM {id:txtInput}
VirusTotal - IDAPython plugin v1.0 (c) Hex-Rays

{FormChangeCb}
<#API key#~A~pi key:{txtApiKey}>

Options:
<#Open results in a chooser when form closes#~P~opout results on close:{rOptRemember}>
<#Use MD5 checksum#~M~D5Sum:{rOptMD5}>
<#Use file on disk#~F~ile:{rOptFile}>{grpOptions}>

<#Type input (file or MD5 string)#~I~nput:{txtInput}>
<Results:{cEChooser}>
<#Get reports from VT#~R~eport:{btnReport}>
""", {
            'FormChangeCb': Form.FormChangeCb(self.OnFormChange),
            'txtApiKey'   : Form.StringInput(swidth=80),
            'grpOptions'  : Form.ChkGroupControl(("rOptRemember", "rOptMD5", "rOptFile")),
            'txtInput'    : Form.FileInput(open=True),
            'btnReport'   : Form.ButtonInput(self.OnReportClick),
            'cEChooser'   : Form.EmbeddedChooserControl(self.EChooser)
        })



    def OnReportClick(self, code=0):
        pass


    def OnFormChange(self, fid):
        if fid == self.rOptMD5.id or fid == self.rOptFile.id:
            input = (self.cfg.md5sum, self.cfg.infile)
            if fid == self.rOptMD5.id:
                c1 = self.rOptMD5
                c2 = self.rOptFile
                idx = 0
            else:
                c1 = self.rOptFile
                c2 = self.rOptMD5
                idx = 1

            v = not self.GetControlValue(c1)
            if v: idx = not idx

            # Uncheck the opposite input type
            self.SetControlValue(c2, v)

            # Set input field depending on input type
            self.SetControlValue(self.txtInput, input[idx])
        #
        # Report button
        #
        elif fid == self.btnReport.id:
            input = self.GetControlValue(self.txtInput)
            as_file = self.GetControlValue(self.rOptFile)
            apikey = self.GetControlValue(self.txtApiKey)

            ok, r = VtReport(self.cfg.apikey,
                        filename=input if as_file else None,
                        md5sum=None if as_file else input)

            # Error?
            if not ok:
                idc.Warning(r)
                return 1

            # Pass the result
            self.EChooser.SetItems(r)

            # We have results and it was a file? Print its MD5
            if r and as_file:
                print("%s: %s" % (vt.LAST_FILE_HASH, input))

            # Refresh the embedded chooser control
            # (Could also clear previous results if not were retrieved during this run)
            self.RefreshField(self.cEChooser)

            # Store the input for the caller
            self.cfg.input = input

            # No results and file as input was supplied?
            if r is None:
                if as_file:
                    # Propose to upload
                    if idc.AskYN(0, "HIDECANCEL\nNo previous results. Do you want to submit the file:\n\n'%s'\n\nto VirusTotal?" % input) == 0:
                        return 1

                    try:
                        r = vt.scan_file(input)
                    except Exception as e:
                        idc.Warning("Exceptio during upload: %s" % str(e))
                    else:
                        if r is None:
                            idc.Warning("Failed to upload the file!")
                        else:
                            idc.Warning("File uploaded. Check again later to get the analysis report. Scan id: %s" % r)
                else:
                    idc.Warning("No results found for hash: %s" % input)

        return 1


    def Show(self, cfg):
        # Compile the form once
        if not self.Compiled():
            _, args = self.Compile()
            #print args[0]

        # Populate the form
        self.txtApiKey.value  = cfg.apikey
        self.grpOptions.value = cfg.options
        self.txtInput.value   = cfg.infile if self.rOptFile.checked else cfg.md5sum

        # Remember the config
        self.cfg = cfg

        # Execute the form
        ok = self.Execute()

        # Forget the cfg
        del self.cfg

        # Success?
        if ok != 0:
            # Update config
            cfg.options = self.grpOptions.value
            cfg.apikey  = self.txtApiKey.value

            # Popup results?
            if self.rOptRemember.checked:
                ok = 2

        return ok


# -----------------------------------------------------------------------
class VirusTotalPlugin_t(plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = "VirusTotal plugin for IDA"
    help = ""
    wanted_name = "VirusTotal report"
    wanted_hotkey = "Alt-F8"


    def init(self):
        # Some initialization
        self.icon_id = 0
        return idaapi.PLUGIN_OK


    def run(self, arg=0):
        # Load icon from the memory and save its id
        self.icon_id = idaapi.load_custom_icon(data=VT_ICON, format="png")
        if self.icon_id == 0:
            raise RuntimeError("Failed to load icon data!")

        # Create config object
        cfg = VirusTotalConfig()

        # Read previous config
        cfg.Read()

        # Create form
        f = VirusTotalForm(self.icon_id)

        # Show the form
        ok = f.Show(cfg)
        if ok == 0:
            f.Free()
            return

        # Save configuration
        cfg.Write()

        # Spawn a non-modal chooser w/ the results if any
        if ok == 2 and f.EChooser.GetItems():
            VirusTotalChooser(
                "VirusTotal results [%s]" % cfg.input,
                f.EChooser.GetItems(),
                self.icon_id).Show()

        f.Free()
        return


    def term(self):
        # Free the custom icon
        if self.icon_id != 0:
            idaapi.free_custom_icon(self.icon_id)


# -----------------------------------------------------------------------
def PLUGIN_ENTRY():
    return VirusTotalPlugin_t()

# --------------------------------------------------------------------------
if PLUGIN_TEST:
    # Create form
    f = PLUGIN_ENTRY()
    f.init()
    f.run()
    f.term()




repository: mxmssh/idapython
file: Scripts/VirusTotal.py
description: Automatically exported from code.google.com/p/idapython
wanted_name: VirusTotal report
comment: VirusTotal plugin for IDA
created_at: 2015-04-10 06:53:38+00:00
pushed_at: 2015-04-10 06:56:13+00:00
url:
https://github.com/mxmssh/idapython/blob/081b988a03b88867786ad4131269db6930637a5b/Scripts/VirusTotal.py

content:
# -----------------------------------------------------------------------
# VirusTotal IDA Plugin
# By Elias Bachaalany <elias at hex-rays.com>
# (c) Hex-Rays 2011
#
# Special thanks:
# - VirusTotal team
# - Bryce Boe for his VirusTotal Python code
#
import idaapi
import idc
from idaapi import Choose2, plugin_t
import BboeVt as vt
import webbrowser
import urllib
import os


PLUGIN_TEST = 0

# -----------------------------------------------------------------------
# Configuration file
VT_CFGFILE = os.path.join(idaapi.get_user_idadir(), "virustotal.cfg")

# -----------------------------------------------------------------------
# VirusTotal Icon in PNG format
VT_ICON = (
    "\x89\x50\x4E\x47\x0D\x0A\x1A\x0A\x00\x00\x00\x0D\x49\x48\x44\x52"
    "\x00\x00\x00\x10\x00\x00\x00\x10\x04\x03\x00\x00\x00\xED\xDD\xE2"
    "\x52\x00\x00\x00\x30\x50\x4C\x54\x45\x03\x8B\xD3\x5C\xB4\xE3\x9C"
    "\xD1\xED\xF7\xFB\xFD\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
    "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
    "\x00\x00\x00\x00\x00\x00\x00\x00\x00\xD3\xF2\x42\x61\x00\x00\x00"
    "\x4B\x49\x44\x41\x54\x78\x9C\x2D\xCA\xC1\x0D\x80\x30\x0C\x43\x51"
    "\x27\x2C\x50\x89\x05\x40\x2C\x40\xEB\xFD\x77\xC3\x76\xC9\xE9\xEB"
    "\xC5\x20\x5F\xE8\x1A\x0F\x97\xA3\xD0\xE4\x1D\xF9\x49\xD1\x59\x29"
    "\x4C\x43\x9B\xD0\x15\x01\xB5\x4A\x9C\xE4\x70\x14\x39\xB3\x31\xF8"
    "\x15\x70\x04\xF4\xDA\x20\x39\x02\x8A\x0D\xA8\x0F\x94\xA7\x09\x0E"
    "\xC5\x16\x2D\x54\x00\x00\x00\x00\x49\x45\x4E\x44\xAE\x42\x60\x82")


# -----------------------------------------------------------------------
class VirusTotalConfig(object):
    def __init__(self):
        self.Default()


    def Default(self):
        self.md5sum = GetInputMD5()
        self.infile = idaapi.dbg_get_input_path()
        if not self.infile:
            self.infile = ""

        # Persistent options
        self.apikey = ""
        self.options = 1 | 2


    def Read(self):
        """
        Read configuration from file
        """
        if not os.path.exists(VT_CFGFILE):
            return
        f = open(VT_CFGFILE, 'r')
        lines = f.readlines()
        for i, line in enumerate(lines):
            line = line.strip()
            if i == 0:
                self.apikey = line
            elif i == 1:
                self.options = int(line)
            else:
                break


    def Write(self):
        """
        Write configuration to file
        """
        lines = (self.apikey.strip(), str(self.options))
        try:
            f = open(VT_CFGFILE, 'w')
            f.write("\n".join(lines))
            f.close()
        except:
            pass


# -----------------------------------------------------------------------
def VtReport(apikey, filename=None, md5sum=None):
    if filename is None and md5sum is None:
        return (False, "No parameters passed!")

    # Check filename existance
    if filename is not None and not os.path.exists(filename):
        return (False, "Input file '%s' does not exist!" % filename)

    #print("fn=%s md5=%s" % (filename, md5sum))
    # Get file report from VirusTotal
    try:
        vt.set_apikey(apikey)
        result = vt.get_file_report(filename=filename, md5sum=md5sum)
    except Exception as e:
        return (False, "Exception:\n%s" % str(e))

    # Already analyzed?
    if result is not None:
        # Transform the results
        items = []
        for av, mwname in result.items():
            mwname = str(mwname) if mwname else "n/a"
            av = str(av)
            items.append([av, mwname])
        result = items

    return (True, result)


# -----------------------------------------------------------------------
class VirusTotalChooser(Choose2):
    """
    Chooser class to display results from VT
    """
    def __init__(self, title, items, icon, embedded=False):
        Choose2.__init__(self,
                         title,
                         [ ["Antivirus", 20], ["Result", 40] ],
                         embedded=embedded)
        self.items = items
        self.icon = icon


    def GetItems(self):
        return self.items


    def SetItems(self, items):
        self.items = [] if items is None else items


    def OnClose(self):
        pass


    def OnGetLine(self, n):
        return self.items[n]


    def OnGetSize(self):
        return len(self.items)


    def OnSelectLine(self, n):
        # Google search for the malware name and the antivirus name
        s = urllib.urlencode({"q" : " ".join(self.items[n])})
        webbrowser.open_new_tab("http://www.google.com/search?%s" % s)


# --------------------------------------------------------------------------
class VirusTotalForm(Form):
    def __init__(self, icon):
        self.EChooser = VirusTotalChooser("E1", [], icon, embedded=True)
        Form.__init__(self, r"""STARTITEM {id:txtInput}
VirusTotal - IDAPython plugin v1.0 (c) Hex-Rays

{FormChangeCb}
<#API key#~A~pi key:{txtApiKey}>

Options:
<#Open results in a chooser when form closes#~P~opout results on close:{rOptRemember}>
<#Use MD5 checksum#~M~D5Sum:{rOptMD5}>
<#Use file on disk#~F~ile:{rOptFile}>{grpOptions}>

<#Type input (file or MD5 string)#~I~nput:{txtInput}>
<Results:{cEChooser}>
<#Get reports from VT#~R~eport:{btnReport}>
""", {
            'FormChangeCb': Form.FormChangeCb(self.OnFormChange),
            'txtApiKey'   : Form.StringInput(swidth=80),
            'grpOptions'  : Form.ChkGroupControl(("rOptRemember", "rOptMD5", "rOptFile")),
            'txtInput'    : Form.FileInput(open=True),
            'btnReport'   : Form.ButtonInput(self.OnReportClick),
            'cEChooser'   : Form.EmbeddedChooserControl(self.EChooser)
        })



    def OnReportClick(self, code=0):
        pass


    def OnFormChange(self, fid):
        if fid == self.rOptMD5.id or fid == self.rOptFile.id:
            input = (self.cfg.md5sum, self.cfg.infile)
            if fid == self.rOptMD5.id:
                c1 = self.rOptMD5
                c2 = self.rOptFile
                idx = 0
            else:
                c1 = self.rOptFile
                c2 = self.rOptMD5
                idx = 1

            v = not self.GetControlValue(c1)
            if v: idx = not idx

            # Uncheck the opposite input type
            self.SetControlValue(c2, v)

            # Set input field depending on input type
            self.SetControlValue(self.txtInput, input[idx])
        #
        # Report button
        #
        elif fid == self.btnReport.id:
            input = self.GetControlValue(self.txtInput)
            as_file = self.GetControlValue(self.rOptFile)
            apikey = self.GetControlValue(self.txtApiKey)

            ok, r = VtReport(self.cfg.apikey,
                        filename=input if as_file else None,
                        md5sum=None if as_file else input)

            # Error?
            if not ok:
                idc.Warning(r)
                return 1

            # Pass the result
            self.EChooser.SetItems(r)

            # We have results and it was a file? Print its MD5
            if r and as_file:
                print("%s: %s" % (vt.LAST_FILE_HASH, input))

            # Refresh the embedded chooser control
            # (Could also clear previous results if not were retrieved during this run)
            self.RefreshField(self.cEChooser)

            # Store the input for the caller
            self.cfg.input = input

            # No results and file as input was supplied?
            if r is None:
                if as_file:
                    # Propose to upload
                    if idc.AskYN(0, "HIDECANCEL\nNo previous results. Do you want to submit the file:\n\n'%s'\n\nto VirusTotal?" % input) == 0:
                        return 1

                    try:
                        r = vt.scan_file(input)
                    except Exception as e:
                        idc.Warning("Exceptio during upload: %s" % str(e))
                    else:
                        if r is None:
                            idc.Warning("Failed to upload the file!")
                        else:
                            idc.Warning("File uploaded. Check again later to get the analysis report. Scan id: %s" % r)
                else:
                    idc.Warning("No results found for hash: %s" % input)

        return 1


    def Show(self, cfg):
        # Compile the form once
        if not self.Compiled():
            _, args = self.Compile()
            #print args[0]

        # Populate the form
        self.txtApiKey.value  = cfg.apikey
        self.grpOptions.value = cfg.options
        self.txtInput.value   = cfg.infile if self.rOptFile.checked else cfg.md5sum

        # Remember the config
        self.cfg = cfg

        # Execute the form
        ok = self.Execute()

        # Forget the cfg
        del self.cfg

        # Success?
        if ok != 0:
            # Update config
            cfg.options = self.grpOptions.value
            cfg.apikey  = self.txtApiKey.value

            # Popup results?
            if self.rOptRemember.checked:
                ok = 2

        return ok


# -----------------------------------------------------------------------
class VirusTotalPlugin_t(plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = "VirusTotal plugin for IDA"
    help = ""
    wanted_name = "VirusTotal report"
    wanted_hotkey = "Alt-F8"


    def init(self):
        # Some initialization
        self.icon_id = 0
        return idaapi.PLUGIN_OK


    def run(self, arg=0):
        # Load icon from the memory and save its id
        self.icon_id = idaapi.load_custom_icon(data=VT_ICON, format="png")
        if self.icon_id == 0:
            raise RuntimeError("Failed to load icon data!")

        # Create config object
        cfg = VirusTotalConfig()

        # Read previous config
        cfg.Read()

        # Create form
        f = VirusTotalForm(self.icon_id)

        # Show the form
        ok = f.Show(cfg)
        if ok == 0:
            f.Free()
            return

        # Save configuration
        cfg.Write()

        # Spawn a non-modal chooser w/ the results if any
        if ok == 2 and f.EChooser.GetItems():
            VirusTotalChooser(
                "VirusTotal results [%s]" % cfg.input,
                f.EChooser.GetItems(),
                self.icon_id).Show()

        f.Free()
        return


    def term(self):
        # Free the custom icon
        if self.icon_id != 0:
            idaapi.free_custom_icon(self.icon_id)


# -----------------------------------------------------------------------
def PLUGIN_ENTRY():
    return VirusTotalPlugin_t()

# --------------------------------------------------------------------------
if PLUGIN_TEST:
    # Create form
    f = PLUGIN_ENTRY()
    f.init()
    f.run()
    f.term()




repository: hoyori/idapython
file: Scripts/VirusTotal.py
description: Automatically exported from code.google.com/p/idapython
wanted_name: VirusTotal report
comment: VirusTotal plugin for IDA
created_at: 2015-04-04 07:26:57+00:00
pushed_at: 2015-04-04 07:30:32+00:00
url:
https://github.com/hoyori/idapython/blob/081b988a03b88867786ad4131269db6930637a5b/Scripts/VirusTotal.py


repository: Hehouhua/idapython
file: Scripts/VirusTotal.py
description: Automatically exported from code.google.com/p/idapython
wanted_name: VirusTotal report
comment: VirusTotal plugin for IDA
created_at: 2016-04-08 08:07:27+00:00
pushed_at: 2016-04-08 08:10:01+00:00
url:
https://github.com/Hehouhua/idapython/blob/518dc295c8a9747834ddbe19566c2d2924d8ad01/Scripts/VirusTotal.py


repository: HACK-THE-WORLD/IDAPluginList
file: auto_re/auto_re.py
description: IDA插件集合，含项目名称及项目地址，每日定时Clone项目。
wanted_name: Auto RE
created_at: 2022-03-21 02:05:27+00:00
pushed_at: 2024-09-21 02:08:36+00:00
url:
https://github.com/HACK-THE-WORLD/IDAPluginList/blob/a6367e2126fca45260d1052b0c54bd332490fd60/auto_re/auto_re.py

content:
# -*- coding: utf-8 -*
__author__ = 'Trafimchuk Aliaksandr'
__version__ = '1.9'

from collections import defaultdict
import idaapi
from idautils import FuncItems, CodeRefsTo
from idaapi import o_reg, o_imm, o_far, o_near, o_mem, o_displ
import os
import re
import sys
import traceback


HAS_PYSIDE = idaapi.IDA_SDK_VERSION < 690
if HAS_PYSIDE:
    from PySide import QtGui, QtCore
    from PySide.QtGui import QTreeView, QVBoxLayout, QLineEdit, QMenu, QInputDialog, QAction, QTabWidget
else:
    from PyQt5 import QtGui, QtCore
    from PyQt5.QtWidgets import QTreeView, QVBoxLayout, QLineEdit, QMenu, QInputDialog, QAction, QTabWidget


try:
    # Python 2.
    xrange
except NameError:
    # Python 3.
    xrange = range


# enable to allow PyCharm remote debug
RDEBUG = False
# adjust this value to be a full path to a debug egg
RDEBUG_EGG = r'c:\Program Files\JetBrains\PyCharm 2017.1.4\debug-eggs\pycharm-debug.egg'
RDEBUG_HOST = 'localhost'
RDEBUG_PORT = 12321


TAGS_IGNORE_LIST = {
    'OpenProcessToken',
    'DisconnectNamedPipe'
}

IGNORE_CALL_LIST = {
    'RtlNtStatusToDosError',
    'GetLastError',
    'SetLastError'
}

TAGS = {
    'net': ['WSAStartup', 'socket', 'recv', 'recvfrom', 'send', 'sendto', 'acccept', 'bind', 'listen', 'select',
            'setsockopt', 'ioctlsocket', 'closesocket', 'WSAAccept', 'WSARecv', 'WSARecvFrom', 'WSASend', 'WSASendTo',
            'WSASocket', 'WSAConnect', 'ConnectEx', 'TransmitFile', 'HTTPOpenRequest', 'HTTPSendRequest',
            'URLDownloadToFile', 'InternetCrackUrl', 'InternetOpen', 'InternetOpen', 'InternetConnect',
            'InternetOpenUrl', 'InternetQueryOption', 'InternetSetOption', 'InternetReadFile', 'InternetWriteFile',
            'InternetGetConnectedState', 'InternetSetStatusCallback', 'DnsQuery', 'getaddrinfo', 'GetAddrInfo',
            'GetAdaptersInfo', 'GetAdaptersAddresses', 'HttpQueryInfo', 'ObtainUserAgentString', 'WNetGetProviderName',
            'GetBestInterfaceEx', 'gethostbyname', 'getsockname', 'connect', 'WinHttpOpen', 'WinHttpSetTimeouts',
            'WinHttpSendRequest', 'WinHttpConnect', 'WinHttpCrackUrl', 'WinHttpReadData', 'WinHttpOpenRequest',
            'WinHttpReceiveResponse', 'WinHttpQueryHeaders', 'HttpSendRequestW', 'HttpSendRequestA', 'HttpAddRequestHeadersW', 'HttpAddRequestHeadersA', 'HttpOpenRequestW', 'HttpOpenRequestA', 'NetServerGetInfo', 'NetApiBufferFree', 'NetWkstaGetInfo'],
    'spawn': ['CreateProcess', 'ShellExecute', 'ShellExecuteEx', 'system', 'CreateProcessInternal', 'NtCreateProcess',
              'ZwCreateProcess', 'NtCreateProcessEx', 'ZwCreateProcessEx', 'NtCreateUserProcess', 'ZwCreateUserProcess',
              'RtlCreateUserProcess', 'NtCreateSection', 'ZwCreateSection', 'NtOpenSection', 'ZwOpenSection',
              'NtAllocateVirtualMemory', 'ZwAllocateVirtualMemory', 'NtWriteVirtualMemory', 'ZwWriteVirtualMemory',
              'NtMapViewOfSection', 'ZwMapViewOfSection', 'OpenSCManager', 'CreateService', 'OpenService',
              'StartService', 'ControlService', 'ShellExecuteExA', 'ShellExecuteExW'],
    'inject': ['OpenProcess-disabled', 'ZwOpenProcess', 'NtOpenProcess', 'WriteProcessMemory', 'NtWriteVirtualMemory',
               'ZwWriteVirtualMemory', 'CreateRemoteThread', 'QueueUserAPC', 'ZwUnmapViewOfSection', 'NtUnmapViewOfSection'],
    'com': ['CoCreateInstance', 'CoInitializeSecurity', 'CoGetClassObject', 'OleConvertOLESTREAMToIStorage', 'CreateBindCtx', 'CoSetProxyBlanket', 'VariantClear'],
    'crypto': ['CryptAcquireContext', 'CryptProtectData', 'CryptUnprotectData', 'CryptProtectMemory',
               'CryptUnprotectMemory', 'CryptDecrypt', 'CryptEncrypt', 'CryptHashData', 'CryptDecodeMessage',
               'CryptDecryptMessage', 'CryptEncryptMessage', 'CryptHashMessage', 'CryptExportKey', 'CryptGenKey',
               'CryptCreateHash', 'CryptDecodeObjectEx', 'EncryptMessage', 'DecryptMessage'],
    'kbd': ['SendInput', 'VkKeyScanA', 'VkKeyScanW'],
    'file': ['_open64', 'open64', 'open', 'open64', 'fopen', 'fread', 'fclose', 'fwrite', 'flock', 'read', 'write',
             'fstat', 'lstat', 'stat', 'chmod', 'chown', 'lchown', 'link', 'symlink', 'readdir', 'readdir64'],
    'reg': ['RegOpenKeyExW', 'RegQueryValueExW', 'RegSetValueExW', 'RegCreateKeyExW', 'RegDeleteValueW', 'RegEnumKeyW', 'RegCloseKey', 'RegQueryInfoKeyW', 'RegOpenKeyExA', 'RegQueryValueExA', 'RegSetValueExA', 'RegCreateKeyExA', 'RegDeleteValueA', 'RegEnumKeyA',  'RegQueryInfoKeyA'],
    'dev': ['DeviceIoControl'],
    'wow': ['Wow64DisableWow64FsRedirection', 'Wow64RevertWow64FsRedirection']
}

STRICT_TAG_NAME_CHECKING = {'file'}

blacklist = {'@__security_check_cookie@4', '__SEH_prolog4', '__SEH_epilog4'}
replacements = [
    ('??3@YAXPAX@Z', 'alloc'),
    ('?', '')
]

def inf_is_64bit():
    return (idaapi.inf_is_64bit if idaapi.IDA_SDK_VERSION >= 900 else idaapi.cvar.inf.is_64bit)()


def get_addr_width():
    return '16' if inf_is_64bit() else '8'


def decode_insn(ea):
    if idaapi.IDA_SDK_VERSION >= 700 and sys.maxsize > 2**32:
        insn = idaapi.insn_t()
        if idaapi.decode_insn(insn, ea) > 0:
            return insn
    else:
        if idaapi.decode_insn(ea):
            return idaapi.cmd.copy()


def force_name(ea, new_name):
    if not ea or ea == idaapi.BADADDR:
        return
    if idaapi.IDA_SDK_VERSION >= 700:
        return idaapi.force_name(ea, new_name, idaapi.SN_NOCHECK)
    return idaapi.do_name_anyway(ea, new_name, 0)


class AutoReIDPHooks(idaapi.IDP_Hooks):
    """
    Hooks to keep view updated if some function is updated
    """
    def __init__(self, view, *args):
        super(AutoReIDPHooks, self).__init__(*args)
        self._view = view

    def __on_rename(self, ea, new_name):
        if not self._view:
            return
        items = self._view._model.findItems(('%0' + get_addr_width() + 'X') % ea, QtCore.Qt.MatchRecursive)
        if len(items) != 1:
            return

        item = items[0]
        index = self._view._model.indexFromItem(item)
        if not index.isValid():
            return

        name_index = index.sibling(index.row(), 1)
        if not name_index.isValid():
            return

        self._view._model.setData(name_index, new_name)

    def ev_rename(self, ea, new_name):
        """ callback for IDA >= 700 """
        self.__on_rename(ea, new_name)
        return super(AutoReIDPHooks, self).ev_rename(ea, new_name)

    def rename(self, ea, new_name):
        """ callback for IDA < 700 """
        self.__on_rename(ea, new_name)
        return super(AutoReIDPHooks, self).rename(ea, new_name)


class AutoREView(idaapi.PluginForm):
    ADDR_ROLE = QtCore.Qt.UserRole + 1

    OPT_FORM_PERSIST = idaapi.PluginForm.FORM_PERSIST if hasattr(idaapi.PluginForm, 'FORM_PERSIST') else idaapi.PluginForm.WOPN_PERSIST
    OPT_FORM_NO_CONTEXT = idaapi.PluginForm.FORM_NO_CONTEXT if hasattr(idaapi.PluginForm, 'FORM_NO_CONTEXT') else idaapi.PluginForm.WCLS_NO_CONTEXT

    def __init__(self, data):
        super(AutoREView, self).__init__()
        self._data = data
        self.tv = None
        self._model = None
        self._idp_hooks = None

    def Show(self):
        return idaapi.PluginForm.Show(self, 'AutoRE', options=self.OPT_FORM_PERSIST)

    def _get_parent_widget(self, form):
        if HAS_PYSIDE:
            return self.FormToPySideWidget(form)
        return self.FormToPyQtWidget(form)

    def OnCreate(self, form):
        self.parent = self._get_parent_widget(form)

        self._idp_hooks = AutoReIDPHooks(self)
        if not self._idp_hooks.hook():
            print('IDP_Hooks.hook() failed')

        self.tv = QTreeView()
        self.tv.setExpandsOnDoubleClick(False)

        root_layout = QVBoxLayout(self.parent)
        # self.le_filter = QLineEdit(self.parent)

        # root_layout.addWidget(self.le_filter)
        root_layout.addWidget(self.tv)

        self.parent.setLayout(root_layout)

        self._model = QtGui.QStandardItemModel()
        self._init_model()
        self.tv.setModel(self._model)

        self.tv.setColumnWidth(0, 200)
        self.tv.setColumnWidth(1, 300)
        self.tv.header().setStretchLastSection(True)

        self.tv.expandAll()

        self.tv.doubleClicked.connect(self.on_navigate_to_method_requested)
        # self.le_filter.textChanged.connect(self.on_filter_text_changed)
        self.tv.setContextMenuPolicy(QtCore.Qt.CustomContextMenu)
        self.tv.customContextMenuRequested.connect(self._tree_customContextMenuRequesssted)

        rename_action = QAction('Rename...', self.tv)
        rename_action.setShortcut('n')
        rename_action.triggered.connect(self._tv_rename_action_triggered)
        self.tv.addAction(rename_action)

    def _tree_customContextMenuRequesssted(self, pos):
        idx = self.tv.indexAt(pos)
        if not idx.isValid():
            return

        addr = idx.data(role=self.ADDR_ROLE)
        if not addr:
            return

        name_idx = idx.sibling(idx.row(), 1)
        old_name = name_idx.data()

        menu = QMenu()
        rename_action = menu.addAction('Rename `%s`...' % old_name)
        rename_action.setShortcut('n')
        action = menu.exec_(self.tv.mapToGlobal(pos))
        if action == rename_action:
            return self._rename_ea_requested(addr, name_idx)

    def _tv_rename_action_triggered(self):
        selected = self.tv.selectionModel().selectedIndexes()
        if not selected:
            return

        idx = selected[0]
        if not idx.isValid():
            return

        addr = idx.data(role=self.ADDR_ROLE)
        if not addr:
            return

        name_idx = idx.sibling(idx.row(), 1)
        if not name_idx.isValid():
            return

        return self._rename_ea_requested(addr, name_idx)

    def _rename_ea_requested(self, addr, name_idx):
        old_name = name_idx.data()

        if idaapi.IDA_SDK_VERSION >= 700:
            new_name = idaapi.ask_str(str(old_name), 0, 'New name:')
        else:
            new_name = idaapi.askstr(0, str(old_name), 'New name:')

        if new_name is None:
            return

        force_name(addr, new_name)
        renamed_name = idaapi.get_ea_name(addr)
        name_idx.model().setData(name_idx, renamed_name)

    def OnClose(self, form):
        if self._idp_hooks:
            self._idp_hooks.unhook()

    def _tv_init_header(self, model):
        item_header = QtGui.QStandardItem("EA")
        item_header.setToolTip("Address")
        model.setHorizontalHeaderItem(0, item_header)

        item_header = QtGui.QStandardItem("Function name")
        model.setHorizontalHeaderItem(1, item_header)

        item_header = QtGui.QStandardItem("API called")
        model.setHorizontalHeaderItem(2, item_header)

    # noinspection PyMethodMayBeStatic
    def _tv_make_tag_item(self, name):
        rv = QtGui.QStandardItem(name)

        rv.setEditable(False)
        return [rv, QtGui.QStandardItem(), QtGui.QStandardItem()]

    def _tv_make_ref_item(self, tag, ref):
        ea_item = QtGui.QStandardItem(('%0' + get_addr_width() + 'X') % ref['ea'])
        ea_item.setEditable(False)
        ea_item.setData(ref['ea'], self.ADDR_ROLE)

        name_item = QtGui.QStandardItem(ref['name'])
        name_item.setEditable(False)
        name_item.setData(ref['ea'], self.ADDR_ROLE)

        apis = ', '.join(ref['tags'][tag])
        api_name = QtGui.QStandardItem(apis)
        api_name.setEditable(False)
        api_name.setData(ref['ea'], self.ADDR_ROLE)
        api_name.setToolTip(apis)

        return [ea_item, name_item, api_name]

    def _init_model(self):
        self._model.clear()

        root_node = self._model.invisibleRootItem()
        self._tv_init_header(self._model)

        for tag, refs in self._data.items():
            item_tag_list = self._tv_make_tag_item(tag)
            item_tag = item_tag_list[0]

            root_node.appendRow(item_tag_list)

            for ref in refs:
                ref_item_list = self._tv_make_ref_item(tag, ref)

                item_tag.appendRow(ref_item_list)

    def on_navigate_to_method_requested(self, index):
        addr = index.data(role=self.ADDR_ROLE)
        if addr is not None:
            idaapi.jumpto(addr)

    # def on_filter_text_changed(self, text):
    #     print('on_text_changed: %s' % text)


class auto_re_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = ""

    help = ""
    wanted_name = "Auto RE"
    wanted_hotkey = "Ctrl+Shift+M"

    _PREFIX_NAME = 'au_re_'
    _MIN_MAX_MATH_OPS_TO_ALLOW_RENAME = 10

    _CALLEE_NODE_NAMES = {
        idaapi.PLFM_MIPS: '$ mips',
        idaapi.PLFM_ARM: '$ arm'
    }
    _DEFAULT_CALLEE_NODE_NAME = '$ vmm functions'

    _JMP_TYPES = {idaapi.NN_jmp, idaapi.NN_jmpni, idaapi.NN_jmpfi, idaapi.NN_jmpshort}

    def __init__(self):
        super(auto_re_t, self).__init__()
        self._data = None
        self.view = None

    def init(self):
        # self._cfg = None
        self.view = None
        # self._load_config()

        return idaapi.PLUGIN_OK

    # def _load_config(self):
    #     self._cfg = {'auto_rename': False}

    # def _store_config(self, cfg):
    #     pass

    def _handle_tags(self, fn, fn_an, known_refs):
        if known_refs:
            known_refs = dict(known_refs)
            for k, names in known_refs.items():
                existing = set(fn_an['tags'][k])
                new = set(names) - existing
                if new:
                    fn_an['tags'][k] += list(new)

        tags = dict(fn_an['tags'])
        if not tags:
            return
        print('fn: %#08x tags: %s' % (self.start_ea_of(fn), tags))
        cmt = idaapi.get_func_cmt(fn, True)
        if cmt:
            cmt += '\n'
        s = str(tags.keys())
        name = idaapi.get_ea_name(self.start_ea_of(fn))
        item = {'ea': self.start_ea_of(fn), 'name': name, 'tags': tags}
        if not cmt or s not in cmt:
            idaapi.set_func_cmt(fn, '%sTAGS: %s' % (cmt or '', s), True)
        # self.mark_position(self.start_ea_of(fn), 'TAGS: %s' % s)
        for tag in tags:
            if tag not in self._data:
                self._data[tag] = list()
            self._data[tag].append(item)

    def _handle_calls(self, fn, fn_an):
        num_calls = len(fn_an['calls'])
        if num_calls != 1:
            return

        dis = fn_an['calls'][0]
        if dis.Op1.type not in (o_imm, o_far, o_near, o_mem):
            return

        ea = dis.Op1.value
        if not ea and dis.Op1.addr:
            ea = dis.Op1.addr

        if idaapi.has_dummy_name(self.get_flags_at(ea)):
            return

        # TODO: check is there jmp, push+retn then don't rename the func
        if fn_an['strange_flow']:
            return

        possible_name = idaapi.get_ea_name(ea)
        if not possible_name or possible_name in blacklist:
            return

        normalized = self.normalize_name(possible_name)

        # if self._cfg.get('auto_rename'):
        if len(fn_an['math']) < self._MIN_MAX_MATH_OPS_TO_ALLOW_RENAME:
            force_name(self.start_ea_of(fn), normalized)
        # TODO: add an API to the view
        print('fn: %#08x: %d calls, %d math%s possible name: %s, normalized: %s' % (
            self.start_ea_of(fn), len(fn_an['calls']), len(fn_an['math']), 'has bads' if fn_an['has_bads'] else '',
            possible_name, normalized))

    # noinspection PyMethodMayBeStatic
    def _check_is_jmp_wrapper(self, dis):
        # checks instructions like `jmp API`
        if dis.itype not in self._JMP_TYPES:
            return

        # handle call wrappers like jmp GetProcAddress
        if dis.Op1.type == idaapi.o_mem and dis.Op1.addr:
            # TODO: check is there better way to determine is the function a wrapper
            v = dis.Op1.addr
            flags = self.get_flags_at(v)
            if v and dis.itype == idaapi.NN_jmpni and self.is_data(flags) and self.__is_ptr_val(flags):
                v = self.__get_ptr_val(v)
            return v

    # noinspection PyMethodMayBeStatic
    def _check_is_push_retn_wrapper(self, dis0, dis1):
        """
        Checks for sequence of push IMM32/retn
        :param dis0: the first insn
        :param dis1: the second insn
        :return: value of IMM32
        """
        if dis0.itype != idaapi.NN_push or dis0.Op1.type != idaapi.o_imm or not dis0.Op1.value:
            return

        if dis1.itype not in (idaapi.NN_retn,):
            return

        return dis0.Op1.value

    def _preprocess_api_wrappers(self, fnqty):
        rv = defaultdict(dict)

        for i in xrange(fnqty):
            fn = idaapi.getn_func(i)
            items = list(FuncItems(self.start_ea_of(fn)))
            if len(items) not in (1, 2):
                continue

            dis0 = decode_insn(items[0])
            if dis0 is None:
                continue
            addr = self._check_is_jmp_wrapper(dis0)

            if not addr and len(items) > 1:
                dis1 = decode_insn(items[1])
                if dis1 is not None:
                    addr = self._check_is_push_retn_wrapper(dis0, dis1)

            if not addr:
                continue

            name = idaapi.get_ea_name(addr)
            name = name.replace(idaapi.FUNC_IMPORT_PREFIX, '')
            if not name:
                continue

            imp_stripped_name = name.lstrip('_')

            for tag, names in TAGS.items():
                for tag_api in names:
                    if tag in STRICT_TAG_NAME_CHECKING:
                        match = tag_api in (name, imp_stripped_name)
                    else:
                        match = tag_api in name
                    if not match:
                        continue

                    refs = list(CodeRefsTo(self.start_ea_of(fn), 1))

                    for ref in refs:
                        ref_fn = idaapi.get_func(ref)
                        if not ref_fn:
                            # idaapi.msg('AutoRE: there is no func for ref: %08x for api: %s' % (ref, name))
                            continue
                        if tag not in rv[self.start_ea_of(ref_fn)]:
                            rv[self.start_ea_of(ref_fn)][tag] = list()
                        if name not in rv[self.start_ea_of(ref_fn)][tag]:
                            rv[self.start_ea_of(ref_fn)][tag].append(name)
        return dict(rv)

    def run(self, arg):
        if RDEBUG and RDEBUG_EGG:
            if not os.path.isfile(RDEBUG_EGG):
                idaapi.msg('AutoRE: Remote debug is enabled, but I cannot find the debug egg: %s' % RDEBUG_EGG)
            else:
                import sys

                if RDEBUG_EGG not in sys.path:
                    sys.path.append(RDEBUG_EGG)

                import pydevd
                pydevd.settrace(RDEBUG_HOST, port=RDEBUG_PORT, stdoutToServer=True, stderrToServer=True)

        try:
            self._data = dict()
            count = idaapi.get_func_qty()

            # pre-process of api wrapper functions
            known_refs_tags = self._preprocess_api_wrappers(count)

            for i in xrange(count):
                fn = idaapi.getn_func(i)
                fn_an = self.analyze_func(fn)

                # if fn_an['math']:
                # 	print('fn: %#08x has math' % self.start_ea_of(fn))

                if idaapi.has_dummy_name(self.get_flags_at(self.start_ea_of(fn))):
                    self._handle_calls(fn, fn_an)

                known_refs = known_refs_tags.get(self.start_ea_of(fn))
                self._handle_tags(fn, fn_an, known_refs)

            if self.view:
                self.view.Close(AutoREView.OPT_FORM_NO_CONTEXT)
            self.view = AutoREView(self._data)
            self.view.Show()
        except:
            idaapi.msg('AutoRE: error: %s\n' % traceback.format_exc())

    def term(self):
        self._data = None

    @classmethod
    def disasm_func(cls, fn):
        rv = list()
        items = list(FuncItems(cls.start_ea_of(fn)))
        for item_ea in items:
            obj = {'ea': item_ea, 'fn_ea': cls.start_ea_of(fn), 'dis': None}
            insn = decode_insn(item_ea)
            if insn is not None:
                obj['dis'] = insn
            rv.append(obj)
        return rv

    @classmethod
    def get_callee_netnode(cls):
        node_name = cls._CALLEE_NODE_NAMES.get(idaapi.ph.id, cls._DEFAULT_CALLEE_NODE_NAME)
        n = idaapi.netnode(node_name)
        return n

    @classmethod
    def get_callee(cls, ea):
        n = cls.get_callee_netnode()
        v = n.altval(ea)
        v -= 1
        if v == idaapi.BADNODE:
            return
        return v

    @classmethod
    def _analysis_handle_call_insn(cls, dis, rv):
        rv['calls'].append(dis)
        if dis.Op1.type != o_mem or not dis.Op1.addr:
            callee = cls.get_callee(dis.ip)
            if not callee:
                return
        else:
            callee = dis.Op1.addr

        cls._apply_tag_on_callee(callee, rv, is_call=True)

    @classmethod
    def _apply_tag_on_callee(cls, callee_ea, rv, is_call=False):
        name = idaapi.get_ea_name(callee_ea)
        name = name.replace(idaapi.FUNC_IMPORT_PREFIX, '')

        if '@' in name:
            name = name.split('@')[0]

        if not name:
            return

        if name in IGNORE_CALL_LIST:
            if is_call:
                rv['calls'].pop()
            return

        if name in TAGS_IGNORE_LIST:
            return

        for tag, names in TAGS.items():
            for tag_api in names:
                if tag in STRICT_TAG_NAME_CHECKING:
                    match = tag_api in (name, name.lstrip('_'))
                else:
                    match = tag_api in name
                if not match or name in rv['tags'][tag]:
                    continue

                # print('%#08x: %s, tag: %s' % (dis.ea, name, tag))
                rv['tags'][tag].append(name)
                break

    @classmethod
    def __is_ptr_val(cls, flags):
        if idaapi.IDA_SDK_VERSION >= 700:
            return (idaapi.is_qword if inf_is_64bit() else idaapi.is_dword)(flags)
        return (idaapi.isQwrd if inf_is_64bit() else idaapi.isDwrd)(flags)

    @classmethod
    def __get_ptr_val(cls, ea):
        if inf_is_64bit():
            return idaapi.get_qword(ea)

        return (idaapi.get_dword if idaapi.IDA_SDK_VERSION >= 700 else idaapi.get_long)(ea)
        
    @classmethod
    def start_ea_of(cls, o):
        return getattr(o, 'start_ea' if idaapi.IDA_SDK_VERSION >= 700 else 'startEA')

    @classmethod
    def end_ea_of(cls, o):
        return getattr(o, 'end_ea' if idaapi.IDA_SDK_VERSION >= 700 else 'endEA')

    @classmethod
    def get_flags_at(cls, ea):
        return getattr(idaapi, 'get_flags' if idaapi.IDA_SDK_VERSION >= 700 else 'getFlags')(ea)

    @classmethod
    def is_data(cls, flags):
        return getattr(idaapi, 'is_data' if idaapi.IDA_SDK_VERSION >= 700 else 'isData')(flags)

    @classmethod
    def analyze_func(cls, fn):
        rv = {
            'fn': fn,
            'calls': [],
            'math': [],
            'has_bads': False,
            'strange_flow': False,
            'tags': defaultdict(list)
        }
        items = cls.disasm_func(fn)
        items_set = set(map(lambda x: x['ea'], items))

        for item in items:
            dis = item['dis']
            if dis is None:
                rv['has_bads'] = True
                continue

            if dis.itype in (idaapi.NN_call, idaapi.NN_callfi, idaapi.NN_callni):
                cls._analysis_handle_call_insn(dis, rv)
            elif dis.itype == idaapi.NN_xor:
                if dis.Op1.type == o_reg and dis.Op2.type == o_reg and dis.Op1.reg == dis.Op2.reg:
                    continue
                rv['math'].append(dis)
            elif dis.itype in (idaapi.NN_shr, idaapi.NN_shl, idaapi.NN_sal, idaapi.NN_sar, idaapi.NN_ror,
                               idaapi.NN_rol, idaapi.NN_rcl, idaapi.NN_rcl):
                # TODO
                rv['math'].append(dis)
            elif dis.itype in cls._JMP_TYPES:
                if dis.Op1.type not in (o_far, o_near, o_mem, o_displ):
                    continue

                if dis.Op1.type == o_displ:
                    rv['strange_flow'] = True
                    continue

                ea = dis.Op1.value
                if not ea and dis.Op1.addr:
                    ea = dis.Op1.addr
                if ea not in items_set:
                    rv['strange_flow'] = True

                # flags = self.get_flags_at(ea)
                # if dis.itype == idaapi.NN_jmpni and dis.Op1.type == o_mem and ea and self.is_data(flags):
                #     if cls.__is_ptr_val(flags):
                #         val = cls.__get_ptr_val(ea)
                #         if val:
                #             cls._apply_tag_on_callee(val, rv, is_call=False)

        return rv

    @classmethod
    def normalize_name(cls, n):
        for repl in replacements:
            n = n.replace(*repl)
        if '@' in n:
            n = n.split('@')[0]
        if len(n) < 3:
            return ''
        if not n.startswith(cls._PREFIX_NAME):
            n = cls._PREFIX_NAME + n
        return n

    # @classmethod
    # def mark_position(cls, ea, name, slot=[0]):
    #     curloc = idaapi.curloc()
    #     curloc.ea = ea
    #     curloc.lnnum = 0
    #     curloc.x = 0
    #     curloc.y = 0
    #     slot[0] += 1
    #     curloc.mark(slot[0], name, name)


# noinspection PyPep8Naming
def PLUGIN_ENTRY():
    return auto_re_t()


repository: shaoxi2010/useful-scripts
file: IDA/auto_re.py
description: 一些实用脚本
wanted_name: Auto RE
created_at: 2024-03-06 16:47:38+00:00
pushed_at: 2025-03-16 16:48:32+00:00
url:
https://github.com/shaoxi2010/useful-scripts/blob/e947816ffb7c9658f7b374276e4f3e5865a6a7a3/IDA/auto_re.py


repository: tbxy09/mdpickme
file: ftplugin/markdown/mdnav.py
description: A easy use Markdown format tool in Vim
created_at: 2018-09-11 01:56:36+00:00
pushed_at: 2018-09-11 02:21:26+00:00
url:
https://github.com/tbxy09/mdpickme/blob/8dbdaa4ab0e8b62798d40d301898a34e3c9c699a/ftplugin/markdown/mdnav.py

content:
from __future__ import print_function

import collections
import json
import os.path
import re
import sys
import subprocess
import webbrowser

try:
    from urllib.parse import urlparse
except ImportError:
    from urlparse import urlparse


class FakeLogger(object):
    def __init__(self, active=False):
        self.active = active

    def info(self, fmt, *args):
        if not self.active:
            return

        print(fmt % args)


_logger = FakeLogger()


def plugin_entry_point():
    import vim

    if int(vim.eval("exists('g:mdnav#Extensions')")):
        extensions = vim.eval('g:mdnav#Extensions')
        extensions = [ext.strip() for ext in extensions.split(',')]

    else:
        extensions = []

    if int(vim.eval("exists('g:mdnav#DebugMode')")):
        _logger.active = vim.eval('g:mdnav#DebugMode') == 'true'

    row, col = vim.current.window.cursor
    cursor = (row - 1, col)
    lines = vim.current.buffer

    target = parse_link(cursor, lines)
    _logger.info('open %s', target)
    action = open_link(
        target,
        current_file=vim.eval("expand('%:p')"),
        open_in_vim_extensions=extensions,
    )
    action()


def open_link(target, current_file, open_in_vim_extensions=set()):
    """
    :returns: a callable that encapsulates the action to perform
    """
    if target is not None:
        target = target.strip()

    if not target:
        _logger.info('no target')
        return NoOp(target)

    if target.startswith('#'):
        return JumpToAnchor(target)

    if has_scheme(target):
        _logger.info('has scheme -> open in browser')
        return BrowserOpen(target)

    if not has_extension(target, open_in_vim_extensions):
        _logger.info('has no extension for opening in vim')
        return OSOpen(anchor_path(target, current_file))

    if target.startswith('|filename|'):
        target = target[len('|filename|'):]

    return VimOpen(anchor_path(target, current_file))


def anchor_path(target, current_file):
    if os.path.isabs(target):
        return target

    _logger.info('anchor path relative to %s', current_file)
    return os.path.join(os.path.dirname(current_file), target)


def has_extension(path, extensions):
    if not extensions:
        return True

    path = parse_path(path)
    _, ext = os.path.splitext(path.path)
    return ext in extensions


def has_scheme(target):
    return bool(urlparse(target).scheme)


class Action(object):
    def __init__(self, target):
        self.target = target

    def __eq__(self, other):
        return type(self) == type(other) and self.target == other.target

    def __repr__(self):
        return '{}({!r})'.format(type(self).__name__, self.target)


class NoOp(Action):
    def __call__(self):
        print('<mdnav: no link>')


class BrowserOpen(Action):
    def __call__(self):
        print('<mdnav: open browser tab>')
        webbrowser.open_new_tab(self.target)


class OSOpen(Action):
    def __call__(self):
        if sys.platform.startswith('linux'):
            call(['xdg-open', self.target])

        elif sys.platform.startswith('darwin'):
            call(['open', self.target])

        else:
            os.startfile(self.target)


class VimOpen(Action):
    def __call__(self):
        import vim

        path = parse_path(self.target)

        # TODO: make space handling more robust?
        vim.command('e {}'.format(path.path.replace(' ', '\\ ')))
        if path.line is not None:
            try:
                line = int(path.line)

            except:
                print('invalid line number')
                return

            else:
                vim.current.window.cursor = (line, 0)

        if path.anchor is not None:
            JumpToAnchor(path.anchor)()


class JumpToAnchor(Action):
    heading_pattern = re.compile(r'^#+(?P<title>.*)$')
    attr_list_pattern = re.compile(r'{:\s+#(?P<id>\S+)\s')

    def __call__(self):
        import vim
        line = self.find_anchor(self.target, vim.current.buffer)

        if line is None:
            return

        vim.current.window.cursor = (line + 1, 0)

    @classmethod
    def find_anchor(cls, target, buffer):
        needle = cls.norm_target(target)

        for (idx, line) in enumerate(buffer):
            m = cls.heading_pattern.match(line)
            if (
                m is not None and
                cls.title_to_anchor(m.group('title')) == needle
            ):
                return idx

            m = cls.attr_list_pattern.search(line)
            if m is not None and needle == m.group('id'):
                return idx

    @staticmethod
    def title_to_anchor(title):
        return '-'.join(fragment.lower() for fragment in title.split())

    @staticmethod
    def norm_target(target):
        if target.startswith('#'):
            target = target[1:]

        return target.lower()


def call(args):
    """If available use vims shell mechanism to work around display issues
    """
    try:
        import vim

    except ImportError:
        subprocess.call(args)

    else:
        args = ['shellescape(' + json.dumps(arg) + ')' for arg in args]
        vim.command('execute "! " . ' + ' . " " . '.join(args))


def parse_path(path):
    """Parse a path with optional line number of anchor into its parts.

    For example::

        parse_path('foo.md:30') == ParsedPath('foo.md', line=30)
        parse_path('foo.md#anchor') == ParsedPath('foo.md', anchor='anchor')

    """
    path, ext = os.path.splitext(path)
    if '#' in ext:
        ext, anchor = ext.rsplit('#', 1)
        return ParsedPath(path + ext, anchor=anchor)

    if ':' in ext:
        ext, line = ext.rsplit(':', 1)
        return ParsedPath(path + ext, line=line)

    return ParsedPath(path + ext)


class ParsedPath(object):
    def __init__(self, path, line=None, anchor=None):
        self.path = path
        self.line = line
        self.anchor = anchor

    def __repr__(self):
        return 'ParsedPath({!r}, line={}, anchor={!r})'.format(self.path, self.line, self.anchor)


def parse_link(cursor, lines):
    row, column = cursor
    line = lines[row]

    _logger.info('handle line %s (%s, %s)', line, row, column)
    m = reference_definition_pattern.match(line)
    if m is not None:
        return m.group('link').strip()

    link_text, rel_column = select_from_start_of_link(line, column)

    if not link_text:
        _logger.info('could not find link text')
        return None

    m = link_pattern.match(link_text)

    if not m:
        _logger.info('does not match link pattern')
        return None

    if m.end('link') <= rel_column:
        _logger.info('cursor outside link')
        return None

    _logger.info('found match: %s', m.groups())
    assert (m.group('direct') is None) != (m.group('indirect') is None)

    if m.group('direct') is not None:
        _logger.info('found direct link: %s', m.group('direct'))
        return m.group('direct')

    _logger.info('follow indirect link %s', m.group('indirect'))
    indirect_ref = m.group('indirect')
    if not indirect_ref:
        indirect_ref = m.group('text')

    indrect_link_pattern = re.compile(
        r'^\[' + re.escape(indirect_ref) + r'\]:(.*)$'
    )

    for line in lines:
        m = indrect_link_pattern.match(line)

        if m:
            return m.group(1).strip()

    _logger.info('could not match for indirect link')
    return None


reference_definition_pattern = re.compile(r'''
    ^
        \[[^\]]*\]:             # reference def at start of line
        (?P<link>.*)            # interpret everything else as link text
    $
''', re.VERBOSE)

link_pattern = re.compile(r'''
    ^
    (?P<link>
        \[                      # start of link text
            (?P<text>[^\]]*)    # link text
        \]                      # end of link text
        (?:
            \(                  # start of target
                (?P<direct>
                    [^\)]*
                )
            \)                  # collect
            |
            \[
                (?P<indirect>
                    [^\]]*
                )
            \]
        )
    )
    .*                  # any non matching characters
    $
''', re.VERBOSE)


def select_from_start_of_link(line, pos):
    """Return the start of the link string and the new cursor
    """
    if pos < len(line) and line[pos] == '[':
        start = pos

    else:
        start = line[:pos].rfind('[')

    # TODO: handle escapes

    if start < 0:
        return None, pos

    # check for indirect links
    if start != 0 and line[start - 1] == ']':
        alt_start = line[:start].rfind('[')
        if alt_start >= 0:
            start = alt_start

    return line[start:], pos - start


if __name__ == "__main__":
    plugin_entry_point()


repository: CKCat/IDAPlugInContest
file: 2016/PlugIn/hexrays_hlight.py
description: IDAPlugIn Contest 备份
wanted_name: HexLight
comment: highlights the matching brace in Pseudocode-View
created_at: 2021-09-23 15:11:52+00:00
pushed_at: 2024-01-29 08:05:30+00:00
url:
https://github.com/CKCat/IDAPlugInContest/blob/7ee5e83dee340fdbce81e86d310287f0a34d7250/2016/PlugIn/hexrays_hlight.py

content:
# highlighting plugin for Hex-Rays Decompiler
# Copyright (c) 2016
# Milan Bohacek <milan.bohacek+hexlight@gmail.com>
# All rights reserved.
# 
# ==============================================================================
# 
# This file is part of Hexlight.
# 
# Hexlight is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
# 
# ==============================================================================


import idautils
import idaapi
import idc


import traceback

hexlight_cb_info = None
hexlight_cb = None

posledni = 0

def jump(custom_viewer, line):
    (pl, x, y)=idaapi.get_custom_viewer_place(custom_viewer, False)
    pl2 = idaapi.place_t_as_simpleline_place_t(pl.clone())
    oldline = pl2.n
    pl2.n = line
    idaapi.jumpto(custom_viewer, pl2, x, y)
    return oldline

class hexrays_callback_info(object):
    
    def __init__(self):
        self.vu = None
        self.highlights = {}
        self.hicolor = 0xF2E8BF #0x00ffff00
        self.theotherline = None
        self.safe = False
        return
    
    def clearall(self, ps, refresh=True):
        ctr = 0
        for i in self.highlights:
            try:            
                ps[i].bgcolor = self.highlights[i]
                ctr += 1
            except:
                pass
#               place_t_as_simpleline_place_t
#              ps[i].line.replace("\x04{\x04", "{")
#              ps[i].line.replace("\x04}\x04", "}")
            
        self.highlights = {}
        self.theotherline = None
        if((ctr > 0) and refresh):
            idaapi.refresh_idaview_anyway()



    def event_callback(self, event, *args):
        try:
#            print "event: %d"%event
            if event == idaapi.hxe_keyboard:
                vu, keycode, shift = args

                if idaapi.lookup_key_code(keycode, shift, True) == idaapi.get_key_code("B") and shift == 0:
                    if self.theotherline:
                        self.theotherline = jump(vu.ct, self.theotherline)
                    return 0


            if event <= idaapi.hxe_print_func:
                self.safe = False

            if event == idaapi.hxe_switch_pseudocode:
                self.safe = False

            if event == idaapi.hxe_func_printed:
                self.safe = True

            if event == idaapi.hxe_text_ready:
                self.safe = True

            if event == idaapi.hxe_curpos:
                if not self.safe:
                    return 0
                #print "1"
                self.vu = args[0]

                if not self.vu:
                    return 0
                #print "2"

                if self.vu.cfunc.maturity != idaapi.CMAT_FINAL:
                    return 0
                #print "3"

                if not self.vu.visible():
                    return 0
                #print "4"
                if not self.vu.refresh_cpos(idaapi.USE_KEYBOARD):
                 #   print "refresh_cpos failed"
                    return 0
                pos = self.vu.cpos
                ypos = pos.lnnum
                #print "cursor click %d %d %d" % (pos.x, pos.y, pos.lnnum)

                if self.highlights.has_key(ypos):
                    return 0
                #print "5"

                ps = self.vu.cfunc.get_pseudocode()
                #print "6"
                #print "ypos:%d"%ypos
                #print "ps[ypos].line: %s"%(ps[ypos].line)

                #line = [idaapi.COLSTR("[%02d]"%i, chr(i)) for i in range(1,0x40) ]
                #ps[0].line = ''.join(line);
                #ps[1].line = '\x04'.join(line);
                #line = [idaapi.COLSTR(  idaapi.COLSTR("[  \x04%02d\x04  ]"%i, chr(i)), chr(i+1)) for i in range(1,0x40) ]
                #ps[2].line = ''.join(line);
                #ps[3].line = '\x04'.join(line);
                ln = ps[ypos].line[:]
                #print "ln"
                curline = idaapi.tag_remove(ln)
                #print "7"

                idxO = curline.find('{')
                idxC = curline.find('}')
                #print "O:", idxO, " C: ",idxC
                #there is no need to highlight first and last {
                #print "8"

                if (idxO >= 0) or (idxC >= 0):
                #   print "9"
                    self.clearall(ps, False)

                    self.highlights[ypos] = ps[ypos].bgcolor

                    ps[ypos].bgcolor = self.hicolor
                    
                    dir = 1
                    bracechar = '}'
                    idx = idxO

                    if (idxC >= 0):
                        dir = -1
                        bracechar = '{'
                        idx = idxC

                    j = ypos + dir

                    max = len(ps)
                 #   print "max: ",max

                    while (j>=0) and (j<max):
                #       print "10"
                        #print "j:", j
                        ln = idaapi.tag_remove(ps[j].line)
                        if ln.find(bracechar) == idx:
                            if not(self.highlights.has_key(j) ):
                                self.highlights[j] = ps[j].bgcolor
                            #ps[j].line = ps[j].line.replace(bracechar, idaapi.COLSTR("\x04"+bracechar+"\x04", "\x27"))
                            #ps[j].line = ps[j].line.replace(bracechar, idaapi.COLSTR(bracechar, chr(52)))
                            ps[j].bgcolor = self.hicolor
                            self.theotherline = j
                            break
                        j+=dir
                    
                    idaapi.refresh_idaview_anyway()
                else:
                    self.clearall(ps)
                #print "11"
                return 0
        except:
            traceback.print_exc()
        
        return 0

def remove():
    if hexlight_cb:
        idaapi.remove_hexrays_callback(hexlight_cb)

class HexHLightPlugin_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_HIDE
    comment = "highlights the matching brace in Pseudocode-View"
    help = "press B to jump to the matching brace"
    wanted_name = "HexLight"
    wanted_hotkey = ""

    def init(self):
        # Some initialization
        global hexlight_cb_info, hexlight_cb

        if idaapi.init_hexrays_plugin():
            hexlight_cb_info = hexrays_callback_info()
            hexlight_cb = hexlight_cb_info.event_callback
            if not idaapi.install_hexrays_callback(hexlight_cb):
            #    print "could not install hexrays_callback"
                return idaapi.PLUGIN_SKIP
            print "Hexlight plugin installed"
            addon = idaapi.addon_info_t();
            addon.id = "milan.bohacek.hexlight";
            addon.name = "Hexlight";
            addon.producer = "Milan Bohacek";
            addon.url = "milan.bohacek+hexlight@gmail.com";
            addon.version = "6.95";
            idaapi.register_addon( addon );
            return idaapi.PLUGIN_KEEP
        #print "init_hexrays_plugin failed"
        return idaapi.PLUGIN_SKIP

    def run(self, arg=0):
        return

    def term(self):
        remove()

def PLUGIN_ENTRY():
    return HexHLightPlugin_t()


repository: rmmenosfbarra/arsenal
file: Reversing&Exploiting/IDA/Plugins/IdaSploitter/idasploiter.py
wanted_name: IDA Sploiter
comment: Exploit development and vulnerability research toolkit.
created_at: 2017-02-27 17:52:25+00:00
pushed_at: 2019-06-30 11:14:50+00:00
url:
https://github.com/rmmenosfbarra/arsenal/blob/67060284c34e411646beacb0b304338a1eabc18b/Reversing&Exploiting/IDA/Plugins/IdaSploitter/idasploiter.py

content:
#!/usr/bin/env python
#
# IDA Sploiter is an exploit development and vulnerability research environment
# implemented as a plugin for Hex-Ray's IDA Pro disassembler.

IDASPLOITER_VERSION = "1.0"

# Copyright (C) 2014 Peter Kacherginsky
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met: 
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer. 
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
# 3. Neither the name of the copyright holder nor the names of its contributors
#    may be used to endorse or promote products derived from this software without 
#    specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

# Performance profiling
import cProfile
import pstats

# IDA libraries
import idaapi
import idautils
import idc
from idaapi import Form, Choose2, plugin_t

# Python libraries
import os
import binascii
import string
import textwrap
import copy
import csv
import itertools

from struct import pack, unpack
from ctypes import *

###############################################################################
# Data Tables and Structures
###############################################################################

###############################################################################
# 00-FF Single byte to unicode transforms
# NOTE: Some unicode characters can have two bytes
#       http://www.phenoelit.org/stuff/Phenoelit20c3.pdf

# Ascii_to_Unicode_transforms
ASCII_TO_UNICODE_GENERAL = [
    "\x00\x00", "\x00\x01", "\x00\x02", "\x00\x03", "\x00\x04", "\x00\x05",
    "\x00\x06", "\x00\x07", "\x00\x08", "\x00\x09", "\x00\x0a", "\x00\x0b", 
    "\x00\x0c", "\x00\x0d", "\x00\x0e", "\x00\x0f", "\x00\x10", "\x00\x11", 
    "\x00\x12", "\x00\x13", "\x00\x14", "\x00\x15", "\x00\x16", "\x00\x17", 
    "\x00\x18", "\x00\x19", "\x00\x1a", "\x00\x1b", "\x00\x1c", "\x00\x1d", 
    "\x00\x1e", "\x00\x1f", "\x00\x20", "\x00\x21", "\x00\x22", "\x00\x23", 
    "\x00\x24", "\x00\x25", "\x00\x26", "\x00\x27", "\x00\x28", "\x00\x29", 
    "\x00\x2a", "\x00\x2b", "\x00\x2c", "\x00\x2d", "\x00\x2e", "\x00\x2f", 
    "\x00\x30", "\x00\x31", "\x00\x32", "\x00\x33", "\x00\x34", "\x00\x35", 
    "\x00\x36", "\x00\x37", "\x00\x38", "\x00\x39", "\x00\x3a", "\x00\x3b", 
    "\x00\x3c", "\x00\x3d", "\x00\x3e", "\x00\x3f", "\x00\x40", "\x00\x41", 
    "\x00\x42", "\x00\x43", "\x00\x44", "\x00\x45", "\x00\x46", "\x00\x47", 
    "\x00\x48", "\x00\x49", "\x00\x4a", "\x00\x4b", "\x00\x4c", "\x00\x4d", 
    "\x00\x4e", "\x00\x4f", "\x00\x50", "\x00\x51", "\x00\x52", "\x00\x53", 
    "\x00\x54", "\x00\x55", "\x00\x56", "\x00\x57", "\x00\x58", "\x00\x59", 
    "\x00\x5a", "\x00\x5b", "\x00\x5c", "\x00\x5d", "\x00\x5e", "\x00\x5f", 
    "\x00\x60", "\x00\x61", "\x00\x62", "\x00\x63", "\x00\x64", "\x00\x65", 
    "\x00\x66", "\x00\x67", "\x00\x68", "\x00\x69", "\x00\x6a", "\x00\x6b", 
    "\x00\x6c", "\x00\x6d", "\x00\x6e", "\x00\x6f", "\x00\x70", "\x00\x71", 
    "\x00\x72", "\x00\x73", "\x00\x74", "\x00\x75", "\x00\x76", "\x00\x77", 
    "\x00\x78", "\x00\x79", "\x00\x7a", "\x00\x7b", "\x00\x7c", "\x00\x7d", 
    "\x00\x7e", "\x00\x7f"
]

ASCII_TO_UNICODE_ANSI = ASCII_TO_UNICODE_GENERAL + [
    "\x20\xac", "\x00\x81", "\x20\x1a", "\x01\x92", "\x20\x1e", "\x20\x26",
    "\x20\x20", "\x20\x21", "\x02\xc6", "\x20\x30", "\x01\x60", "\x20\x39",
    "\x01\x52", "\x00\x8d", "\x01\x7d", "\x00\x8f", "\x90\x00", "\x20\x18",
    "\x20\x19", "\x20\x1c", "\x20\x1d", "\x20\x22", "\x20\x13", "\x20\x14",
    "\x02\xdc", "\x21\x22", "\x01\x61", "\x3a\x20", "\x01\x53", "\x00\x9d",
    "\x01\x7e", "\x01\x78", "\x00\xa0", "\x00\xa1", "\x00\xa2", "\x00\xa3",
    "\x00\xa4", "\x00\xa5", "\x00\xa6", "\x00\xa7", "\x00\xa8", "\x00\xa9",
    "\x00\xaa", "\x00\xab", "\x00\xac", "\x00\xad", "\x00\xae", "\x00\xaf",
    "\x00\xb0", "\x00\xb1", "\x00\xb2", "\x00\xb3", "\x00\xb4", "\x00\xb5",
    "\x00\xb6", "\x00\xb7", "\x00\xb8", "\x00\xb9", "\x00\xba", "\x00\xbb",
    "\x00\xbc", "\x00\xbd", "\x00\xbe", "\x00\xbf", "\x00\xc0", "\x00\xc1",
    "\x00\xc2", "\x00\xc3", "\x00\xc4", "\x00\xc5", "\x00\xc6", "\x00\xc7",
    "\x00\xc8", "\x00\xc9", "\x00\xca", "\x00\xcb", "\x00\xcc", "\x00\xcd",
    "\x00\xce", "\x00\xcf", "\x00\xd0", "\x00\xd1", "\x00\xd2", "\x00\xd3",
    "\x00\xd4", "\x00\xd5", "\x00\xd6", "\x00\xd7", "\x00\xd8", "\x00\xd9",
    "\x00\xda", "\x00\xdb", "\x00\xdc", "\x00\xdd", "\x00\xde", "\x00\xdf",
    "\x00\xe0", "\x00\xe1", "\x00\xe2", "\x00\xe3", "\x00\xe4", "\x00\xe5",
    "\x00\xe6", "\x00\xe7", "\x00\xe8", "\x00\xe9", "\x00\xea", "\x00\xeb",
    "\x00\xec", "\x00\xed", "\x00\xee", "\x00\xef", "\x00\xf0", "\x00\xf1",
    "\x00\xf2", "\x00\xf3", "\x00\xf4", "\x00\xf5", "\x00\xf6", "\x00\xf7",
    "\x00\xf8", "\x00\xf9", "\x00\xfa", "\x00\xfb", "\x00\xfc", "\x00\xfd",
    "\x00\xfe", "\x00\xff"
]

ASCII_TO_UNICODE_OEM = ASCII_TO_UNICODE_GENERAL + [
    "\x00\xc7", "\x00\xfc", "\x00\xe9", "\x00\xe2", "\x00\xe4", "\x00\xe0",
    "\x00\xe5", "\x00\xe7", "\x00\xea", "\x00\xeb", "\x00\xe8", "\x00\xef",
    "\x00\xee", "\x00\xec", "\x00\xc4", "\x00\xc5", "\x00\xc9", "\x00\xe6",
    "\x00\xc6", "\x00\xf4", "\x00\xf6", "\x00\xf2", "\x00\xfb", "\x00\xf9",
    "\x00\xff", "\x00\xd6", "\x00\xdc", "\x00\xf8", "\x00\xa3", "\x00\xd8",
    "\x00\xd7", "\x01\x92", "\x00\xe1", "\x00\xed", "\x00\xf3", "\x00\xfa",
    "\x00\xf1", "\x00\xd1", "\x00\xaa", "\x00\xba", "\x00\xbf", "\x00\xae",
    "\x00\xac", "\x00\xbd", "\x00\xbc", "\x00\xa1", "\x00\xab", "\x00\xbb",
    "\x25\x91", "\x25\x92", "\x25\x93", "\x25\x02", "\x25\x24", "\x00\xc1",
    "\x00\xc2", "\x00\xc0", "\x00\xa9", "\x25\x63", "\x25\x51", "\x25\x57",
    "\x25\x5d", "\x00\xa2", "\x00\xa5", "\x25\x10", "\x25\x14", "\x25\x34",
    "\x25\x2c", "\x25\x1c", "\x25\x00", "\x25\x3c", "\x00\xe3", "\x00\xc3",
    "\x25\x5a", "\x25\x54", "\x25\x69", "\x25\x66", "\x25\x60", "\x25\x50",
    "\x25\x6c", "\x00\xa4", "\x00\xf0", "\x00\xd0", "\x00\xca", "\x00\xcb",
    "\x00\xc8", "\x01\x31", "\x00\xcd", "\x00\xce", "\x00\xcf", "\x25\x18",
    "\x25\x0c", "\x25\x88", "\x25\x84", "\x00\xa6", "\x00\xcc", "\x25\x80",
    "\x00\xd3", "\x00\xdf", "\x00\xd4", "\x00\xd2", "\x00\xf5", "\x00\xd5",
    "\x00\xb5", "\x00\xfe", "\x00\xde", "\x00\xda", "\x00\xdb", "\x00\xd9",
    "\x00\xfd", "\x00\xdd", "\x00\xaf", "\x00\xb4", "\x00\xad", "\x00\xb1",
    "\x20\x17", "\x00\xbe", "\x00\xb6", "\x00\xa7", "\x00\xf7", "\x00\xb8",
    "\x00\xb0", "\x00\xa8", "\x00\xb7", "\x00\xb9", "\x00\xb3", "\x00\xb2",
    "\x25\xa0", "\x00\xa0"
]

ASCII_TO_UNICODE_UTF7 = ASCII_TO_UNICODE_GENERAL + [
    "\xff\x80", "\xff\x81", "\xff\x82", "\xff\x83", "\xff\x84", "\xff\x85",
    "\xff\x86", "\xff\x87", "\xff\x88", "\xff\x89", "\xff\x8a", "\xff\x8b",
    "\xff\x8c", "\xff\x8d", "\xff\x8e", "\xff\x8f", "\xff\x90", "\xff\x91",
    "\xff\x92", "\xff\x93", "\xff\x94", "\xff\x95", "\xff\x96", "\xff\x97",
    "\xff\x98", "\xff\x99", "\xff\x9a", "\xff\x9b", "\xff\x9c", "\xff\x9d",
    "\xff\x9e", "\xff\x9f", "\xff\xa0", "\xff\xa1", "\xff\xa2", "\xff\xa3",
    "\xff\xa4", "\xff\xa5", "\xff\xa6", "\xff\xa7", "\xff\xa8", "\xff\xa9",
    "\xff\xaa", "\xff\xab", "\xff\xac", "\xff\xad", "\xff\xae", "\xff\xaf",
    "\xff\xb0", "\xff\xb1", "\xff\xb2", "\xff\xb3", "\xff\xb4", "\xff\xb5",
    "\xff\xb6", "\xff\xb7", "\xff\xb8", "\xff\xb9", "\xff\xba", "\xff\xbb",
    "\xff\xbc", "\xff\xbd", "\xff\xbe", "\xff\xbf", "\xff\xc0", "\xff\xc1",
    "\xff\xc2", "\xff\xc3", "\xff\xc4", "\xff\xc5", "\xff\xc6", "\xff\xc7",
    "\xff\xc8", "\xff\xc9", "\xff\xca", "\xff\xcb", "\xff\xcc", "\xff\xcd",
    "\xff\xce", "\xff\xcf", "\xff\xd0", "\xff\xd1", "\xff\xd2", "\xff\xd3",
    "\xff\xd4", "\xff\xd5", "\xff\xd6", "\xff\xd7", "\xff\xd8", "\xff\xd9",
    "\xff\xda", "\xff\xdb", "\xff\xdc", "\xff\xdd", "\xff\xde", "\xff\xdf",
    "\xff\xe0", "\xff\xe1", "\xff\xe2", "\xff\xe3", "\xff\xe4", "\xff\xe5",
    "\xff\xe6", "\xff\xe7", "\xff\xe8", "\xff\xe9", "\xff\xea", "\xff\xeb",
    "\xff\xec", "\xff\xed", "\xff\xee", "\xff\xef", "\xff\xf0", "\xff\xf1",
    "\xff\xf2", "\xff\xf3", "\xff\xf4", "\xff\xf5", "\xff\xf6", "\xff\xf7",
    "\xff\xf8", "\xff\xf9", "\xff\xfa", "\xff\xfb", "\xff\xfc", "\xff\xfd",
    "\xff\xfe", "\xff\xff"
]

ASCII_TO_UNICODE_UTF8 = ASCII_TO_UNICODE_GENERAL + [
    "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00",
    "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00",
    "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00",
    "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00",
    "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00",
    "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00",
    "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00",
    "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00",
    "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00",
    "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00",
    "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x00", "\x00\x01",
    "\x00\x02", "\x00\x03", "\x00\x04", "\x00\x05", "\x00\x06", "\x00\x07",
    "\x00\x08", "\x00\x09", "\x00\x0a", "\x00\x0b", "\x00\x0c", "\x00\x0d",
    "\x00\x0e", "\x00\x0f", "\x00\x10", "\x00\x11", "\x00\x12", "\x00\x13",
    "\x00\x14", "\x00\x15", "\x00\x16", "\x00\x17", "\x00\x18", "\x00\x19",
    "\x00\x1a", "\x00\x1b", "\x00\x1c", "\x00\x1d", "\x00\x1e", "\x00\x1f",
    "\x00\x00", "\x00\x01", "\x00\x02", "\x00\x03", "\x00\x04", "\x00\x05",
    "\x00\x06", "\x00\x07", "\x00\x08", "\x00\x09", "\x00\x0a", "\x00\x0b",
    "\x00\x0c", "\x00\x0d", "\x00\x0e", "\x00\x0f", "\x00\x00", "\x00\x01",
    "\x00\x02", "\x00\x03", "\x00\x04", "\x00\x05", "\x00\x06", "\x00\x07",
    "\x00\x08", "\x00\x09", "\x00\x0a", "\x00\x0b", "\x00\x0c", "\x00\x0d",
    "\x00\x0e", "\x00\x0f"
]

ASCII_TO_UNICODE = [ASCII_TO_UNICODE_ANSI,
                    ASCII_TO_UNICODE_OEM, 
                    ASCII_TO_UNICODE_UTF7, 
                    ASCII_TO_UNICODE_UTF8]

###############################################################################
# Microsoft Portable Executable and Common Object File Format Specification
# Revision 8.3 - February 6, 2013

class IMAGE_DOS_HEADER(Structure):
    _fields_ = [
                    ("signature" , c_char * 2),
                    ("lastsize"  , c_short),
                    ("nblocks"   , c_short),
                    ("nreloc"    , c_short),
                    ("hdrsize"   , c_short),
                    ("minalloc"  , c_short),
                    ("maxalloc"  , c_short),
                    ("ss"        , c_short),
                    ("sp"        , c_short),
                    ("checksum"  , c_short),
                    ("ip"        , c_short),
                    ("cs"        , c_short),
                    ("relocpos"  , c_short),
                    ("noverlay"  , c_short),
                    ("reserved1" , c_short * 4),
                    ("oem_id"    , c_short),
                    ("reserved2" , c_short * 10),
                    ("e_lfanew"  , c_long)
                ]

class IMAGE_FILE_HEADER(Structure):
    _fields_ = [
                    ("Machine"              , c_short),
                    ("NumberOfSections"     , c_short),
                    ("TimeDateStamp"        , c_long),
                    ("PointerToSymbolTable" , c_long),
                    ("NumberOfSymbols"      , c_long),
                    ("SizeOfOptionalHeader" , c_short),
                    ("Characteristics"      , c_short)
                ]

class IMAGE_DATA_DIRECTORY(Structure):
    _fields_ = [
                    ("VirtualAddress", c_long),
                    ("Size", c_long)
                ]

class IMAGE_OPTIONAL_HEADER(Structure):
    _fields_ = [
                    ("signature"               , c_short),
                    ("MajorLinkerVersion"      , c_byte),
                    ("MinorLinkerVersion"      , c_byte),
                    ("SizeOfCode"              , c_long),
                    ("SizeOfInitializedData"   , c_long),
                    ("SizeOfUninitializedData" , c_long),
                    ("AddressOfEntryPoint"     , c_long),
                    ("BaseOfCode"              , c_long),
                    ("BaseOfData"              , c_long),
                    ("ImageBase"               , c_long),
                    ("SectionAlignment"        , c_long),
                    ("FileAlignment"           , c_long),
                    ("MajorOSVersion"          , c_short),
                    ("MinorOSVersion"          , c_short),
                    ("MajorImageVersion"       , c_short),
                    ("MinorImageVersion"       , c_short),
                    ("MajorSubsystemVersion"   , c_short),
                    ("MinorSubsystemVersion"   , c_short),
                    ("Reserved"                , c_long),
                    ("SizeOfImage"             , c_long),
                    ("SizeOfHeaders"           , c_long),
                    ("Checksum"                , c_long),
                    ("Subsystem"               , c_short),
                    ("DLLCharacteristics"      , c_short),
                    ("SizeOfStackReserve"      , c_long),
                    ("SizeOfStackCommit"       , c_long),
                    ("SizeOfHeapReserve"       , c_long),
                    ("SizeOfHeapCommit"        , c_long),
                    ("LoaderFlags"             , c_long),
                    ("NumberOfRvaAndSizes"     , c_long),
                    ("DataDirectory"           , IMAGE_DATA_DIRECTORY * 16)
                ]

class IMAGE_OPTIONAL_HEADER64(Structure):
    _fields_ = [
                    ("signature"               , c_short),
                    ("MajorLinkerVersion"      , c_byte),
                    ("MinorLinkerVersion"      , c_byte),
                    ("SizeOfCode"              , c_long),
                    ("SizeOfInitializedData"   , c_long),
                    ("SizeOfUninitializedData" , c_long),
                    ("AddressOfEntryPoint"     , c_long),
                    ("BaseOfCode"              , c_long),
                    ("ImageBase"               , c_longlong),
                    ("SectionAlignment"        , c_long),
                    ("FileAlignment"           , c_long),
                    ("MajorOSVersion"          , c_short),
                    ("MinorOSVersion"          , c_short),
                    ("MajorImageVersion"       , c_short),
                    ("MinorImageVersion"       , c_short),
                    ("MajorSubsystemVersion"   , c_short),
                    ("MinorSubsystemVersion"   , c_short),
                    ("Reserved"                , c_long),
                    ("SizeOfImage"             , c_long),
                    ("SizeOfHeaders"           , c_long),
                    ("Checksum"                , c_long),
                    ("Subsystem"               , c_short),
                    ("DLLCharacteristics"      , c_short),
                    ("SizeOfStackReserve"      , c_longlong),
                    ("SizeOfStackCommit"       , c_longlong),
                    ("SizeOfHeapReserve"       , c_longlong),
                    ("SizeOfHeapCommit"        , c_longlong),
                    ("LoaderFlags"             , c_long),
                    ("NumberOfRvaAndSizes"     , c_long),
                    ("DataDirectory"           , IMAGE_DATA_DIRECTORY * 16)
                ]

class IMAGE_LOAD_CONFIG_DIRECTORY(Structure):
    _fields_ = [
                    ("Size"                          , c_long),
                    ("TimeDateStamp"                 , c_long),
                    ("MajorVersion"                  , c_short),
                    ("MinorVersion"                  , c_short),
                    ("GlobalFlagsClear"              , c_long),
                    ("GlobalFlagsSet"                , c_long),
                    ("CriticalSectionDefaultTimeout" , c_long),
                    ("DeCommitFreeBlockThreshold"    , c_long),
                    ("DeCommitTotalFreeThreshold"    , c_long),
                    ("LockPrefixTable"               , c_long),
                    ("MaximumAllocationSize"         , c_long),
                    ("VirtualMemoryThreshold"        , c_long),
                    ("ProcessHeapFlags"              , c_long),
                    ("ProcessAffinityMask"           , c_long),
                    ("CSDVersion"                    , c_short),
                    ("Reserved1"                     , c_short),
                    ("EditList"                      , c_long),
                    ("SecurityCookie"                , c_long),
                    ("SEHandlerTable"                , c_long),
                    ("SEHandlerCount"                , c_long)
                ]

class IMAGE_LOAD_CONFIG_DIRECTORY64(Structure):
    _fields_ = [
                    ("Size"                          , c_long),
                    ("TimeDateStamp"                 , c_long),
                    ("MajorVersion"                  , c_short),
                    ("MinorVersion"                  , c_short),
                    ("GlobalFlagsClear"              , c_long),
                    ("GlobalFlagsSet"                , c_long),
                    ("CriticalSectionDefaultTimeout" , c_long),
                    ("DeCommitFreeBlockThreshold"    , c_longlong),
                    ("DeCommitTotalFreeThreshold"    , c_longlong),
                    ("LockPrefixTable"               , c_longlong),
                    ("MaximumAllocationSize"         , c_longlong),
                    ("VirtualMemoryThreshold"        , c_longlong),
                    ("ProcessHeapFlags"              , c_long),
                    ("ProcessAffinityMask"           , c_long),
                    ("CSDVersion"                    , c_short),
                    ("Reserved1"                     , c_short),
                    ("EditList"                      , c_longlong),
                    ("SecurityCookie"                , c_longlong),
                    ("SEHandlerTable"                , c_longlong),
                    ("SEHandlerCount"                , c_longlong),
                ]

###############################################################################
# Module Class - Manages module characteristics

class Module():

    def __init__(self, name, size, base, rebase_to):

            self.addr = base

            # BUG: IDA's API does not always zero out the SWIG buffer.
            # BUG: ntdll does not have path information.
            if "\x00" in name:
                name,junk = name.split("\x00",1)

            self.file = os.path.basename(name)
            self.path = os.path.dirname(name)

            self.size = size

            # Parse the module as a PE file
            pe = PE(self.addr)

            if pe:
                self.NXCompat = pe.isNXCompat()
                self.ASLR     = pe.isDynamicBase()
                self.SafeSEH  = pe.isSafeSEH()
                self.GS       = pe.isGS()

###############################################################################
# PE Module Parser

class PE():

    # PE flags
    IMAGE_FILE_MACHINE_I386 = 0x14c

    MAGIC_PE32      = 0x10b
    MAGIC_PE32_PLUS = 0x20b

    IMAGE_DLL_CHARACTERISTICS_DYNAMIC_BASE = 0x0040
    IMAGE_DLL_CHARACTERISTICS_NX_COMPAT    = 0x0100
    IMAGE_DLLCHARACTERISTICS_NO_SEH        = 0x0400

    def __init__(self, base):

        # PE headers
        self.dos_header            = None
        self.file_header           = None
        self.optional_header       = None
        self.load_config_directory = None
        self.arch64                = False

        # Verify DOS header magic
        magic_dos = idaapi.dbg_read_memory( base, 0x2)
        if magic_dos != "MZ":
            print "[idasploiter] Invalid DOS header magic."
            return None

        # Parse the DOS header
        self.dos_header = IMAGE_DOS_HEADER.from_buffer_copy(
            idaapi.dbg_read_memory(base, 0x40) )

        # Get offset to PE header
        base_pe = base + self.dos_header.e_lfanew

        # Verify PE header magic
        magic_pe = idaapi.dbg_read_memory( base_pe, 0x4 )
        if magic_pe != "PE\x00\x00":
            print "[idasploiter] Invalid PE header magic: %x" % unpack("I",magic_pe)[0]
            return None

        # Parse the FILE header
        # NOTE: IMAGE_FILE_HEADER size is 0x14
        self.file_header = IMAGE_FILE_HEADER.from_buffer_copy( 
            idaapi.dbg_read_memory( base_pe + 0x4, sizeof(IMAGE_FILE_HEADER) ) )

        # Parse the OPTIONAL header
        base_pe_optional =  base_pe + 0x4 + 0x14
        magic_pe_opt = idaapi.dbg_read_memory ( base_pe_optional, 0x2)
        magic_pe_opt = unpack("H", magic_pe_opt)[0]

        # PE32 Optional Header
        if magic_pe_opt == self.MAGIC_PE32:
            self.arch64 = False
            self.optional_header = IMAGE_OPTIONAL_HEADER.from_buffer_copy(
                idaapi.dbg_read_memory( base_pe_optional, self.file_header.SizeOfOptionalHeader ))

        # PE32+ Optional Header
        elif magic_pe_opt == self.MAGIC_PE32_PLUS:
            self.arch64 = True
            self.optional_header = IMAGE_OPTIONAL_HEADER64.from_buffer_copy(
                idaapi.dbg_read_memory( base_pe_optional, self.file_header.SizeOfOptionalHeader ))

        # Invalid PE Header
        else:
            print "[idasploiter] Invalid IMAGE_OPTIONAL_HEADER magic: %x" % unpack("H",magic_pe_opt)[0]
            return None

        # Load Configuration Table
        load_config_directory_rva = self.optional_header.DataDirectory[10].VirtualAddress

        # Parse Load Configuration Table if present
        if load_config_directory_rva:

            # Read LOAD CONFIG DIRECTORY size before parsing it
            load_config_directory_size = idaapi.dbg_read_memory( base + load_config_directory_rva, 0x4)
            load_config_directory_size = unpack("I",load_config_directory_size)[0]

            # Parse LOAD CONFIG DIRECTORY based on PE32 format
            if self.optional_header.signature == self.MAGIC_PE32:
                self.load_config_directory = IMAGE_LOAD_CONFIG_DIRECTORY.from_buffer_copy( idaapi.dbg_read_memory( base + load_config_directory_rva, load_config_directory_size ))

            # Parse LOAD CONFIG DIRECTORY based on PE32+ format
            elif self.optional_header.signature == self.MAGIC_PE32_PLUS:
                self.load_config_directory = IMAGE_LOAD_CONFIG_DIRECTORY64.from_buffer_copy( idaapi.dbg_read_memory( base + load_config_directory_rva, load_config_directory_size ))

            # Invalid PE header
            else:
                print "[idasploiter] Invalid IMAGE_OPTIONAL_HEADER magic: %x" % unpack("H",magic)[0]
                return None

    def isDynamicBase(self):
        if self.optional_header:
            if self.optional_header.DLLCharacteristics & self.IMAGE_DLL_CHARACTERISTICS_DYNAMIC_BASE:
                return "Yes"
            else:
                return "No"
        else:
            return "Unk"

    def isNXCompat(self):
        if self.optional_header:
            if self.optional_header.DLLCharacteristics & self.IMAGE_DLL_CHARACTERISTICS_NX_COMPAT:
                return "Yes"
            else:
                return "No"
        else:
            return "Unk"

    def isSafeSEH(self):

        if self.file_header and self.optional_header:

            # SafeSEH is only applicable to 32-bit Windows.
            if self.file_header.Machine != self.IMAGE_FILE_MACHINE_I386:
                return "N/A"

            # NO_SEH flag indicates that the module does not use SEH and no handlers can be called in this image.
            elif self.optional_header.DLLCharacteristics & self.IMAGE_DLLCHARACTERISTICS_NO_SEH:
                return "No SEH"

            # SafeSEH is enabled when there is SEHandlerTable entry in the LOAD_CONFIG_DIRECTORY
            elif self.load_config_directory and self.load_config_directory.SEHandlerTable != 0:
                return "Yes"

            # Otherwise SafeSEH is not used
            else:
                return "No"

        else:
            return "Unk"

    def isGS(self):

        if self.file_header and self.optional_header:

            # NOTE: PE32+ does not necessarily provide a pointer to the Security Cookie in the
            #       LoadConfigDirectory. The address appears to be still present in the binary.

            # TODO: Need to do additional research into reliably detecting /GS in PE32+ files.

            # PE32+ binaries with missing LoadConfigDirectory may still used security cookies.
            if self.arch64 and not self.load_config_directory:
                return "Unk"

            if self.load_config_directory and self.load_config_directory.SecurityCookie != 0:

                # Attempt to read the security cookie
                # NOTE: The extra exception handler was added to account for cookie being stored in memory locations
                #       that are not accessible by the debugger (e.g. manually loading win32k.sys, security cookie
                #       address will point to kernel space which we can't read)
                try:
                    # Confirm the Security Cookie is not zero
                    if self.arch64:
                        cookie = idaapi.dbg_read_memory(self.load_config_directory.SecurityCookie, 8)
                        cookie = struct.unpack("<Q", cookie)[0]
                    else:
                        cookie = idaapi.dbg_read_memory(self.load_config_directory.SecurityCookie, 4)
                        cookie = struct.unpack("<I", cookie)[0]

                    if cookie != 0:
                        return "Yes"
                    else:
                        return "Null"

                except Exception, e:
                    return "Inv"

            else:
                return "No"

        else:
            return "Unk"

###############################################################################
# Pointer Class - Manager writable pointer characteristics

class Ptr():

    def __init__(self, module, ptr_ea, ptr_offset, ptr_charset, call_ea, insn_disas):
        self.module = module

        self.ptr_ea = ptr_ea
        self.ptr_offset = ptr_offset
        self.ptr_charset = ptr_charset

        self.call_ea = call_ea
        self.insn_disas = insn_disas

        self.name = idaapi.get_name(self.call_ea + self.ptr_offset, self.ptr_ea) or ""
        self.call_name = idaapi.get_func_name(self.call_ea) or ""

        self.p2p_ea = None
        self.p2p_offset = None
        self.p2p_charset = None

###############################################################################
# Function Pointer Search Engine

class FuncPtr():

    def __init__(self,sploiter):

        self.sploiter   = sploiter
        self.ptr_calls  = list()
        self.ptrs       = list()

    # NOTE: Pointer Offsets are interpreted as follows:
    # Example 1: Consider we have loaded fptr with a user controlled value
    #
    #              mov [fptr + offset], value
    #
    #            The fptr value we need is (fptr - offset) to overwrite
    #            the actual fptr location. Do this early so that applied address
    #            filters consider fptr bytes before offset adjustment.
    #
    #            Specify this positive offset in the Function Pointer form
    #            so that it would be applied to all discovered function pointers.
    #
    # Example 2: Consider we have a pointer to a pointer write condition:
    #
    #              mov fptr, [p2p + offset2]
    #              mov [fptr + offset1], value
    #
    #            The value (p2p - offset2) must point to the fptr location,
    #            but since we have another offset the target of (p2p - offset2)
    #            must be actually (fptr - offset1). Also do this early so that
    #            applied address filters consider p2p before offset adjustment.
    #
    #            Specify this positive p2p offset2 so it gets applied to the listing
    #            of pointers-to-pointers. All pointer-to-pointer(s) will also
    #            be calculated relative to the ptr - offset1
    #
    # NOTE: Only search for pointers to pointers within the same module. This
    #       was done purely for performance considerations since any readable
    #       pointer to pointer may be used.


    def search_pointers(self):
        
        # HACK: A separate flag is used to track user canceling the search,
        #       because multiple calls to idaapi.wasBreak() do not properly
        #       detect cancellations.
        breakFlag = False

        # Show wait dialog
        idaapi.show_wait_box("Searching writable function pointers...")

        for m in self.modules:
 
            ###################################################################
            # Locate all of the CALL and JMP instructions in the current module
            # which use an immediate operand.

            # List of call/jmp pointer calls in a given module
            ptr_calls = list()

            # Iterate over segments in the module
            # BUG: Iterating over all loaded segments is more stable than looking up by address
            for n in xrange(idaapi.get_segm_qty()):
                seg = idaapi.getnseg(n)

                # Segment in a selected modules
                if seg and seg.startEA >= m.addr and seg.endEA <= (m.addr + m.size):

                    # Locate executable segments
                    # NOTE: Each module may have multiple executable segments
                    # TODO: Search for "MOV REG, PTR # CALL REG"
                    if seg.perm & idaapi.SEGPERM_EXEC:

                        # Search all instances of CALL /2 imm32/64 - FF 15
                        # TODO: Alternative pointer calls using SIB: FF 14 E5 11 22 33 44 - call dword/qword ptr [0x44332211]
                        #                                            FF 14 65 11 22 33 44
                        #                                            FF 14 25 11 22 33 44
                        call_ea = seg.startEA
                        while True:
                            call_ea = idaapi.find_binary(call_ea + 1, seg.endEA, "FF 15", 16, idaapi.SEARCH_DOWN)
                            if call_ea == idaapi.BADADDR: break
                            ptr_calls.append( call_ea )

                        # Search all instances of JMP /2 imm32/64 - FF 25
                        # TODO: Alternative pointer calls using SIB: FF 24 E5 11 22 33 44 - jmp dword/qword ptr [0x44332211]
                        #                                            FF 24 65 11 22 33 44
                        #                                            FF 24 25 11 22 33 44
                        call_ea = seg.startEA
                        while True:
                            call_ea = idaapi.find_binary(call_ea + 1, seg.endEA, "FF 25", 16, idaapi.SEARCH_DOWN)
                            if call_ea == idaapi.BADADDR: break
                            ptr_calls.append( call_ea )

            ###################################################################
            # Extract all of the function pointers and make sure they are
            # are writable.

            # List of writable function pointer objects in a given module
            ptrs = list()

            for call_ea in ptr_calls:

                # Decode CALL/JMP instruction
                # NOTE: May result in invalid disassembly of split instructions
                insn_size = idaapi.decode_insn(call_ea)

                if insn_size:

                    insn = idaapi.cmd
                    insn_op1 = insn.Operands[0].type

                    # Verify first operand is a direct memory reference
                    if insn.Operands[0].type == idaapi.o_mem:

                        # Get operand address
                        ptr_ea = insn.Operands[0].addr

                        # Apply pointer offset
                        ptr_ea -= self.ptrOffset

                        # Locate segment where the pointer is located
                        ptr_seg = idaapi.getseg(ptr_ea)

                        # Make sure a valid segment writeable segment was found
                        if ptr_seg and ptr_seg.perm & idaapi.SEGPERM_WRITE:

                            # Get pointer charset
                            ptr_charset = self.sploiter.get_ptr_charset(ptr_ea)

                            # Filter the pointer
                            if not self.filterP2P:
                                if ptr_charset == None:                                    continue
                                if self.ptrNonull     and not "nonull"     in ptr_charset: continue
                                if self.ptrUnicode    and not "unicode"    in ptr_charset: continue
                                if self.ptrAscii      and not "ascii"      in ptr_charset: continue
                                if self.ptrAsciiPrint and not "asciiprint" in ptr_charset: continue
                                if self.ptrAlphaNum   and not "alphanum"   in ptr_charset: continue
                                if self.ptrNum        and not "numeric"    in ptr_charset: continue
                                if self.ptrAlpha      and not "alpha"      in ptr_charset: continue

                            # Increment the fptr counter

                            # Get pointer disassembly
                            insn_disas = idc.GetDisasmEx(call_ea, idaapi.GENDSM_FORCE_CODE)

                            # Add pointer to the list
                            ptr = Ptr(m.file, ptr_ea, self.ptrOffset, ptr_charset, call_ea, insn_disas)
                            ptrs.append(ptr)

            ###################################################################
            # Cache Pointers to Pointers

            ptr_ea_prefix_cache = dict()

            if self.searchP2P:

                # CACHE: Running repeated searches over the entire memory space is 
                #        very expensive. Let's cache all of the addresses containing
                #        bytes corresponding to discovered function pointers in a
                #        single search and simply reference this cache for each
                #        function pointer. Specifically running idaapi.find_binary()
                #        is much more expensive than idaapi.dbg_read_memory().
                #
                # NOTE:  For performance considerations, the cache works on a per 
                #        module basis, but could be expanded for the entire memory
                #        space.
                #
                # prefix_offset - how many bytes of discovered function
                #        pointers to cache.
                #                 
                #        Example: For function pointers 0x00401234, 0x00404321, 0x000405678
                #        we are going to use prefix_offset 2, so we will cache all of the
                #        values located at addresses 0x0040XXXX

                if self.sploiter.addr64:
                    pack_format   = "<Q"
                    addr_bytes    = 8
                    prefix_offset = 6
                else:
                    pack_format   = "<I"
                    addr_bytes    = 4
                    prefix_offset = 2

                # Set of unique N-byte address prefixes to search in memory
                ea_prefix_set = set()

                for ptr in ptrs:
                    ptr_ea = ptr.ptr_ea

                    ptr_bytes = struct.pack( pack_format, ptr_ea )
                    ptr_bytes = ptr_bytes[-prefix_offset:]

                    ea_prefix_set.add( ptr_bytes)

                # Search the module for all bytes corresponding to the prefix
                # and use them as candidates for pointers-to-pointers

                for ea_prefix in ea_prefix_set:

                    # NOTE: Make sure you search using 44 33 22 11 format and not 11223344
                    ea_prefix_str = " ".join(  ["%02x" % ord(b) for b in ea_prefix]  )

                    # Initialize search parameters for a given module
                    ea = m.addr
                    maxea = m.addr + m.size

                    while True:
                        ea = idaapi.find_binary(ea + 1, maxea, ea_prefix_str, 16, idaapi.SEARCH_DOWN)
                        if ea == idaapi.BADADDR: break

                        p2p_ea = ea - (addr_bytes - prefix_offset)

                        dbg_mem = idaapi.dbg_read_memory(p2p_ea, addr_bytes)
                        ptr_ea_prefix = unpack(pack_format, dbg_mem)[0]

                        if ptr_ea_prefix in ptr_ea_prefix_cache:
                            ptr_ea_prefix_cache[ptr_ea_prefix].add(p2p_ea)
                        else:
                            ptr_ea_prefix_cache[ptr_ea_prefix] = set([p2p_ea,])

                        # Detect search cancellation, but allow the loop below
                        # to run to create already cached/found function pointers

                        # Canceled
                        if breakFlag or idaapi.wasBreak():
                            breakFlag = True
                            break

                    # Canceled
                    if breakFlag or idaapi.wasBreak():
                        breakFlag = True
                        break


            ###################################################################
            # Locate Pointer to Pointers

            for ptr in ptrs:

                ptr_ea = ptr.ptr_ea

                # Locate pointers-to-pointers for a given function pointer in the cache
                if self.searchP2P and ptr_ea in ptr_ea_prefix_cache:

                    for p2p_ea in ptr_ea_prefix_cache[ptr_ea]:

                        # Apply pointer-to-pointer offset
                        p2p_ea -= self.p2pOffset

                        p2p_charset = self.sploiter.get_ptr_charset(p2p_ea)

                        # Filter the pointer
                        if self.filterP2P:
                            if p2p_charset == None:                                    continue
                            if self.ptrNonull     and not "nonull"     in p2p_charset: continue
                            if self.ptrUnicode    and not "unicode"    in p2p_charset: continue
                            if self.ptrAscii      and not "ascii"      in p2p_charset: continue
                            if self.ptrAsciiPrint and not "asciiprint" in p2p_charset: continue
                            if self.ptrAlphaNum   and not "alphanum"   in p2p_charset: continue
                            if self.ptrNum        and not "numeric"    in p2p_charset: continue
                            if self.ptrAlpha      and not "alpha"      in p2p_charset: continue

                        # Copy existing pointer object to modify it for the particular p
                        p2p = copy.copy(ptr)
                        p2p.p2p_ea = p2p_ea
                        p2p.p2p_offset = self.p2pOffset
                        p2p.p2p_charset = p2p_charset

                        # Apppend p2p specific pointer object to the global list
                        self.ptrs.append(p2p)

                        # Exceeded maximum number of pointers
                        if self.maxPtrs and len(self.ptrs) >= self.maxPtrs:
                            breakFlag = True
                            print "[idasploiter] Maximum number of pointers exceeded."
                            break

                # Simply append pointer object to the global list
                else:
                    self.ptrs.append(ptr)

                    # Exceeded maximum number of pointers
                    if self.maxPtrs and len(self.ptrs) >= self.maxPtrs:
                        breakFlag = True
                        print "[idasploiter] Maximum number of pointers exceeded."
                        break

                if breakFlag or idaapi.wasBreak():
                    breakFlag = True
                    break

            # Canceled
            # NOTE: Only works when started from GUI not script.
            if breakFlag or idaapi.wasBreak():
                breakFlag = True
                print "[idasploiter] Canceled."
                break

        print "[idasploiter] Found %d total pointers." % len(self.ptrs)
        idaapi.hide_wait_box()

###############################################################################
# Gadget Class - manages ROP gadget characteristics

class Gadget():

    def __init__(self, instructions, pivot, operations, chg_registers, use_registers):

        self.address = 0x0

        self.module = ""

        self.instructions = instructions
        self.size = len(instructions)

        self.pivot = pivot
        self.operations = operations

        self.chg_registers = chg_registers
        self.use_registers = use_registers

        self.ptr_charset = []

###############################################################################
# ROP Search Engine

class Rop():

    def __init__(self, sploiter):
        
        self.maxRopOffset = 40 # Maximum offset from the return instruction to look for gadgets. default: 40
        self.maxRopSize   = 6  # Maximum number of instructions to look for gadgets. default: 6
        self.maxRetnImm   = 64 # Maximum imm16 value in retn. default: 64
        self.maxJopImm    = 255 # Maximum jop [reg + IMM] value. default: 64
        self.maxRops      = 0  # Maximum number of ROP chains to find. default: 0 (unlimited)

        self.debug        = False

        self.regnames     = idaapi.ph_get_regnames()

        self.sploiter     = sploiter
        self.retns        = list()
        self.gadgets      = list()

        # Decoded instruction cache
        self.insn_cache = dict()

        # Extra bytes to read to ensure correct decoding of
        # RETN, RETN imm16, CALL /2, and JMP /4 instructions.
        self.dbg_read_extra = 6 # FF + ModR/M + SIB + disp32 

        self.insn_arithmetic_ops = ["inc","dec","neg", "add","sub","mul","imul","div","idiv","adc","sbb","lea"]
        self.insn_bit_ops = ["not","and","or","xor","shr","shl","sar","sal","shld","shrd","ror","rcr","rcl"]

    def get_o_reg_name(self, insn, n):

        reg_num = insn.Operands[n].reg
        reg_name = self.regnames[reg_num]

        # NOTE: IDA's x86/x86-64 regname array contains only register root names
        #       (e.g ax,cx,dx,etc.). However we can still figure out exact register
        #       size by looking at the operand 'dtyp' property.
        if reg_num < 8:

            # 32-bit register
            if insn.Operands[n].dtyp == idaapi.dt_dword:
                reg_name = 'e'+reg_name

            # 64-bit register
            elif insn.Operands[n].dtyp == idaapi.dt_qword:
                reg_name = 'r'+reg_name

            # 16-bit register otherwise

        return reg_name

    def search_retns(self):

        for m in self.modules:

            # Iterate over segments in the module
            # BUG: Iterating over all loaded segments is more stable than looking up by address
            for n in xrange(idaapi.get_segm_qty()):
                seg = idaapi.getnseg(n)

                # Locate executable segments in a selected modules
                # NOTE: Each module may have multiple executable segments
                if seg and \
                   seg.startEA >= m.addr and seg.endEA <= (m.addr + m.size) and \
                   seg.perm & idaapi.SEGPERM_EXEC:

                    #######################################################
                    # Search for ROP gadgets
                    if self.searchRop:

                        #Search all instances of RETN 
                        ea = seg.startEA
                        while True:
                            ea = idaapi.find_binary(ea + 1, seg.endEA, "C3", 16, idaapi.SEARCH_DOWN)
                            if ea == idaapi.BADADDR: break
                            self.retns.append( (ea, m.file))

                        # Search all instances of RETN imm16
                        ea = seg.startEA
                        while True:
                            ea = idaapi.find_binary(ea + 1, seg.endEA, "C2", 16, idaapi.SEARCH_DOWN)
                            if ea == idaapi.BADADDR: break

                            # Read imm16 value and filter large values
                            retn_imm16 = idaapi.dbg_read_memory ( ea + 1, 0x2)
                            retn_imm16 = unpack("<H", retn_imm16)[0]

                            if retn_imm16 <= self.maxRetnImm:

                                self.retns.append( (ea, m.file))

                    #######################################################
                    # Search for JOP gadgets
                    if self.searchJop:

                        # Search all instances of JMP reg (FF /4) and CALL reg (FF /2)
                        ea = seg.startEA
                        while True:
                            ea = idaapi.find_binary(ea + 1, seg.endEA, "FF", 16, idaapi.SEARCH_DOWN)
                            if ea == idaapi.BADADDR: break

                            # Read possible ModR/M, SIB, and IMM8/IMM32 bytes
                            jop = idaapi.dbg_read_memory ( ea + 1, 0x6)

                            ###################################################
                            # JMP/CALL reg
                            if jop[0] in ["\xe0","\xe1","\xe2","\xe3","\xe4","\xe5","\xe6","\xe7",
                                          "\xd0","\xd1","\xd2","\xd3","\xd4","\xd5","\xd6","\xd7"]:
                                self.retns.append( (ea, m.file))

                            ###################################################
                            # JMP/CALL [reg] no SIB
                            # NOTE: Do not include pure [disp] instruction.

                            # JMP/CALL [reg] no *SP,*BP
                            elif jop[0] in ["\x20","\x21","\x22","\x23","\x26","\x27", 
                                            "\x10","\x11","\x12","\x13","\x16","\x17"]:
                                self.retns.append( (ea, m.file))

                            # JMP/CALL [reg + imm8] no *SP
                            elif jop[0] in ["\x60","\x61","\x62","\x63","\x65","\x66","\x67",
                                            "\x50","\x51","\x52","\x53","\x55","\x56","\x57"]:
                                jop_imm8 = jop[1]
                                jop_imm8 = unpack("b", jop_imm8)[0] # signed

                                if jop_imm8 <= self.maxJopImm:
                                    self.retns.append( (ea, m.file))


                            # JMP/CALL [reg + imm32] no *SP
                            elif jop[0] in ["\xa0","\xa1","\xa2","\xa3","\xa5","\xa6","\xa7",
                                            "\x90","\x91","\x92","\x93","\x95","\x96","\x97"]:
                                jop_imm32 = jop[1:5]
                                jop_imm32 = unpack("<i", jop_imm32)[0] # signed

                                if jop_imm32 <= self.maxJopImm:
                                    self.retns.append( (ea, m.file))

                            ###################################################
                            # JMP/CALL [reg] with SIB
                            # NOTE: Do no include pure [disp] instructions in SIB ([*] - none)
                            elif (jop[0] in ["\x24","\x64","\xa4"] and not jop[1] in ["\x25","\x65","\xad","\xe5"]) or \
                                 (jop[0] in ["\x14","\x54","\x94"] and not jop[1] in ["\x25","\x65","\xad","\xe5"]):

                                 # Check for displacement
                                if jop[0] in ["\x64","\x54"]:
                                    jop_imm8 = jop[2]
                                    jop_imm8 = unpack("b", jop_imm8)[0] # signed

                                    if jop_imm8 <= self.maxJopImm:
                                        self.retns.append( (ea, m.file))

                                elif jop[0] in ["\xa4","\x94"]:
                                    jop_imm32 = jop[2:6]
                                    jop_imm32 = unpack("<i", jop_imm32)[0] # signed

                                    if jop_imm32 <= self.maxJopImm:
                                        self.retns.append( (ea, m.file))

                                else:
                                    self.retns.append( (ea, m.file))

        print "[idasploiter] Found %d returns" % len(self.retns)

    def search_gadgets(self):

        count_total = len(self.retns)
        count_notify = 0
        count_curr  = 0

        # BUG: A separate flag is used to track user canceling the search,
        #      because multiple calls to idaapi.wasBreak() do not properly
        #      detect cancellations.
        breakFlag = False

        # Show wait dialog
        if not self.debug: idaapi.show_wait_box("Searching gadgets: 00%%%%")

        for (ea_end, module) in self.retns:

            # Flush the gadgets cache for each new retn pointer
            self.gadgets_cache = dict()

            # Flush memory cache for each new retn pointer
            self.dbg_mem_cache = None

            # CACHE: It is faster to read as much memory in one blob than to make incremental reads backwards.
            #        Try to read and cache self.maxRopOffset bytes back. In cases where it is not possible,
            #        then simply try to read the largest chunk.
            
            # NOTE: Read a bit extra to cover correct decoding of RETN, RETN imm16, CALL /2, and JMP /4 instructions.

            for i in range(self.maxRopOffset):
                self.dbg_mem_cache = idaapi.dbg_read_memory(ea_end - self.maxRopOffset + i, self.maxRopOffset - i + self.dbg_read_extra)
                if self.dbg_mem_cache != None:
                    break

            # Search all possible gadgets up to maxoffset bytes back
            # NOTE: Try all byte combinations to capture longer/more instructions
            #       even with bad bytes in the middle.
            for i in range(1, len(self.dbg_mem_cache) - self.dbg_read_extra):

                ea = ea_end - i

                # Get pointer charset
                ptr_charset = self.sploiter.get_ptr_charset(ea)

                # Filter the pointer
                if ptr_charset == None:                                    continue
                if self.ptrNonull     and not "nonull"     in ptr_charset: continue
                if self.ptrUnicode    and not "unicode"    in ptr_charset: continue
                if self.ptrAscii      and not "ascii"      in ptr_charset: continue
                if self.ptrAsciiPrint and not "asciiprint" in ptr_charset: continue
                if self.ptrAlphaNum   and not "alphanum"   in ptr_charset: continue
                if self.ptrNum        and not "numeric"    in ptr_charset: continue
                if self.ptrAlpha      and not "alpha"      in ptr_charset: continue

                # Try to build a gadget at the pointer
                gadget = self.build_gadget(ea, ea_end)

                # Successfully built the gadget
                if gadget:

                    # Populate gadget object with more data
                    gadget.address = ea
                    gadget.module = module
                    gadget.ptr_charset = ptr_charset

                    # Filter gadgets with too many instruction
                    if gadget.size > self.maxRopSize: 
                        break

                    # Append newly built gadget
                    self.gadgets.append(gadget)
                    self.gadgets_cache[ea] = gadget

                    # Exceeded maximum number of gadgets
                    if self.maxRops and len(self.gadgets) >= self.maxRops:
                        breakFlag = True
                        print "[idasploiter] Maximum number of gadgets exceeded."
                        break
                else:
                    self.gadgets_cache[ea] = None

                if breakFlag or idaapi.wasBreak():
                    breakFlag = True
                    break

            # Canceled
            # NOTE: Only works when started from GUI not script.
            if breakFlag or idaapi.wasBreak():
                breakFlag = True
                print "[idasploiter] Canceled."
                break

            # Progress report
            if not self.debug and count_curr >= count_notify:

                # NOTE: Need to use %%%% to escape both Python and IDA's format strings
                idaapi.replace_wait_box("Searching gadgets: %02d%%%%" % (count_curr*100/count_total) ) 

                count_notify += 0.10 * count_total

            count_curr += 1            

        print "[idasploiter] Found %d gadgets." % len(self.gadgets)
        if not self.debug: idaapi.hide_wait_box()


    # Attempt to build a gadget at the provided start address
    # by verifying it properly terminates at the expected RETN.
    def build_gadget(self, ea, ea_end):

        instructions  = list()
        chg_registers = set()
        use_registers = set()
        operations    = set()
        pivot         = 0

        # Process each instruction in the gadget
        while ea <= ea_end:

            ###################################################################
            # Gadget Level Cache:
            #
            # Locate a gadget (failed or built) starting at this address.
            # If one is located, then we don't need to process any further
            # instructions and just get necessary data from the cached
            # gadget to never have to process the same address twice.
            if ea in self.gadgets_cache:

                # Check if the gadget was build successfully
                gadget_cache = self.gadgets_cache[ea]
                
                # Build the reset of the gadget from cache
                if gadget_cache:    

                    for insn in gadget_cache.instructions:
                        instructions.append(insn)

                    for reg in gadget_cache.chg_registers:
                        chg_registers.add(reg)

                    for reg in gadget_cache.use_registers:
                        use_registers.add(reg)

                    for op in gadget_cache.operations:
                        operations.add(op)

                    pivot += gadget_cache.pivot

                    gadget = Gadget(instructions, pivot, operations, chg_registers, use_registers)
                    return gadget

                # Previous attempt to build gadget at this address failed
                else:
                    return None

            # Process new instruction
            else:

                # Instruction length
                # NOTE: decode_insn also sets global idaapi.cmd
                #       which contains insn_t structure
                insn_size = idaapi.decode_insn(ea)

                # Check successful decoding of the instruction
                if insn_size:

                    # Decoded instruction is too big to be a RETN or RETN imm16
                    if ea + insn_size > ea_end + self.dbg_read_extra:
                        return None

                    ###############################################################
                    # Instruction Level Cache
                    #
                    # Most instructions are repetitive so we can just cache
                    # unique byte combinations to avoid costly decoding more
                    # than once
                    
                    # Read instruction from memory cache
                    dbg_mem_offset = ea - (ea_end - (len(self.dbg_mem_cache) - self.dbg_read_extra) )
                    dbg_mem = self.dbg_mem_cache[dbg_mem_offset:dbg_mem_offset + insn_size]

                    # Create instruction cache if it doesn't already exist
                    if not dbg_mem in self.insn_cache:

                        ###########################################################
                        # Decode instruction
                        ###########################################################

                        # Get global insn_t structure describing the instruction
                        # NOTE: copy() is expensive, so we keep this single-threaded
                        insn = idaapi.cmd                  

                        #######################################################
                        # Decode and Cache instruction characteristics
                        self.insn_cache[dbg_mem] = self.decode_instruction(insn, ea, ea_end)

                    ##################################################################
                    # Retrieve cached instruction and apply it to the gadget    

                    # Check that cached instruction contains valid data
                    if self.insn_cache[dbg_mem]:

                        # Retrieve basic instruction characteristics
                        insn_mnem = self.insn_cache[dbg_mem]["insn_mnem"]
                        insn_disas = self.insn_cache[dbg_mem]["insn_disas"]
                        instructions.append(insn_disas)

                        #######################################################
                        # Expected ending instruction of the gadget
                        if ea == ea_end:
                            gadget = Gadget(instructions, pivot, operations, chg_registers, use_registers)
                            return gadget

                        #######################################################
                        # Filter out of place ROP/JOP/COP terminators
                        # NOTE: retn/jmp/call are allowed, but only in the last position

                        # Unexpected return instruction
                        elif insn_mnem == "retn":
                            return None

                        # Unexpected call/jmp instruction
                        elif insn_mnem in ["jmp","call"]:
                            return None

                        #######################################################
                        # Add instruction instruction characteristics to the gadget
                        else:

                            for reg in self.insn_cache[dbg_mem]["insn_chg_registers"]:
                                chg_registers.add(reg)

                            for reg in self.insn_cache[dbg_mem]["insn_use_registers"]:
                                use_registers.add(reg)

                            for op in self.insn_cache[dbg_mem]["insn_operations"]:
                                operations.add(op)

                            pivot += self.insn_cache[dbg_mem]["insn_pivot"]

                    # Previous attempt to decode the instruction invalidated the gadget
                    else:                        
                        return None
                        
                    ###############################################################
                    # Next instruction
                    # NOTE: This is outside cache
                    ea += insn_size

                ###################################################################
                # Failed decoding of the instruction
                # NOTE: Gadgets may have bad instructions in the middle which
                #       can be tolerated as long as we can find a useful instruction
                #       further out.
                else:

                    # HACK: IDA does not disassemble "\x00\x00" unless you enable
                    #       "Disassemble zero opcode instructions" in Processor Options.
                    #       Since this option is normally disabled, I will attempt
                    #       to get this instruction manually.

                    # Read two bytes from memory cache at current instruction candidate
                    dbg_mem_offset = ea - (ea_end - self.maxRopOffset)
                    dbg_mem = self.dbg_mem_cache[dbg_mem_offset:dbg_mem_offset + 2]

                    # Compare to two zero bytes
                    if dbg_mem[:2] == "\x00\x00":

                        if self.sploiter.addr64:
                            instructions.append("add [rax],al")
                        else:
                            instructions.append("add [eax],al")

                        use_registers.add("al")
                        operations.add("reg-to-mem")

                        ea += 2

                    # "MOV Sreg, r/m16" instructions will result in illegal instruction exception: c000001d
                    # or the memory couldn't be read exception: c0000005 which we don't want in our gadgets.
                    elif dbg_mem[0] == "\x8E":
                        return None

                    # Record a "bad byte" if allowed
                    elif dbg_mem and not self.ropNoBadBytes:
                        byte = dbg_mem[0]

                        instructions.append("db %sh" % binascii.hexlify(byte))

                        ea += 1

                    # Invalidate the gadget
                    else:
                        return None

        # Failed to build a gadget, because RETN instruction was not found
        else:    
            return None

    ###############################################################
    # Decode instruction

    def decode_instruction(self, insn, ea, ea_end):

        # Instruction specific characteristics
        insn_chg_registers = set()
        insn_use_registers = set()
        insn_operations = set()
        insn_pivot = 0

        # Instruction feature
        #
        # instruc_t.feature
        #
        # CF_STOP = 0x00001 #  Instruction doesn't pass execution to the next instruction
        # CF_CALL = 0x00002 #  CALL instruction (should make a procedure here)
        # CF_CHG1 = 0x00004 #  The instruction modifies the first operand
        # CF_CHG2 = 0x00008 #  The instruction modifies the second operand
        # CF_CHG3 = 0x00010 #  The instruction modifies the third operand
        # CF_CHG4 = 0x00020 #  The instruction modifies 4 operand
        # CF_CHG5 = 0x00040 #  The instruction modifies 5 operand
        # CF_CHG6 = 0x00080 #  The instruction modifies 6 operand
        # CF_USE1 = 0x00100 #  The instruction uses value of the first operand
        # CF_USE2 = 0x00200 #  The instruction uses value of the second operand
        # CF_USE3 = 0x00400 #  The instruction uses value of the third operand
        # CF_USE4 = 0x00800 #  The instruction uses value of the 4 operand
        # CF_USE5 = 0x01000 #  The instruction uses value of the 5 operand
        # CF_USE6 = 0x02000 #  The instruction uses value of the 6 operand
        # CF_JUMP = 0x04000 #  The instruction passes execution using indirect jump or call (thus needs additional analysis)
        # CF_SHFT = 0x08000 #  Bit-shift instruction (shl,shr...)
        # CF_HLL  = 0x10000 #  Instruction may be present in a high level language function.
        insn_feature = insn.get_canon_feature()

        # Instruction mnemonic name
        insn_mnem = insn.get_canon_mnem()

        #if insn_mnem in self.mnems: self.mnems[insn_mnem] += 1
        #else:                       self.mnems[insn_mnem]  = 1

        # Get instruction operand types
        #
        # op_t.type
        #                    Description                          Data field
        # o_void     =  0 #  No Operand                           ----------
        # o_reg      =  1 #  General Register (al,ax,es,ds...)    reg
        # o_mem      =  2 #  Direct Memory Reference  (DATA)      addr
        # o_phrase   =  3 #  Memory Ref [Base Reg + Index Reg]    phrase
        # o_displ    =  4 #  Memory Reg [Base Reg + Index Reg + Displacement] phrase+addr
        # o_imm      =  5 #  Immediate Value                      value
        # o_far      =  6 #  Immediate Far Address  (CODE)        addr
        # o_near     =  7 #  Immediate Near Address (CODE)        addr
        insn_op1 = insn.Operands[0].type
        insn_op2 = insn.Operands[1].type

        ###############################################################
        # Filter gadget
        ###############################################################

        # Do not filter ROP, JOP, COP, always decode them
        # NOTE: A separate check must be done to check if they are out of place.
        if not insn_mnem in ["retn","jmp","call"]:

            # Filter gadgets with instructions that don't forward execution to the next address
            if insn_feature & idaapi.CF_STOP:
                return None

            # Filter gadgets with instructions in a bad list
            elif insn_mnem in self.ropBadMnems:
                return None

            # Filter gadgets with jump instructions
            # Note: conditional jumps may still be useful if we can
            #       set flags prior to calling them.
            elif not self.ropAllowJcc and insn_mnem[0] == "j":
                return None

        ###############################################################
        # Get disassembly
        ###############################################################
        # NOTE: GENDSM_FORCE_CODE ensures correct decoding
        #       of split instructions.
        insn_disas = idc.GetDisasmEx(ea, idaapi.GENDSM_FORCE_CODE)                
        insn_disas = insn_disas.partition(';')[0]       # Remove comments from disassembly                       
        insn_disas = ' '.join(insn_disas.split())       # Remove extraneous space from disassembly

        ###############################################################
        # Analyze instruction
        ###############################################################

        # Standalone instruction
        if insn_op1 == idaapi.o_void:

            # TODO: Determine and test how these instructions affect the stack
            #       in 32-bit and 64-bit modes.
            if insn_mnem in ["pusha","pushad","popa","popad","pushf","pushfd","pushfq","popf","popfd","popfq"]:
                insn_operations.add("stack")

                if insn_mnem in ["popa","popad"]:
                    insn_pivot += 7*4
                elif insn_mnem in ["pusha","pushad"]:
                    insn_pivot -= 8*4
                elif insn_mnem in ["popf","popfd"]:
                    insn_pivot += 4
                elif insn_mnem in ["pushf","pushfd"]:
                    insn_pivot -= 4
                elif insn_mnem == "popfq":  # TODO: Needs testing
                    insn_pivot += 8
                elif insn_mnem == "pushfq": # TODO: Needs testing
                    insn_pivot -= 8

        # Single operand instruction
        elif insn_op2 == idaapi.o_void:

            # Single operand register
            if insn_op1 == idaapi.o_reg:
                insn_operations.add("one-reg")

                if insn_feature & idaapi.CF_CHG1:
                    reg_name = self.get_o_reg_name(insn, 0)
                    insn_chg_registers.add(reg_name)

                    # Check for stack operation
                    if reg_name[1:] == "sp":
                        insn_operations.add("stack")

                        if insn_mnem == "inc":
                            insn_pivot += 1

                        elif insn_mnem == "dec":
                            insn_pivot -= 1

                elif insn_feature & idaapi.CF_USE1:
                    reg_name = self.get_o_reg_name(insn, 0)
                    insn_use_registers.add(reg_name)

            # Single operand immediate
            elif insn_op1 == idaapi.o_imm:
                insn_operations.add("one-imm")

            # Single operand reference
            # TODO: determine the [reg + ...] value if present
            elif insn_op1 == idaapi.o_phrase or insn_op1 == idaapi.o_displ:
                insn_operations.add("one-mem")

            # PUSH/POP mnemonic with a any operand type
            if insn_mnem in ["push","pop"]:
                insn_operations.add("stack")

                # Adjust pivot based on operand size (32bit vs 64bit)
                if insn_mnem == "pop":
                    if   insn.Operands[0].dtyp == idaapi.dt_dword: insn_pivot += 4
                    elif insn.Operands[0].dtyp == idaapi.dt_qword: insn_pivot += 8
                elif insn_mnem == "push":                   
                    if   insn.Operands[0].dtyp == idaapi.dt_dword: insn_pivot -= 4
                    elif insn.Operands[0].dtyp == idaapi.dt_qword: insn_pivot -= 8

            # Check for arithmetic operation:
            if insn_mnem in self.insn_arithmetic_ops:
                insn_operations.add("math")

            # Check for bit-wise operations:
            if insn_mnem in self.insn_bit_ops:
                insn_operations.add("bit")

        # Two operand instruction
        else:

            # Check for arithmetic operations
            if insn_mnem in self.insn_arithmetic_ops:
                insn_operations.add("math")

            # Check for bit-wise operations
            if insn_mnem in self.insn_bit_ops:
                insn_operations.add("bit")

            # Two operand instruction with the first operand a register
            if insn_op1 == idaapi.o_reg:

                reg_name = self.get_o_reg_name(insn, 0)

                if insn_feature & idaapi.CF_CHG1:
                    insn_chg_registers.add(reg_name)

                    # Check for stack operation
                    if reg_name[1:] == "sp":
                        insn_operations.add("stack")

                        # Determine stack pivot distance
                        if insn_op2 == idaapi.o_imm:

                            # NOTE: adb and sbb may also be useful, but let the user
                            #       determine their use by locating the operations "stack"
                            if insn_mnem == "add":
                                insn_pivot += insn.Operands[1].value

                            elif insn_mnem == "sub":
                                insn_pivot -= insn.Operands[1].value

                    # Check for operations
                    if insn_op2 == idaapi.o_reg:
                        insn_operations.add("reg-to-reg")
                    elif insn_op2 == idaapi.o_imm:
                        insn_operations.add("imm-to-reg")

                    # TODO: determine the [reg + ...] value if present
                    elif insn_op2 == idaapi.o_phrase or insn_op2 == idaapi.o_displ:
                        insn_operations.add("mem-to-reg")

                if insn_feature & idaapi.CF_USE1:
                    insn_use_registers.add(reg_name)


            # Two operand instruction with the second operand a register
            if insn_op2 == idaapi.o_reg:

                reg_name = self.get_o_reg_name(insn, 1)

                if insn_feature & idaapi.CF_CHG2:
                    insn_chg_registers.add(reg_name)

                    # Check for stack operation
                    if reg_name[1:] == "sp":
                        insn_operations.add("stack")

                if insn_feature & idaapi.CF_USE2:
                    insn_use_registers.add(reg_name)

                # Check for operations
                # TODO: determine the [reg + ...] value if present
                if insn_op1 == idaapi.o_phrase or insn_op1 == idaapi.o_displ:
                    insn_operations.add("reg-to-mem")

        # Build instruction dictionary
        insn = dict()
        insn["insn_mnem"] = insn_mnem
        insn["insn_disas"] = insn_disas
        insn["insn_operations"] = insn_operations
        insn["insn_chg_registers"] = insn_chg_registers
        insn["insn_use_registers"] = insn_use_registers
        insn["insn_pivot"] = insn_pivot

        return insn

###############################################################################
# Sploiter Engine

class Sploiter():

    def __init__(self):
        


        # Process modules list
        self.modules = list()
        self.rop     = None
        self.funcptr = None

        # Default patterns
        self.c4_list = string.punctuation
        self.c3_list = string.uppercase
        self.c2_list = string.lowercase
        self.c1_list = string.digits

        if not idaapi.ph.id == idaapi.PLFM_386:
            print "[idasploiter] Only Intel 80x86 processors are supported."
            sys.exit(1)

        # Check if processor supports 64-bit addressing
        if idaapi.ph.flag & idaapi.PR_USE64:
            self.addr64 = True
            self.addr_format = "%016X"
            self.pack_format_be = ">Q"
            self.pack_format_le = "<Q"
        else:
            self.addr64 = False
            self.addr_format = "%08X"
            self.pack_format_be = ">I"
            self.pack_format_le = "<I"

    def get_ptr_charset(self, ea):

        ptr_charset = []

        # Flags. True until proven otherwise.
        nonull     = True
        unicode    = True
        ascii      = True
        asciiprint = True        
        alphanum   = True
        alpha      = True
        numeric    = True
        
        ptr_bytes = pack(self.pack_format_be, ea)
        for i,b in enumerate( ptr_bytes ):

            b_int = ord(b)

            if b in self.ptrBadChars:
                return None

            # Locate any null bytes
            if nonull and b == "\x00":
                nonull = False

            # Unicode compatible address must have
            # null bytes at even points in the address
            if unicode and not i % 2:

                if not ptr_bytes[i:i+2] in ASCII_TO_UNICODE[ self.unicodeTable ]:
                    unicode = False

            # Find any non-ascii characters
            if ascii and not b_int > 127:

                # Find any non-ascii printable characters
                if asciiprint and not (b_int < 32 or b_int > 126):

                    # Find any non-alphanumeric characters
                    if alphanum and (b in string.ascii_letters or b in string.digits):

                        # Find any non-numeric characters
                        if numeric and not b in string.digits:
                            numeric = False

                        # Find any non-letter characters
                        if alpha and not b in string.ascii_letters:
                            alpha = False
                    else:                        
                        alphanum = False
                        numeric  = False
                        alpha     = False
                else:                    
                    asciiprint = False
                    alphanum   = False
                    numeric    = False
                    alpha      = False
            else:                
                ascii      = False
                asciiprint = False
                alphanum   = False
                numeric    = False
                alpha      = False

                # NOTE: You can continue to filter for upper/lower here if necessary
        
        if nonull:     ptr_charset.append("nonull")

        if unicode:    ptr_charset.append("unicode")

        if ascii:      
            ptr_charset.append("ascii")

            if asciiprint: 
                ptr_charset.append("asciiprint")

                if alphanum:   
                    ptr_charset.append("alphanum")

                    if numeric:    ptr_charset.append("numeric")
                    if alpha:      ptr_charset.append("alpha")
        

        return ptr_charset  

    def process_modules(self):

        # Reset modules list
        self.modules = list()

        for m in idautils.Modules():

            module = Module(m.name, m.size, m.base, m.rebase_to)
            self.modules.append(module)

    def show_modules_view(self):
        mod = ModuleView(self)
        mod.show()

    def process_rop(self, select_list = None):

        # Initialize ROP gadget search engine
        self.rop = Rop(self)

        # Prompt user for ROP search settings
        f = RopForm(self, select_list)
        ok = f.Execute()
        if ok == 1:

            # Configure ROP gadget search engine

            # Get selected modules
            self.rop.modules = [self.modules[i] for i in f.mod.GetEmbSelection()]

            if len(self.rop.modules) > 0:

                # Pointer filters
                self.rop.ptrNonull     = f.cPtrNonull.checked
                self.rop.ptrUnicode    = f.cPtrUnicode.checked
                self.rop.ptrAscii      = f.cPtrAscii.checked
                self.rop.ptrAsciiPrint = f.cPtrAsciiPrint.checked
                self.rop.ptrAlphaNum   = f.cPtrAlphaNum.checked
                self.rop.ptrAlpha      = f.cPtrAlpha.checked
                self.rop.ptrNum        = f.cPtrNum.checked

                # Filter bad characters
                buf                    = f.strBadChars.value
                buf = buf.replace(' ','')         # remove spaces
                buf = buf.replace('\\x','')       # remove '\x' prefixes
                buf = buf.replace('0x','')        # remove '0x' prefixes
                try:
                    buf = binascii.unhexlify(buf) # convert to bytes
                    self.ptrBadChars   = buf
                except Exception, e:
                    idaapi.warning("Invalid input: %s" % e)
                    self.ptrBadChars   = ""

                # Ascii_to_Unicode_transformation table
                # BUG: DropdownControl does not work on IDA 6.5
                self.unicodeTable      = f.radUnicode.value

                # ROP instruction filters
                self.rop.ropBadMnems   = [mnem.strip().lower() for mnem in f.strBadMnems.value.split(',')]
                self.rop.ropAllowJcc   = f.cRopAllowJcc.checked
                self.rop.ropNoBadBytes = f.cRopNoBadBytes.checked

                # Get ROP engine settings
                self.rop.maxRopSize    = f.intMaxRopSize.value
                self.rop.maxRopOffset  = f.intMaxRopOffset.value
                self.rop.maxRops       = f.intMaxRops.value
                self.rop.maxRetnImm    = f.intMaxRetnImm.value

                # Gadget search values
                self.rop.searchRop     = f.cRopSearch.checked
                self.rop.searchJop     = f.cJopSearch.checked

                # Search for returns and ROP gadgets
                self.rop.search_retns()
                self.rop.search_gadgets()

                # Show the ROP gadgets view
                ropView = RopView(self)
                ropView.show()

            else:
                idaapi.warning("No modules selected.")

        f.Free()

    def process_funcptr(self, select_list = None):

        # Prompt user for function pointer search settings
        f = PtrForm(self, select_list)
        ok = f.Execute()
        if ok == 1:

            # Initialize Function Pointer search engine
            self.funcptr = FuncPtr(self)

            # Configure Function Pointer search engine

            # Get selected modules
            self.funcptr.modules = [self.modules[i] for i in f.mod.GetEmbSelection()]

            if len(self.funcptr.modules) > 0:

                # Pointer filters
                self.funcptr.ptrNonull     = f.cPtrNonull.checked
                self.funcptr.ptrUnicode    = f.cPtrUnicode.checked
                self.funcptr.ptrAscii      = f.cPtrAscii.checked
                self.funcptr.ptrAsciiPrint = f.cPtrAsciiPrint.checked
                self.funcptr.ptrAlphaNum   = f.cPtrAlphaNum.checked
                self.funcptr.ptrAlpha      = f.cPtrAlpha.checked
                self.funcptr.ptrNum        = f.cPtrNum.checked

                # Apply pointer filters to the funcptr or ptr-to-ptr
                if f.rFilterP2P.selected:
                    self.funcptr.filterP2P = True
                else:
                    self.funcptr.filterP2P = False

                # Ascii_to_Unicode_transformation table
                # BUG: DropdownControl does not work on IDA 6.5
                self.unicodeTable      = f.radUnicode.value

                # Filter bad characters
                buf                    = f.strBadChars.value
                buf = buf.replace(' ','')         # remove spaces
                buf = buf.replace('\\x','')       # remove '\x' prefixes
                buf = buf.replace('0x','')        # remove '0x' prefixes
                try:
                    buf = binascii.unhexlify(buf) # convert to bytes
                    self.ptrBadChars = buf
                except Exception, e:
                    idaapi.warning("Invalid input: %s" % e)
                    self.ptrBadChars = ""

                # Offsets
                self.funcptr.ptrOffset = int(f.intPtrOffset.value, 16)
                self.funcptr.p2pOffset = int(f.intP2POffset.value, 16)

                # Pointer search engine settings
                self.funcptr.searchP2P = not f.rSearchPtrOnly.checked
                self.funcptr.maxPtrs   = f.intMaxPtrs.value

                # Search for writable pointers
                self.funcptr.search_pointers()

                # Show the writable pointers view
                ptrView = PtrView(self)
                ptrView.show()

            else:
                idaapi.warning("No modules selected.")

        f.Free()

    def pattern_create(self):

        f = PatternCreateForm(self)
        ok = f.Execute()
        if ok == 1:  
            pass

        f.Free()

    def pattern_detect(self, debugger=True):

        f = PatternDetectForm(self, debugger)
        ok = f.Execute()
        if ok == 1:  
            pass

        f.Free()

    def process_compare(self, debugger=True):

        f = CompareForm(self, debugger)
        ok = f.Execute()
        if ok == 1:
            pass

        f.Free()


###############################################################################
# Sploiter UI
###############################################################################

###############################################################################
# Module UI
###############################################################################

class ModuleView(Choose2):
    """
    Chooser class to display security characteristics of loaded modules.
    """
    def __init__(self, sploiter, embedded = False):

        self.sploiter = sploiter

        Choose2.__init__(self,
                         "Modules",
                         [ ["Address",  13 | Choose2.CHCOL_HEX], 
                           ["Name",     10 | Choose2.CHCOL_PLAIN], 
                           ["Size",     10 | Choose2.CHCOL_HEX],
                           ["SafeSEH",   6 | Choose2.CHCOL_PLAIN],
                           ["ASLR",      6 | Choose2.CHCOL_PLAIN], 
                           ["DEP",       6 | Choose2.CHCOL_PLAIN],
                           ["Canary",    6 | Choose2.CHCOL_PLAIN],
                           ["Path",     40 | Choose2.CHCOL_PLAIN], 
                         ],
                         flags = Choose2.CH_MULTI,  # Select multiple modules
                         embedded=embedded)

        self.icon = 150

        # Items for display
        self.items = list()

        # Initialize/Refresh the view
        self.refreshitems()

        # Selected items
        self.select_list = list()

        # Command callbacks
        self.cmd_load_module     = None
        self.cmd_search_gadgets  = None
        self.cmd_search_pointers = None

    def show(self):
        # Attempt to open the view
        if self.Show() < 0: return False

        # Add extra context menu commands
        # NOTE: Make sure you check for duplicates
        if self.cmd_load_module == None:
            self.cmd_load_module = self.AddCommand("Load module...", flags = idaapi.CHOOSER_POPUP_MENU, icon=135)
        if self.cmd_search_gadgets == None:
            self.cmd_search_gadgets = self.AddCommand("Search gadgets...", flags = idaapi.CHOOSER_POPUP_MENU | idaapi.CHOOSER_MULTI_SELECTION, icon=182 )
        if self.cmd_search_pointers == None:
            self.cmd_search_pointers = self.AddCommand("Search function pointers...", flags = idaapi.CHOOSER_POPUP_MENU | idaapi.CHOOSER_MULTI_SELECTION, icon=143 )


        return True

    def refreshitems(self):
        self.items = list()

        for m in self.sploiter.modules:
            self.items.append([ self.sploiter.addr_format % m.addr, 
                                m.file, 
                                "%08X" % m.size,
                                m.SafeSEH, m.ASLR,
                                m.NXCompat,
                                m.GS,
                                m.path])

    def OnCommand(self, n, cmd_id):

        # Search ROP gadgets
        if cmd_id == self.cmd_search_gadgets:

            # Empty selection
            if n == -1:

                # Initialize ROP gadget form with empty selection
                self.sploiter.process_rop(select_list = self.select_list)

            # Selection start
            elif n == -2:
                self.select_list = list()

            # Selection end
            elif n == -3:

                # Initialize ROP gadget form with user selection
                self.sploiter.process_rop(select_list = self.select_list)
                self.select_list = list()

            # Selection number
            else:
                self.select_list.append(n)

        # Search function pointers
        elif cmd_id == self.cmd_search_pointers:

            # Empty selection
            if n == -1:

                # Initialize function pointers form with empty selection
                self.sploiter.process_funcptr(select_list = self.select_list)

            # Selection start
            elif n == -2:
                self.select_list = list()

            # Selection end
            elif n == -3:

                # Initialize function pointers form with user selection
                self.sploiter.process_funcptr(select_list = self.select_list)
                self.select_list = list()

            # Selection number
            else:
                self.select_list.append(n)

        elif cmd_id == self.cmd_load_module:

            module_name = idaapi.askfile_c(0, "*.*", "Please select module to load")
            if module_name:
                print "[idasploiter] Loading module: %s" % module_name     
                loadlib = idaapi.Appcall.proto("kernel32_LoadLibraryA", "int __stdcall loadlib(const char *fn);")
                hmod = loadlib(module_name)
                if hmod:
                    print "[idasploiter] Finished loading module: %s" % module_name
                else:
                    print "[idasploiter] Could not load: %s" % module_name

                self.refreshitems()

        return 1

    def OnSelectLine(self, n):
        pass

    def OnGetLine(self, n):
        return self.items[n]

    def OnGetIcon(self, n):

        if not len(self.items) > 0:
            return -1

        m = self.sploiter.modules[n]
        if m.SafeSEH != "No" and m.ASLR == "Yes" and m.NXCompat == "Yes" and m.GS != "No":
            return 61
        elif m.ASLR == "Yes":
            return 60
        else:
            return 59

    def OnClose(self):
        self.cmd_search_gadgets = None
        self.cmd_search_pointers    = None

    def OnGetSize(self):
        return len(self.items)

    def OnRefresh(self, n):
        self.refreshitems()
        return n

    def OnActivate(self):
        self.refreshitems()

###############################################################################
# ROP UI
###############################################################################

###############################################################################
# ROP/JOP/COP Form

class RopForm(Form):

    def __init__(self, sploiter, select_list = None):

        self.sploiter = sploiter
        self.select_list = select_list

        self.mod = ModuleView(self.sploiter, embedded=True)

        Form.__init__(self, 
r"""BUTTON YES* Search
Search ROP gadgets

{FormChangeCb}<Modules:{cEChooser}>

Pointer Charset:                  Search Settings:
<nonull:{cPtrNonull}>        <Bad Chars        :{strBadChars}>     
<unicode:{cPtrUnicode}>       Unicode Table    <ANSI:{rUnicodeANSI}><OEM:{rUnicodeOEM}><UTF7:{rUnicodeUTF7}><UTF8:{rUnicodeUTF8}>{radUnicode}>
<ascii:{cPtrAscii}>         <Bad Instructions :{strBadMnems}>
<asciiprint:{cPtrAsciiPrint}>    <Max gadget size  :{intMaxRopSize}>
<alphanum:{cPtrAlphaNum}>      <Max gadget offset:{intMaxRopOffset}>
<alpha:{cPtrAlpha}>         <Max RETN imm16   :{intMaxRetnImm}>
<numeric:{cPtrNum}>{ptrGroup}>       <Max JOP imm8/32  :{intMaxJopImm}>
                <Max gadgets      :{intMaxRops}>
                Other settings   <Allow conditional jumps:{cRopAllowJcc}>
                <Do not allow bad bytes:{cRopNoBadBytes}>
                <Search for ROP gadgets:{cRopSearch}>
                <Search for JOP gadgets:{cJopSearch}>{ropGroup}>


""", {
                'cEChooser'       : Form.EmbeddedChooserControl(self.mod, swidth=131),
                'ptrGroup'        : Form.ChkGroupControl(("cPtrNonull", "cPtrAscii", "cPtrAsciiPrint", "cPtrUnicode",'cPtrAlphaNum','cPtrAlpha','cPtrNum')),
                'ropGroup'        : Form.ChkGroupControl(('cRopAllowJcc','cRopNoBadBytes','cRopSearch','cJopSearch')),
                'intMaxRopSize'   : Form.NumericInput(swidth=4,tp=Form.FT_DEC,value=self.sploiter.rop.maxRopSize),
                'intMaxRopOffset' : Form.NumericInput(swidth=4,tp=Form.FT_DEC,value=self.sploiter.rop.maxRopOffset),
                'intMaxRops'      : Form.NumericInput(swidth=4,tp=Form.FT_DEC,value=self.sploiter.rop.maxRops),
                'intMaxRetnImm'   : Form.NumericInput(swidth=4,tp=Form.FT_HEX,value=self.sploiter.rop.maxRetnImm),
                'intMaxJopImm'    : Form.NumericInput(swidth=4,tp=Form.FT_HEX,value=self.sploiter.rop.maxJopImm),
                'strBadChars'     : Form.StringInput(swidth=70,tp=Form.FT_ASCII),
                'radUnicode'      : Form.RadGroupControl(("rUnicodeANSI","rUnicodeOEM","rUnicodeUTF7","rUnicodeUTF8")),
                'strBadMnems'     : Form.StringInput(swidth=80,tp=Form.FT_ASCII,value="leave, int, into, enter, syscall, sysenter, sysexit, sysret, in, out, loop, loope, loopne, lock, rep, repe, repz, repne, repnz"),
                'FormChangeCb'    : Form.FormChangeCb(self.OnFormChange),
            })

        self.Compile()

    def OnFormChange(self, fid):

        # Form initialization
        if fid == -1:
            self.SetFocusedField(self.cEChooser)

            # Preselect non-ASLR modules on startup if none were already specified
            if self.select_list == None:

                self.select_list = list()
                for i, m in enumerate(self.sploiter.modules):
                    if m.ASLR == "No":
                        self.select_list.append(i)

            self.SetControlValue(self.cEChooser, self.select_list)

            # Enable both ROP and JOP search by default
            self.SetControlValue(self.cRopSearch, True)
            self.SetControlValue(self.cJopSearch, True)

            # Skip bad instructions by default
            self.SetControlValue(self.cRopNoBadBytes, True)

        # Form OK pressed
        elif fid == -2:
            pass

        return 1

###############################################################################
# ROP Viewer

class RopView(Choose2):
    """
    Chooser class to display security characteristics of loaded modules.
    """
    def __init__(self, sploiter):

        self.sploiter = sploiter

        Choose2.__init__(self,
                         "ROP gadgets",
                         [ ["Address",           13 | Choose2.CHCOL_HEX], 
                           ["Gadget",            30 | Choose2.CHCOL_PLAIN], 
                           ["Module",            10 | Choose2.CHCOL_PLAIN],
                           ["Size",               3 | Choose2.CHCOL_DEC],
                           ["Pivot",              4 | Choose2.CHCOL_DEC],
                           ["Operations",         12 | Choose2.CHCOL_PLAIN],
                           ["Changed Registers", 12 | Choose2.CHCOL_PLAIN],
                           ["Used Registers",    12 | Choose2.CHCOL_PLAIN],
                           ["Charset",           12 | Choose2.CHCOL_PLAIN],
                         ],
                         flags = Choose2.CH_MULTI)

        self.icon = 182

        # Items for display
        self.items = []

        # Initialize/Refresh the view
        self.refreshitems()

        # Selected items
        self.gadget_chain = list()

        # Command callbacks
        self.cmd_chain_add   = None
        self.cmd_chain_build = None
        self.cmd_chain_clear = None
        self.cmd_export_csv  = None


    def show(self):
        # Attempt to open the view
        if self.Show() < 0: return False

        if self.cmd_chain_add == None:
            self.cmd_chain_add   = self.AddCommand("Add to chain",     flags = idaapi.CHOOSER_POPUP_MENU | idaapi.CHOOSER_MULTI_SELECTION, icon=50)
        if self.cmd_chain_build == None:
            self.cmd_chain_build = self.AddCommand("Build chain",      flags = idaapi.CHOOSER_POPUP_MENU, icon=156)
        if self.cmd_chain_clear == None:
            self.cmd_chain_clear = self.AddCommand("Clear chain",      flags = idaapi.CHOOSER_POPUP_MENU, icon=32)
        if self.cmd_export_csv == None:
            self.cmd_export_csv  = self.AddCommand("Export as csv...", flags = idaapi.CHOOSER_POPUP_MENU, icon=40)

        return True

    def refreshitems(self):
        self.items = []

        for g in self.sploiter.rop.gadgets:

            self.items.append([ self.sploiter.addr_format % g.address, 
                                " # ".join(g.instructions), 
                                g.module, 
                                "%d" % g.size, 
                                "%d" % g.pivot, 
                                ", ".join(g.operations), 
                                ", ".join(g.chg_registers), 
                                ", ".join(g.use_registers), 
                                ", ".join(g.ptr_charset)
                                ])

    def OnCommand(self, n, cmd_id):

        # Build ROP chain
        if cmd_id == self.cmd_chain_build:

            # Display ROP chain builder form
            f = RopChainForm(self)
            ok = f.Execute()
            if ok == 1:
                pass
            f.Free()

        # Clear ROP chain
        elif cmd_id == self.cmd_chain_clear:
            self.gadget_chain = list()

        # Add gadget to ROP chain
        elif cmd_id == self.cmd_chain_add:
            if n >= 0:
                self.gadget_chain.append( self.sploiter.rop.gadgets[n] )

        # Export CSV
        elif cmd_id == self.cmd_export_csv:

            file_name = idaapi.askfile_c(1, "*.csv", "Please enter CSV file name")
            if file_name:
                print "[idasploiter] Exporting gadgets to %s" % file_name
                with open(file_name, 'wb') as csvfile:
                    csvwriter = csv.writer(csvfile, delimiter=',',
                                            quotechar='"', quoting=csv.QUOTE_MINIMAL)
                    csvwriter.writerow(["Address","Gadget","Module","Size","Pivot","Operations","Changed Registers","Used Registers","Charset"])
                    for item in self.items:
                        csvwriter.writerow(item)

        return 1

    def OnSelectLine(self, n):
        idaapi.jumpto( self.sploiter.rop.gadgets[n].address )

    def OnGetLine(self, n):
        return self.items[n]

    def OnClose(self):
        self.cmd_chain_add   = None
        self.cmd_chain_build = None
        self.cmd_chain_clear = None
        self.cmd_export_csv  = None

    def OnGetSize(self):
        return len(self.items)

    def OnRefresh(self, n):
        self.refreshitems()
        return n

    def OnActivate(self):
        self.refreshitems()

###############################################################################
# ROP Chain Builder UI
###############################################################################

###############################################################################
# ROP Chain Viewer

class RopChainView(Choose2):
    """
    Chooser class to display security characteristics of loaded modules.
    """
    def __init__(self, ropview):

        self.ropview = ropview

        Choose2.__init__(self,
                         "ROP gadgets",
                         [ ["Address",           13 | Choose2.CHCOL_HEX], 
                           ["Gadget",            30 | Choose2.CHCOL_PLAIN], 
                           ["Module",            10 | Choose2.CHCOL_PLAIN],
                           ["Size",               3 | Choose2.CHCOL_DEC],
                           ["Pivot",              4 | Choose2.CHCOL_DEC],
                           ["Operations",        12 | Choose2.CHCOL_PLAIN],
                           ["Changed Registers", 12 | Choose2.CHCOL_PLAIN],
                           ["Used Registers",    12 | Choose2.CHCOL_PLAIN],
                           ["Charset",           12 | Choose2.CHCOL_PLAIN],
                         ],
                         flags = Choose2.CH_MULTI,
                         embedded = True)

        self.icon = 182

        # Items for display and corresponding data
        self.items = []

        # Selected items
        self.select_list = list()

        # Initialize/Refresh the view
        self.refreshitems()

    def refreshitems(self):
        self.items = []

        for g in self.ropview.gadget_chain:

            self.items.append([ self.ropview.sploiter.addr_format % g.address, 
                                " # ".join(g.instructions), 
                                g.module, 
                                "%d" % g.size if g.size else "", 
                                "%d" % g.pivot if g.pivot else "", 
                                ", ".join(g.operations), 
                                ", ".join(g.chg_registers), 
                                ", ".join(g.use_registers), 
                                ", ".join(g.ptr_charset)
                                ])

        return

    def OnCommand(self, n, cmd_id):

        if cmd_id == self.cmd_move_up:
            if n > 0:
                self.ropview.gadget_chain.insert(n-1, self.ropview.gadget_chain.pop(n))

        # NOTE: Moving items down should be done in one go in reverse
        #       to avoid self.items from desyncing from gadget_chain list.
        elif cmd_id == self.cmd_move_down:

            # Empty selection
            if n == -1:
                pass

            # Selection start
            elif n == -2:
                self.select_list = list()

            # Selection end
            elif n == -3:
                for i in sorted(self.select_list, reverse=True):
                    self.ropview.gadget_chain.insert(i+1, self.ropview.gadget_chain.pop(i))
                self.select_list = list()

            else:
                self.select_list.append(n)

        # NOTE: Duplicating items should be done in one go in reverse
        #       to avoid self.items desyncing from gadget_chain list.
        elif cmd_id == self.cmd_duplicate:
            
            # Empty selection
            if n == -1:
                pass

            # Selection start
            elif n == -2:
                self.select_list = list()

            # Selection end
            elif n == -3:

                # Duplicate selection after the last index
                last_i = self.select_list[-1]

                # Duplicate selected items in reverse
                for i in sorted(self.select_list, reverse=True):

                    if last_i == len(self.ropview.gadget_chain) - 1:
                        self.ropview.gadget_chain.append(self.ropview.gadget_chain[i])
                    else:
                        self.ropview.gadget_chain.insert(last_i + 1, self.ropview.gadget_chain[i])

                self.select_list = list()

            else:
                self.select_list.append(n)

        # NOTE: Deleting items should be done in one go in reverse
        #       to avoid self.items desyncing from gadget_chain list.
        elif cmd_id == self.cmd_delete:
            
            # Empty selection
            if n == -1:
                pass

            # Selection start
            elif n == -2:
                self.select_list = list()

            # Selection end
            elif n == -3:
                for i in sorted(self.select_list, reverse=True):
                    del self.ropview.gadget_chain[i]
                self.select_list = list()

            else:
                self.select_list.append(n)



        return 1

    def OnSelectLine(self, n):
        pass

    def OnGetLine(self, n):
        return self.items[n]

    def OnClose(self):
        pass

    def OnGetSize(self):
        return len(self.items)

    def OnRefresh(self, n):
        self.refreshitems()
        return n

    def OnRefreshed(self):
        self.refreshitems()
        return len(self.items)

    def OnActivate(self):
        self.refreshitems()


###############################################################################
# ROPChain Form

class RopChainForm(Form):

    def __init__(self, ropview):

        self.ropview = ropview

        self.ropchainview = RopChainView(self.ropview)

        Form.__init__(self, 
r"""BUTTON YES* NONE
BUTTON CANCEL NONE
ROP Chain Builder

{FormChangeCb}<ROP gadgets:{cEChooser}>
<:{strAddr}><:{strComment}><Insert:{iButtonInsert}>
<:{strRopChain}>
Formats<Plain:{rFmtPlain}><Python:{rFmtPython}><Ruby:{rFmtRuby}><Perl:{rFmtPerl}><JavaScript:{rFmtJS}>{rGroup}> 
{spacer}<Generate:{iButtonGenerate}>
""", {
                'spacer'          : Form.StringLabel("                                "),
                'cEChooser'       : Form.EmbeddedChooserControl(self.ropchainview, swidth=125),
                'strAddr'         : Form.StringInput(value="0x4141414141414141" if self.ropview.sploiter.addr64 else "0x41414141", swidth=10),
                'strComment'      : Form.StringInput(value="padding", swidth=20),
                'iButtonInsert'   : Form.ButtonInput(self.OnButtonInsert, swidth=5),
                'strRopChain'     : Form.MultiLineTextControl(text="", swidth=125, flags = Form.MultiLineTextControl.TXTF_FIXEDFONT),
                'rGroup'          : Form.RadGroupControl(("rFmtPlain","rFmtPython", "rFmtRuby", "rFmtPerl", "rFmtJS") ),
                'iButtonGenerate' : Form.ButtonInput(self.OnButtonGenerate, swidth=5),
                'FormChangeCb'    : Form.FormChangeCb(self.OnFormChange),
            })

        self.Compile()

    def OnFormChange(self, fid):

        # Form initialization
        if fid == -1:
            self.SetFocusedField(self.cEChooser)

            # Add commands to embedded view
            self.ropchainview.cmd_move_up   = self.cEChooser.AddCommand("Move Up", flags = idaapi.CHOOSER_POPUP_MENU | idaapi.CHOOSER_MULTI_SELECTION | idaapi.CHOOSER_HOTKEY, icon=86)
            self.ropchainview.cmd_move_down = self.cEChooser.AddCommand("Move Down", flags = idaapi.CHOOSER_POPUP_MENU | idaapi.CHOOSER_MULTI_SELECTION, icon=85)
            self.ropchainview.cmd_duplicate = self.cEChooser.AddCommand("Duplicate", flags = idaapi.CHOOSER_POPUP_MENU | idaapi.CHOOSER_MULTI_SELECTION, icon=111)
            self.ropchainview.cmd_delete    = self.cEChooser.AddCommand("Delete", flags = idaapi.CHOOSER_POPUP_MENU | idaapi.CHOOSER_MULTI_SELECTION, icon=112)

        # Form OK pressed
        elif fid == -2:
            pass

        return 1

    def OnButtonInsert(self, code=0):
        # Get user input
        address = self.GetControlValue(self.strAddr)
        address = int(address,16)

        comment = self.GetControlValue(self.strComment)

        # Create a dummy gadget
        gadget = Gadget(instructions = [comment,], pivot = 0, operations = [], chg_registers = [], use_registers = [])
        gadget.address = address
        gadget.size = None
        gadget.pivot = None

        # Insert dummy gadget at selected positions
        select_list = self.GetControlValue(self.cEChooser)
        if select_list:
            for i in sorted(select_list,reverse=True):
                self.ropview.gadget_chain.insert(i + 1, gadget)
        else:
            self.ropview.gadget_chain.append(gadget)

        # Update the view
        self.RefreshField(self.cEChooser)

    # TODO: Generate relative addresses to a dynamically determined module base
    #       (e.g. ROP chain built based on a leaked or calculated base address)
    def OnButtonGenerate(self, code=0):

        ropchain = ""

        if self.ropview.sploiter.addr64:
            addr_format = "%016X"
            pack_format_python = "<Q"
            pack_format_ruby = "Q<"
            pack_format_perl = "Q<"
            pack_words_js = 4
            pack_bytes_c = 8

        else:
            addr_format = "%08X"
            pack_format_python = "<I"
            pack_format_ruby = "V"
            pack_format_perl = "V"
            pack_words_js = 2
            pack_bytes_c = 4

        # Plain output
        if self.GetControlValue(self.rGroup) == 0:

            ropchain += "// Generated by IDA Sploiter\n"

            for g in self.ropview.gadget_chain:

                bytes =  ["\\x%02x" % b for b in struct.unpack("B"*pack_bytes_c, struct.pack(pack_format_python, g.address))]

                ropchain += "\"%s\" // %s # %s\n" % (
                                ''.join(bytes),
                                addr_format % g.address,
                                " # ".join(g.instructions))

        # Python
        elif self.GetControlValue(self.rGroup) == 1:

            ropchain += "# Generated by IDA Sploiter\n"

            for i, g in enumerate(self.ropview.gadget_chain):

                ropchain += "rop_chain %s= struct.pack('%s', 0x%s) # %s\n" % (
                                '+' if i > 0 else ' ',
                                pack_format_python,
                                addr_format % g.address,
                                " # ".join(g.instructions) )

        # Ruby
        elif self.GetControlValue(self.rGroup) == 2:

            ropchain += "# Generated by IDA Sploiter\n"

            for i, g in enumerate(self.ropview.gadget_chain):

                ropchain += "rop_chain %s [0x%s].pack('%s') # %s\n" % (
                                '<<' if i > 0 else ' =',
                                addr_format % g.address,
                                pack_format_ruby,
                                " # ".join(g.instructions) )

        # Perl
        elif self.GetControlValue(self.rGroup) == 3:

            ropchain += "# Generated by IDA Sploiter\n"

            for i, g in enumerate(self.ropview.gadget_chain):

                ropchain += "%s$rop_chain %s= pack('%s', 0x%s); # %s\n" % (
                                'my ' if i == 0 else '',
                                '.' if i > 0 else ' ',
                                pack_format_perl,
                                addr_format % g.address,
                                " # ".join(g.instructions) )

        # JavaScript
        elif self.GetControlValue(self.rGroup) == 4:

            ropchain += "// Generated by IDA Sploiter\n"

            for i, g in enumerate(self.ropview.gadget_chain):

                words_js = ["%%u%04x" % w for w in struct.unpack("H"*pack_words_js, struct.pack(pack_format_python, g.address))]

                ropchain += "rop_chain %s= unescape(\"%s\"); // %s # %s\n" % (
                                '+' if i > 0 else ' ',
                                ''.join(words_js),
                                addr_format % g.address,
                                " # ".join(g.instructions) )

        strRopChain = idaapi.textctrl_info_t(text=ropchain, flags = Form.MultiLineTextControl.TXTF_FIXEDFONT)
        self.SetControlValue(self.strRopChain, strRopChain)

###############################################################################
# Function Pointer UI
###############################################################################

###############################################################################
# Function Pointer Form

class PtrForm(Form):

    def __init__(self, sploiter, select_list = None):

        self.sploiter = sploiter
        self.select_list = select_list
        self.mod = ModuleView(self.sploiter, embedded=True)

        Form.__init__(self, 
r"""BUTTON YES* Search
Search writable function pointers

{FormChangeCb}<Modules:{cEChooser}>

Pointer Charset:                  Search Settings:
<nonull:{cPtrNonull}>        <Bad chars     :{strBadChars}>
<unicode:{cPtrUnicode}>       Unicode Table <ANSI:{rUnicodeANSI}><OEM:{rUnicodeOEM}><UTF7:{rUnicodeUTF7}><UTF8:{rUnicodeUTF8}>{radUnicode}>
<ascii:{cPtrAscii}>         <PTR Offset    :{intPtrOffset}> 
<asciiprint:{cPtrAsciiPrint}>    <P2P Offset    :{intP2POffset}>
<alphanum:{cPtrAlphaNum}>      <Max Pointers  :{intMaxPtrs}>
<alpha:{cPtrAlpha}>         Filter        <Function Pointers:{rFilterPtr}><Pointers-to-Pointers:{rFilterP2P}>{rGroup}>
<numeric:{cPtrNum}>{ptrGroup}>       Other settings<Do not search for pointers-to-pointers:{rSearchPtrOnly}>{rSearchGroup}>


""", {
                'cEChooser'       : Form.EmbeddedChooserControl(self.mod, swidth=131),
                'ptrGroup'        : Form.ChkGroupControl(("cPtrNonull", "cPtrAscii", "cPtrAsciiPrint", "cPtrUnicode",'cPtrAlphaNum','cPtrAlpha','cPtrNum')),
                'rGroup'          : Form.RadGroupControl(("rFilterPtr","rFilterP2P") ),
                'rSearchGroup'    : Form.ChkGroupControl(("rSearchPtrOnly",) ),
                'strBadChars'     : Form.StringInput(swidth=70,tp=Form.FT_ASCII),
                'radUnicode'      : Form.RadGroupControl(("rUnicodeANSI","rUnicodeOEM","rUnicodeUTF7","rUnicodeUTF8")),
                'intPtrOffset'    : Form.StringInput(tp=Form.FT_ASCII,value="0x0", swidth=10),
                'intP2POffset'    : Form.StringInput(tp=Form.FT_ASCII,value="0x0", swidth=10),
                'intMaxPtrs'      : Form.NumericInput(swidth=4,tp=Form.FT_DEC,value=0),
                'FormChangeCb'    : Form.FormChangeCb(self.OnFormChange),
            })

        self.Compile()

    def OnFormChange(self, fid):

        # Form initialization
        if fid == -1:
            self.SetFocusedField(self.cEChooser)

            # Preselect non-ASLR modules on startup if none were already specified
            if self.select_list == None:
                
                self.select_list = list()
                for i, m in enumerate(self.sploiter.modules):
                    if m.ASLR == "No":
                        self.select_list.append(i)

            self.SetControlValue(self.cEChooser, self.select_list)

        # Form OK pressed
        elif fid == -2:
            pass

        return 1

###############################################################################
# Function Pointer Viewer

class PtrView(Choose2):
    """
    Chooser class to display writable function pointers.
    """
    def __init__(self, sploiter):

        self.sploiter = sploiter

        Choose2.__init__(self,
                         "Writeable function pointers",
                         [ ["Pointer",            13 | Choose2.CHCOL_HEX], 
                           ["Offset",              4 | Choose2.CHCOL_HEX], 
                           ["Name",               10 | Choose2.CHCOL_PLAIN],
                           ["Module",             10 | Choose2.CHCOL_PLAIN],
                           ["Charset",            12 | Choose2.CHCOL_PLAIN],
                           ["Caller Address",     13 | Choose2.CHCOL_HEX], 
                           ["Caller Name",        10 | Choose2.CHCOL_PLAIN],
                           ["Instruction",        15 | Choose2.CHCOL_PLAIN],
                           ["Ptr-to-Ptr",         13 | Choose2.CHCOL_PLAIN],
                           ["Ptr-to-Ptr Offset",   4 | Choose2.CHCOL_HEX], 
                           ["Ptr-to-Ptr Charset", 12 | Choose2.CHCOL_PLAIN],
                         ],
                         flags = Choose2.CH_MULTI_EDIT)

        self.icon = 143

        # Items for display and corresponding data
        # NOTE: Could become desynchronized, so to avoid this
        #       refresh the view after each change.
        self.items = []

        # Initialize/Refresh the view
        self.refreshitems()

        # Selected items
        self.select_caller_breakpoint_list = list()

        # Command callbacks
        self.cmd_select_caller_breakpoint = None
        self.cmd_jump_ptr = None
        self.cmd_jump_p2p = None
        self.cmd_jump_caller = None
        self.cmd_export_csv  = None

    def show(self):
        # Attempt to open the view
        if self.Show() < 0: return False

        # Add extra context menu commands
        # NOTE: Make sure you check for duplicates
        if self.cmd_select_caller_breakpoint == None:
            self.cmd_select_caller_breakpoint = self.AddCommand("Add breakpoint", flags = idaapi.CHOOSER_POPUP_MENU | idaapi.CHOOSER_MULTI_SELECTION, icon=120)

        if self.cmd_jump_ptr == None:
            self.cmd_jump_ptr = self.AddCommand("Jump to pointer", flags = idaapi.CHOOSER_POPUP_MENU | idaapi.CHOOSER_MULTI_SELECTION, icon=141)

        if self.cmd_jump_p2p == None:
            self.cmd_jump_p2p = self.AddCommand("Jump to pointer-to-pointer", flags = idaapi.CHOOSER_POPUP_MENU | idaapi.CHOOSER_MULTI_SELECTION, icon=142)

        if self.cmd_jump_caller == None:
            self.cmd_jump_caller = self.AddCommand("Jump to caller", flags = idaapi.CHOOSER_POPUP_MENU | idaapi.CHOOSER_MULTI_SELECTION, icon=143)

        if self.cmd_export_csv == None:
            self.cmd_export_csv = self.AddCommand("Export as csv...", flags = idaapi.CHOOSER_POPUP_MENU, icon=40)

        return True

    def refreshitems(self):
        self.items = []

        for ptr in self.sploiter.funcptr.ptrs:

            self.items.append([ self.sploiter.addr_format % ptr.ptr_ea, 
                                hex(ptr.ptr_offset),
                                ptr.name, 
                                ptr.module, 
                                ", ".join(ptr.ptr_charset),
                                self.sploiter.addr_format % ptr.call_ea,
                                ptr.call_name,
                                ptr.insn_disas,      
                                self.sploiter.addr_format % ptr.p2p_ea if ptr.p2p_ea else "",
                                hex(ptr.p2p_offset) if ptr.p2p_ea else "",
                                ", ".join(ptr.p2p_charset) if ptr.p2p_ea else "",                          
                                ])

    def OnCommand(self, n, cmd_id):

        if cmd_id == self.cmd_select_caller_breakpoint:

            # Empty selection
            if n == -1:
                pass                

            # Selection start
            elif n == -2:
                self.select_caller_breakpoint_list = list()

            # Selection end
            elif n == -3:               

                # Unique set of caller addresses
                eas = set()

                for i in self.select_caller_breakpoint_list:
                    eas.add( self.sploiter.funcptr.ptrs[i].call_ea )

                print "[idasploiter] Settings %d breakpoints on selected caller addresses." % len(eas)

                for ea in eas:
                    # Set default software breakpoint
                    idaapi.add_bpt(ea, 0, idaapi.BPT_DEFAULT)

                self.select_caller_breakpoint_list = list()

            # Selection number
            else:
                self.select_caller_breakpoint_list.append(n)

        # Jump to Pointer address
        elif cmd_id == self.cmd_jump_ptr:
            if n >= 0 and self.sploiter.funcptr.ptrs[n].ptr_ea:
                idaapi.jumpto( self.sploiter.funcptr.ptrs[n].ptr_ea )

        # Jump to Pointer-to-Pointer address
        elif cmd_id == self.cmd_jump_p2p:
            if n >= 0 and self.sploiter.funcptr.ptrs[n].p2p_ea:
                idaapi.jumpto( self.sploiter.funcptr.ptrs[n].p2p_ea )

        # Jump to Caller address
        elif cmd_id == self.cmd_jump_caller:
            if n >= 0 and self.sploiter.funcptr.ptrs[n].call_ea:
                idaapi.jumpto( self.sploiter.funcptr.ptrs[n].call_ea )

        # Export CSV
        elif cmd_id == self.cmd_export_csv:

            file_name = idaapi.askfile_c(1, "*.csv", "Please enter CSV file name")
            if file_name:
                print "[idasploiter] Exporting function pointers to %s" % file_name
                with open(file_name, 'wb') as csvfile:
                    csvwriter = csv.writer(csvfile, delimiter=',',
                                            quotechar='"', quoting=csv.QUOTE_MINIMAL)
                    csvwriter.writerow(["Pointer","Offset","Name","Module","Charset","Caller Address","Caller Name","Instruction","Ptr-to-Ptr","Ptr-to-Ptr Offset","Ptr-to-Ptr Charset"])
                    for item in self.items:
                        csvwriter.writerow(item)
        return 1

 
    def OnSelectLine(self, n):
        idaapi.jumpto( int(self.items[n][0],16) )

    def OnGetLine(self, n):
        return self.items[n]

    def OnClose(self):
        self.cmd_select_caller_breakpoint = None
        self.cmd_jump_ptr = None
        self.cmd_jump_p2p = None
        self.cmd_jump_caller = None
        self.cmd_export_csv  = None

    def OnGetSize(self):
        return len(self.items)

    def OnRefresh(self, n):
        self.refreshitems()
        return n

    def OnActivate(self):
        self.refreshitems()

###############################################################################
# Pattern UI
###############################################################################

###############################################################################
# Pattern Create Form

class PatternCreateForm(Form):

    def __init__(self, sploiter):

        self.sploiter = sploiter
        self.pattern_complete = ""

        Form.__init__(self, 
r"""BUTTON YES NONE
BUTTON CANCEL NONE
Create pattern
{FormChangeCb}<:{strPattern}>
<Size   :{intSize}>
Formats<Regular:{rFmtReg}><Hex:{rFmtHex}><JavaScript:{rFmtJS}>{rGroup}>
Charset<C4:{cC4}><##:{strC4}>
       <C3:{cC3}><##:{strC3}>
       <C2:{cC2}><##:{strC2}>
       <C1:{cC1}>{cGroup}><##:{strC1}>

{spacer}<Generate:{iButtonGenerate}>
""", {
            'spacer'          : Form.StringLabel("            "),
            'intSize'         : Form.NumericInput(swidth=30,tp=Form.FT_DEC,value=1000),
            'iButtonGenerate' : Form.ButtonInput(self.OnButtonGenerate,swidth=5),
            'rGroup'          : Form.RadGroupControl(("rFmtReg","rFmtHex", "rFmtJS") ),
            'cGroup'          : Form.ChkGroupControl(("cC4","cC3","cC2","cC1")),
            'strC4'           : Form.StringInput(swidth=38,tp=Form.FT_ASCII),
            'strC3'           : Form.StringInput(swidth=38,tp=Form.FT_ASCII),
            'strC2'           : Form.StringInput(swidth=38,tp=Form.FT_ASCII),
            'strC1'           : Form.StringInput(swidth=38,tp=Form.FT_ASCII),
            'strPattern'      : Form.MultiLineTextControl(text="", width=400, flags = Form.MultiLineTextControl.TXTF_FIXEDFONT | Form.MultiLineTextControl.TXTF_SELECTED | Form.MultiLineTextControl.TXTF_READONLY),
            'FormChangeCb'    : Form.FormChangeCb(self.OnFormChange),
            })

        self.Compile()

    def OnFormChange(self, fid):

        # Form initialization
        if fid == -1:

            # Set initial checkboxes for the classic metasploit 3-byte pattern
            self.SetControlValue(self.cC4, False)
            self.SetControlValue(self.cC3, True)
            self.SetControlValue(self.cC2, True)
            self.SetControlValue(self.cC1, True)

            # Set initial charsets values
            self.SetControlValue(self.strC4, self.sploiter.c4_list)
            self.SetControlValue(self.strC3, self.sploiter.c3_list)
            self.SetControlValue(self.strC2, self.sploiter.c2_list)
            self.SetControlValue(self.strC1, self.sploiter.c1_list)

        # Form OK pressed
        elif fid == -2:
            pass

        return 1

    def OnButtonGenerate(self, code=0):

        size = self.GetControlValue(self.intSize) 
        print "[idasploiter] Creating a pattern of size %d." % size

        # Update global charset lists
        self.sploiter.c4_list = self.GetControlValue(self.strC4)
        self.sploiter.c3_list = self.GetControlValue(self.strC3)
        self.sploiter.c2_list = self.GetControlValue(self.strC2)
        self.sploiter.c1_list = self.GetControlValue(self.strC1)

        # Generate pattern only using selected lists
        self.update_complete_pattern(self.GetControlValue(self.cC4), self.GetControlValue(self.cC3), self.GetControlValue(self.cC2), self.GetControlValue(self.cC1))

        pattern = self.pattern_create(size) 

        if pattern:

            # Hex output
            if self.GetControlValue(self.rGroup) == 1:
                pattern = binascii.hexlify(pattern)

            # JavaScript Unicode output
            elif self.GetControlValue(self.rGroup) == 2:
                
                pattern_js = ""

                # Number of WORDs in the pattern
                num_words = len(pattern) / 2

                # Chop off extra characters
                num_extra = len(pattern) % 2
                if num_extra:
                    pattern = pattern[:-num_extra]

                # Unpack individual WORDs from the pattern
                words = struct.unpack("H"*num_words, pattern)
                for word in words:
                    pattern_js += "%%u%s" % hex(word)

                pattern = pattern_js

            strPattern = idaapi.textctrl_info_t(text=pattern, flags = Form.MultiLineTextControl.TXTF_FIXEDFONT | Form.MultiLineTextControl.TXTF_SELECTED | Form.MultiLineTextControl.TXTF_READONLY)
            self.SetControlValue(self.strPattern, strPattern )


    def pattern_create(self, size):

        # Make sure at least one custom charset was specified
        if not len(self.pattern_complete):
            idaapi.warning("Cannot generate an empty pattern.")
            return False

        # Requested pattern size exceeds maximum allowed pattern size for the selected characterset
        elif size > len(self.pattern_complete):
            idaapi.warning("Requested size %d exceeds maximum unique pattern size %d." % (size, len(self.pattern_complete)))
            return False

        # Get a size substring from the complete pattern
        else:
            return self.pattern_complete[:size]

    def update_complete_pattern(self, c4, c3, c2, c1):

        c4_list = self.sploiter.c4_list if c4 else ["",]
        c3_list = self.sploiter.c3_list if c3 else ["",]
        c2_list = self.sploiter.c2_list if c2 else ["",]
        c1_list = self.sploiter.c1_list if c1 else ["",]

        # Generate complete pattern
        # NOTE: Not as pretty as calculating offset based on positions, but
        #       allows custom charactersets with varying positions
        self.pattern_complete = ""

        # Produce a cartesian product of all character lists and make a single string by flattening the list of lists
        self.pattern_complete = ''.join([i for j in itertools.product(c4_list, c3_list, c2_list, c1_list) for i in j])


###############################################################################
# Pattern Offset Form

class PatternDetectForm(Form):

    def __init__(self, sploiter, debugger = True):

        self.sploiter = sploiter
        self.debugger = debugger

        self.pattern_complete = ""

        form_dict = {
            'spacer'         : Form.StringLabel("           "),
            'strPattern'      : Form.StringInput(tp=Form.FT_ASCII, swidth=17),
            'strOffset'       : Form.StringInput(tp=Form.FT_ASCII, swidth=13),
            'iButtonDetect'   : Form.ButtonInput(self.OnButtonDetect, swidth=5),
            'rGroup'          : Form.RadGroupControl(('rFmtAuto', 'rFmtStr', 'rFmtAddr')),
            'cGroup'          : Form.ChkGroupControl(("cC4","cC3","cC2","cC1")),
            'strC4'           : Form.StringInput(swidth=38,tp=Form.FT_ASCII),
            'strC3'           : Form.StringInput(swidth=38,tp=Form.FT_ASCII),
            'strC2'           : Form.StringInput(swidth=38,tp=Form.FT_ASCII),
            'strC1'           : Form.StringInput(swidth=38,tp=Form.FT_ASCII),
            'FormChangeCb'    : Form.FormChangeCb(self.OnFormChange),
            }

        reg_form = ""
        self.reg_ctrls = list()

        # Display register values when executed with a debugger running
        if debugger:

            # Select general purpose registers for a given architecture
            if self.sploiter.addr64:
                self.reg_list = ["RAX", "RBX", "RCX", "RDX", "RSP", "RBP", "RSI", "RDI", "RIP", "R8", "R9", "R10", "R11", "R12", "R13", "R14", "R15"]
            else:
                self.reg_list = ["EAX", "EBX", "ECX", "EDX", "ESP", "EBP", "ESI", "EDI", "EIP"]

            # Populate register values and store form controls in a list
            for reg_name in self.reg_list:
                reg_value = idc.GetRegValue(reg_name)
                reg_value = self.sploiter.addr_format % reg_value

                ctrlRegValue = Form.StringInput(tp=Form.FT_ASCII, value=reg_value,swidth=17)
                form_dict["str%s" % reg_name] = ctrlRegValue

                ctrlRegOffset = Form.StringInput(tp=Form.FT_ASCII, swidth=13)
                form_dict["str%sOffset" % reg_name] = ctrlRegOffset

                reg_form += "  <%(reg)03s  :{str%(reg)s}><Offset:{str%(reg)sOffset}>\n" % {'reg': reg_name}

                self.reg_ctrls.append( (ctrlRegValue,  ctrlRegOffset) )

        Form.__init__(self, 
r"""BUTTON YES NONE
BUTTON CANCEL NONE
Detect pattern
{FormChangeCb}%s<Pattern:{strPattern}><Offset:{strOffset}>

Formats<Auto:{rFmtAuto}><Address:{rFmtAddr}><String:{rFmtStr}>{rGroup}>
Charset<C4:{cC4}><:{strC4}>
       <C3:{cC3}><:{strC3}>
       <C2:{cC2}><:{strC2}>
       <C1:{cC1}>{cGroup}><:{strC1}>

{spacer}<Detect:{iButtonDetect}>
""" % reg_form, form_dict)

        self.Compile()

    def OnFormChange(self, fid):

        # Form initialization
        if fid == -1:

            # Set initial checkboxes for the classic metasploit 3-byte pattern
            self.SetControlValue(self.cC4, False)
            self.SetControlValue(self.cC3, True)
            self.SetControlValue(self.cC2, True)
            self.SetControlValue(self.cC1, True)

            # Set initial charsets values
            self.SetControlValue(self.strC4, self.sploiter.c4_list)
            self.SetControlValue(self.strC3, self.sploiter.c3_list)
            self.SetControlValue(self.strC2, self.sploiter.c2_list)
            self.SetControlValue(self.strC1, self.sploiter.c1_list)

            # NOTE: Trying to detect offset for registry values during
            #       form initialization appears to be very slow and does
            #       not provide an option for a user to specify a custom 
            #       characterset.

        # Form OK pressed
        elif fid == -2:
            pass

        return 1

    def OnButtonDetect(self, code=0):

        # Update global charset lists
        self.sploiter.c4_list = self.GetControlValue(self.strC4)
        self.sploiter.c3_list = self.GetControlValue(self.strC3)
        self.sploiter.c2_list = self.GetControlValue(self.strC2)
        self.sploiter.c1_list = self.GetControlValue(self.strC1)

        # Generate pattern only using selected lists
        self.update_complete_pattern(self.GetControlValue(self.cC4), self.GetControlValue(self.cC3), self.GetControlValue(self.cC2), self.GetControlValue(self.cC1))

        for ctrlRegValue, ctrlRegOffset in self.reg_ctrls:
            reg_value = self.GetControlValue(ctrlRegValue)
            self.SetControlValue(ctrlRegOffset, self.verbalize_pattern(reg_value, address=True) )

        pattern = self.GetControlValue(self.strPattern)

        # Force address format
        if self.GetControlValue(self.rFmtAddr):
            self.SetControlValue(self.strOffset, self.verbalize_pattern(pattern, address=True) )

        # Force string format
        elif self.GetControlValue(self.rFmtStr):
            self.SetControlValue(self.strOffset, self.verbalize_pattern(pattern, string=True) )

        # Auto format
        else:
            self.SetControlValue(self.strOffset, self.verbalize_pattern(pattern) )

    def verbalize_pattern(self, pattern, address=False, string=False):

        # Forced address format
        if address:

            # Cleanup the pattern
            pattern = pattern.replace(' ','')       # remove spaces
            pattern = pattern.replace('\\x','')     # remove '\x' prefixes
            pattern = pattern.replace('0x','')      # remove '0x' prefixes

            if self.sploiter.addr64 and len(pattern) == 8*2:
                try:
                    pattern = struct.pack("Q", int(pattern, 16))
                except Exception, e:
                    return "Invalid address"

            elif not self.sploiter.addr64 and len(pattern) == 4*2:
                try:
                    pattern = struct.pack("I", int(pattern, 16))
                except Exception, e:
                    return "Invalid address"

            else:
                return "Invalid address"

        # Forced string format
        elif string:
            pass

        # Autodetect format
        else:

            # Cleanup the pattern
            pattern = pattern.replace(' ','')       # remove spaces
            pattern = pattern.replace('\\x','')     # remove '\x' prefixes
            pattern = pattern.replace('0x','')      # remove '0x' prefixes

            # Attempt to detect and unpack pattern
            if self.sploiter.addr64 and len(pattern) == 8*2:
                try:
                    pattern = struct.pack("Q", int(pattern, 16))
                except Exception, e:
                    pass

            elif not self.sploiter.addr64 and len(pattern) == 4*2:
                try:
                    pattern = struct.pack("I", int(pattern, 16))
                except Exception, e:
                    pass

        if len(pattern) > 0:
            offset         = self.pattern_offset(pattern)
            offset_reverse = self.pattern_offset(pattern[::-1])
        else:
            return "Blank pattern"

        if offset >= 0:
            return "%d" % offset
        elif offset_reverse >= 0:
            return "%d (reverse)" % offset_reverse
        else:
            return "Not detected"

    def pattern_offset(self, pattern):

        offset = -1

        # Make sure at least one custom charset was specified
        if not len(self.pattern_complete):
            idaapi.warning("Cannot detect an empty pattern.")
            return offset

        # Locate the pattern
        if pattern in self.pattern_complete:
            offset = self.pattern_complete.index(pattern)

        return offset

    def update_complete_pattern(self, c4, c3, c2, c1):

        c4_list = self.sploiter.c4_list if c4 else ["",]
        c3_list = self.sploiter.c3_list if c3 else ["",]
        c2_list = self.sploiter.c2_list if c2 else ["",]
        c1_list = self.sploiter.c1_list if c1 else ["",]

        # Generate complete pattern
        # NOTE: Not as pretty as calculating offset based on positions, but
        #       allows custom character sets with varying positions
        self.pattern_complete = ""

        # Produce a cartesian product of all character lists and make a single string by flattening the list of lists
        self.pattern_complete = ''.join([i for j in itertools.product(c4_list, c3_list, c2_list, c1_list) for i in j])

###############################################################################
# Compare UI
###############################################################################

###############################################################################
# Compare Form

class CompareForm(Form):

    def __init__(self, sploiter, debugger = True):

        self.sploiter = sploiter
        self.debugger = debugger

        Form.__init__(self, 
r"""BUTTON YES NONE
BUTTON CANCEL NONE
Compare file to memory
{FormChangeCb}<File     :{impFile}>
<Memory   :{intAddr}>

<:{strCompare}>

<Bad Chars:{strBadChars}>
<Mem Holes:{strMemHoles}>

{spacer}<Compare:{iButtonCompare}>
""", {
            'spacer'        : Form.StringLabel("                                "),
            'impFile'        : Form.FileInput(swidth=62, open=True),
            'intAddr'        : Form.NumericInput(swidth=62, tp=Form.FT_ADDR, value=idaapi.get_screen_ea()),
            'strCompare'     : Form.MultiLineTextControl(text="", swidth=72, flags = Form.MultiLineTextControl.TXTF_FIXEDFONT | Form.MultiLineTextControl.TXTF_READONLY),
            'strBadChars'    : Form.StringInput(swidth=62,tp=Form.FT_ASCII),
            'strMemHoles'    : Form.StringInput(swidth=62,tp=Form.FT_ASCII),
            'iButtonCompare' : Form.ButtonInput(self.OnButtonCompare,),
            'FormChangeCb'   : Form.FormChangeCb(self.OnFormChange),
            })

        self.Compile()

    def OnFormChange(self, fid):

        # Form initialization
        if fid == -1:
            pass

        # Form OK pressed
        elif fid == -2:
            pass

        return 1

    def OnButtonCompare(self, code=0):
        
        #######################################################################
        # Read file into buffer
        file_name = self.GetControlValue(self.impFile)

        # Verify file name was provided
        if not len(file_name):
            idaapi.warning("Please provide file name.")
            return

        # Attempt to read the file
        print "[idasploiter] Reading file %s" % file_name
        
        try:
            f = open(file_name,'rb')
        except Exception, e:
            idaapi.warning("File I/O error({0}): {1}.".format(e.errno, e.strerror) )
            return
        else:
            file_buf = f.read()
            f.close()

        # Check file content
        if not len(file_buf) > 0:
            idaapi.warning("Provided file was blank.")
            return

        # 64k ought to be enough for anybody =)
        elif len(file_buf) > 0xffff:
            idaapi.warning("File is too big.")
            return

        # Covert file bytes into printable hex view
        file_buf_str = " ".join(["%02x" % ord(b) for b in file_buf])
        file_buf_list = textwrap.wrap(file_buf_str, 48)

        #######################################################################
        # Read memory into buffer and mark non-matching bytes
        address = self.GetControlValue(self.intAddr)

        print "[idasploiter] Reading %d bytes from %s." % ( len(file_buf), self.sploiter.addr_format % address)

        if self.debugger:
            mem_buf = idaapi.dbg_read_memory(address, len(file_buf))
        else:
            mem_buf = idaapi.get_many_bytes(address, len(file_buf))

        if not mem_buf:
            idaapi.warning("Could not read memory.")
            return

        # Convert memory bytes into printable hex view
        mem_buf_str = ""

        # Badchars tracking
        badchars = list()
        badchars_nonseq = set()
        badchars_hole = list()

        # Locate all of the non-sequential badchars. The assumption is that
        # after the first badchar (e.g. 0x0A) the rest may be corrupted.
        for i, b in enumerate(file_buf):
            if file_buf[i] == mem_buf[i]:
                mem_buf_str+= ".. "

            else:
                mem_buf_str+= "%02x " % ord( mem_buf[i])

            # Do differential processing if the two bytes are not the same or
            # if the matching byte is in the bad characters list

            if file_buf[i] != mem_buf[i] or "%02x" % ord(file_buf[i]) in badchars_nonseq:

                ###############################################################
                # Memory holes

                # Check if the badchar appears anywhere else in memory to
                # detect positional holes. If this is the first discrepancy
                # then it is likely it is a badchar (e.g. 0x00, 0x0a, 0x0d)

                if len(badchars_hole):
                    hole_i, hole_size = badchars_hole[-1]

                    # Continuation of a previously discovered hole
                    if i == hole_i + hole_size:
                        badchars_hole[-1] = (hole_i, hole_size + 1)

                    # New hole discovered
                    else:
                        badchars_hole.append( (i,1) )

                # New hole discovered
                else:
                    badchars_hole.append( (i,1) )

                ###############################################################
                # Bad characters

                # Check if previous byte in the file is already a badchar, so
                # this new discrepancy doesn't matter:

                if not "%02x" % ord(file_buf[i-1]) in  badchars_nonseq:

                    # Skip sequential corruption
                    if len(badchars_nonseq):
                        bad_i, bad_b = badchars[-1]

                        # Check if the previous byte was also corrupted
                        # indicating a sequential corruption
                        if i - 1 != bad_i:
                            badchars_nonseq.add( "%02x" % ord(b) )

                    else:                      
                        badchars_nonseq.add( "%02x" % ord(b) )

                # Store badchars in a list
                badchars.append( (i,b) )


        mem_buf_list = textwrap.wrap(mem_buf_str, 48)

        if len(badchars) == 0:
            self.SetControlValue(self.strBadChars, "Identical.")  
        elif len(badchars) == len(file_buf):
            self.SetControlValue(self.strBadChars, "Completely different.")          
        else:
            self.SetControlValue(self.strBadChars, " ".join( sorted(badchars_nonseq) ) )
            self.SetControlValue(self.strMemHoles, ", ".join( [ "0x%02X: %d byte(s)" % h for h in badchars_hole ]  ) )

        #######################################################################
        # Populate the text box

        buf_formatted = "      00 01 02 03 04 05 06 07  08 09 0A 0B 0C 0D 0E 0F\n\n"
        for i in range( len(file_buf_list) ):

            file_line = file_buf_list[i].ljust(48)
            file_line = file_line[:23]+" "+file_line[23:]

            mem_line = mem_buf_list[i].ljust(48)
            mem_line = mem_line[:23]+" "+mem_line[23:]

            buf_formatted += "%04X  %s  F\n      %s  M\n\n" % (i*0x10, file_line, mem_line)

        strCompare = idaapi.textctrl_info_t(text=buf_formatted, flags = Form.MultiLineTextControl.TXTF_FIXEDFONT | Form.MultiLineTextControl.TXTF_READONLY)
        self.SetControlValue(self.strCompare, strCompare)

###############################################################################
# Plugin Manager
###############################################################################

class SploitManager():
    """ Class that manages GUI forms and exploitation methods of the plugin. """
    
    def __init__(self): 
        self.addmenu_item_ctxs = list()
        self.sploiter = Sploiter()

    ###########################################################################
    # Menu Items
    def add_menu_item_helper(self, menupath, name, hotkey, flags, pyfunc, args):

        # add menu item and report on errors
        addmenu_item_ctx = idaapi.add_menu_item(menupath, name, hotkey, flags, pyfunc, args)
        if addmenu_item_ctx is None:
            return 1
        else:
            self.addmenu_item_ctxs.append(addmenu_item_ctx)
            return 0

    def add_menu_items(self):
        if self.add_menu_item_helper("View/Open subviews/Segments", "Modules", "Shift+F6", 0, self.show_modules_view, None): return 1

        if self.add_menu_item_helper("Search/all error operands", "function pointers...", "Alt+f", 1, self.show_funcptr_view, None): return 1
        if self.add_menu_item_helper("Search/all error operands", "gadgets...", "Alt+r", 1, self.show_rop_view, None): return 1

        if self.add_menu_item_helper("Edit/Begin selection", "Create pattern...", "Shift+c", 0, self.show_pattern_create, None): return 1
        if self.add_menu_item_helper("Edit/Begin selection", "Detect pattern...", "Shift+d", 0, self.show_pattern_detect, None): return 1
        if self.add_menu_item_helper("Edit/Begin selection", "Compare file to memory...", "Shift+f", 0, self.show_compare, None): return 1
        
        return 0

    def del_menu_items(self):
        for addmenu_item_ctx in self.addmenu_item_ctxs:
            idaapi.del_menu_item(addmenu_item_ctx)

    ###########################################################################
    # Utility Functions

    # Check debugger is ready for querying
    def check_debugger(self):
        if idaapi.dbg_can_query() and idaapi.get_process_state() < 0:
            return True
        else:
            return False

    ###########################################################################
    # View Callbacks
    
    # Modules View
    def show_modules_view(self):
        if self.check_debugger():
            self.sploiter.process_modules()
            self.sploiter.show_modules_view()
        else:            
            idaapi.warning("Please run and suspend the debugger.")
            

    # ROP View
    def show_rop_view(self):
        if self.check_debugger():
            self.sploiter.process_modules()
            self.sploiter.process_rop()
        else:            
            idaapi.warning("Please run and suspend the debugger.")
            

    # Function Pointer View
    def show_funcptr_view(self):     
        if self.check_debugger():   
            self.sploiter.process_modules()
            self.sploiter.process_funcptr()
        else:            
            idaapi.warning("Please run and suspend the debugger.")

    # Create Pattern Form
    def show_pattern_create(self):
        self.sploiter.pattern_create()

    # Detect Pattern Form
    def show_pattern_detect(self):

        # Debugger Version
        if self.check_debugger():
            self.sploiter.pattern_detect()

        # Static Version
        else:
            self.sploiter.pattern_detect(debugger=False)

    # Compare File to Memory Form
    def show_compare(self):

        # Debugger Version
        if self.check_debugger():
            self.sploiter.process_compare()

        # Static Version
        else:
            self.sploiter.process_compare(debugger=False)


###############################################################################

class idasploiter_t(plugin_t):

    flags = idaapi.PLUGIN_UNL
    comment = "Exploit development and vulnerability research toolkit."
    help = "Exploit development and vulnerability research toolkit."
    wanted_name = "IDA Sploiter"
    wanted_hotkey = ""

    def init(self):  

        # Only Intel x86/x86-64 are supported
        if idaapi.ph_get_id() == idaapi.PLFM_386:

            global idasploiter_manager

            # Check if already initialized
            if not 'idasploiter_manager' in globals():

                idasploiter_manager = SploitManager()
                if idasploiter_manager.add_menu_items():
                    print "Failed to initialize IDA Sploiter."
                    idasploiter_manager.del_menu_items()
                    del idasploiter_manager
                    return idaapi.PLUGIN_SKIP
                else:
                    print("Initialized IDA Sploiter v%s (c) Peter Kacherginsky <iphelix@thesprawl.org>" % IDASPLOITER_VERSION)

            return idaapi.PLUGIN_KEEP
        else:
            return idaapi.PLUGIN_SKIP

    def run(self, arg):
        pass

    def term(self):
        pass        

def PLUGIN_ENTRY():
    return idasploiter_t()

###############################################################################
# Script / Testing
###############################################################################

def idasploiter_main():
    global idasploiter_manager

    if 'idasploiter_manager' in globals():
        idasploiter_manager.del_menu_items()
        del idasploiter_manager

    idasploiter_manager = SploitManager()
    idasploiter_manager.add_menu_items()

    sploiter = idasploiter_manager.sploiter

if __name__ == '__main__':
    #idasploiter_main()
    pass


repository: zhkl0228/d810
file: D810.py
url: https://github.com/zhkl0228/d810/blob/680cc531b4eb91f14b2630268a571e7b68550931/D810.py
wanted_name: D-810
comment: Interface to the D-810 plugin
created_at: 2021-09-26 14:29:11+00:00
pushed_at: 2021-10-05 09:35:45+00:00
description:
D-810 is an IDA Pro plugin which can be used to deobfuscate code at decompilation time by modifying IDA Pro microcode.

content:
import os
import idaapi
import ida_hexrays
import ida_kernwin


from d810.conf import D810Configuration
from d810.manager import D810State, D810_LOG_DIR_NAME
from d810.log import configure_loggers, clear_logs


D810_VERSION = "0.1"

class D810Plugin(idaapi.plugin_t):
    # variables required by IDA
    flags = 0  # normal plugin
    wanted_name = "D-810"
    wanted_hotkey = "Ctrl-Shift-D"
    comment = "Interface to the D-810 plugin"
    help = ""
    initialized = False

    def __init__(self):
        super(D810Plugin, self).__init__()
        self.d810_config = None
        self.state = None
        self.initialized = False


    def reload_plugin(self):
        if self.initialized:
            self.term()

        self.d810_config = D810Configuration()

        #TO-DO: if [...].get raises an exception because log_dir is not found, handle exception
        real_log_dir = os.path.join(self.d810_config.get("log_dir"), D810_LOG_DIR_NAME)

        #TO-DO: if [...].get raises an exception because erase_logs_on_reload is not found, handle exception
        if self.d810_config.get("erase_logs_on_reload"):
            clear_logs(real_log_dir)

        configure_loggers(real_log_dir)
        self.state = D810State(self.d810_config)
        print("D-810 reloading...")
        self.state.start_plugin()
        self.initialized = True


    # IDA API methods: init, run, term
    def init(self):
        if not ida_hexrays.init_hexrays_plugin():
            print("D-810 need Hex-Rays decompiler. Skipping")
            return idaapi.PLUGIN_SKIP

        kv = ida_kernwin.get_kernel_version().split(".")
        if (int(kv[0]) < 7) or (int(kv[1]) < 5):
            print("D-810 need IDA version >= 7.5. Skipping")
            return idaapi.PLUGIN_SKIP
        print("D-810 initialized (version {0})".format(D810_VERSION))
        return idaapi.PLUGIN_OK


    def run(self, args):
        self.reload_plugin()


    def term(self):
        print("Terminating D-810...")
        if self.state is not None:
            self.state.stop_plugin()

        self.initialized = False


def PLUGIN_ENTRY():
    return D810Plugin()


repository: Jamie793/IDA-Plugin-D810-Obfs
file: D810.py
wanted_name: D-810
comment: Interface to the D-810 plugin
created_at: 2022-08-02 01:44:05+00:00
pushed_at: 2022-08-02 01:44:15+00:00
url:
https://github.com/Jamie793/IDA-Plugin-D810-Obfs/blob/f989ab61ad50b6150fea8b6597edd644ad6d774a/D810.py


repository: yashchaurasia667/Binary-Exploitation
file: Tools/ghidra_11.3.1_PUBLIC/Extensions/IDAPro/Python/6xx/plugins/xmlldr.py
wanted_name: XML Importer
comment: Import XML PROGRAM file
created_at: 2023-11-11 13:49:18+00:00
pushed_at: 2025-04-26 07:50:25+00:00
url:
https://github.com/yashchaurasia667/Binary-Exploitation/blob/bf9e64a8a9d630289d3fae68cd2f7cd6f88648eb/Tools/ghidra_11.3.1_PUBLIC/Extensions/IDAPro/Python/6xx/plugins/xmlldr.py

content:
## ###
#  IP: GHIDRA
# 
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#  
#       http://www.apache.org/licenses/LICENSE-2.0
#  
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
##
#---------------------------------------------------------------------
# xmlldr.py - IDA XML Import loader and plugin
#---------------------------------------------------------------------
"""
Loader and plugin for IDA to import a XML PROGRAM file to a database.
"""

import idaapi
import sys
import time
from xml.etree import cElementTree


XML_IMPORTER_VERSION = "2.1.1"
BASELINE_IDA_VERSION = 620
IDA_SDK_VERSION = idaapi.IDA_SDK_VERSION
BADADDR = idaapi.BADADDR


"""
Loader functions
"""
def accept_file(li, n):
    """
    Check if the file is of supported format

    @param li: a file-like object which can be used to access the input data
    @param n : format number. The function will be called with incrementing
                number until it returns zero
    @return: 0 - no more supported formats
                string "name" - format name to display in the chooser dialog
                dictionary { 'format': "name", 'options': integer }
                options: should be 1, possibly ORed with ACCEPT_FIRST (0x8000)
                to indicate preferred format
    """
    if IDA_SDK_VERSION < BASELINE_IDA_VERSION:
        return 0
    # we support only one format per file
    if n > 0: return 0
    # read 16K bytes to allow for the DTD
    data = li.read(0x4000)
    # look for start of <PROGRAM> element
    start = data.find("<PROGRAM")
    if start >= 0: return "XML PROGRAM file"
    return 0


def load_file(li, neflags, format):
    """
    Load the file into database

    @param li: a file-like object which can be used to access the input data
    @param neflags: options selected by the user, see loader.hpp
    @return: 0-failure, 1-ok
    """
    global event, element
    plugin = False
    status = 0
    st = idaapi.setStat(idaapi.st_Work)
    xml = XmlImporter(False, 0)
    try:
        status = xml.import_xml()
    except Cancelled:
        msg = "XML PROGRAM import cancelled!"
        print "\n" + msg
        idaapi.warning(msg)
    except MultipleAddressSpacesNotSupported:
        msg  = "XML Import cancelled!"
        msg += "\n\nXML Import does not currently support"
        msg += "\nimporting multiple address spaces."
        print "\n" + msg
        idaapi.warning(msg)
    except:
        print "\nHouston, we have a problem!"
        msg = "***** Exception occurred: XML loader failed! *****"
        print "\n" + msg + "\n", sys.exc_type, sys.exc_value
        print event, element.tag, element.attrib
        idaapi.warning(msg)
    finally:
        idaapi.setStat(st)
        xml.cleanup()
        return status


class XmlImporterPlugin(idaapi.plugin_t):
    """
    XML Importer plugin class
    """
    flags = 0
    comment = "Import XML PROGRAM file"
    help = "Import XML <PROGRAM> document to database"
    wanted_name = "XML Importer"
    wanted_hotkey = "Ctrl-Alt-l"

    def init(self):
        """
        init function for XML Importer plugin.
        
        Returns:
            Constant PLUGIN_OK if this IDA version supports the plugin,
            else returns PLUGIN_SKIP if this IDA is older than the supported
            baseline version.
        """
        if IDA_SDK_VERSION < BASELINE_IDA_VERSION:
            idaapi.msg('\nXML Importer plugin (xmlldr.py) not supported ' +
                        'by this version of IDA\n')
            return idaapi.PLUGIN_SKIP
        return idaapi.PLUGIN_OK

    def run(self, arg):
        """
        run function for XML Importer plugin.
        
        Args:
            arg: Integer, a non-zero value enables auto-run feature for
                 IDA batch (no gui) processing mode. Default is 0.
        """
        st = idaapi.setStat(idaapi.st_Work)
        xml = XmlImporter(True, arg)
        try:
            try:
                xml.import_xml()
            except Cancelled:
                msg = "XML Import cancelled!"
                print "\n" + msg
                idaapi.warning(msg)
            except MultipleAddressSpacesNotSupported:
                msg  = "XML Import cancelled!"
                msg += "\n\nXML Import does not currently support"
                msg += "\nimporting multiple address spaces."
                print "\n" + msg
                idaapi.warning(msg)
            except:
                msg = "***** Exception occurred: XML Importer failed! *****"
                print "\n" + msg + "\n", sys.exc_type, sys.exc_value
                idaapi.warning(msg)
        finally:
            xml.cleanup()
            idaapi.setStat(st)


    def term(self):
        pass


def PLUGIN_ENTRY():
    return XmlImporterPlugin()


class Cancelled(Exception):
    pass


class FileError(Exception):
    pass


class MultipleAddressSpacesNotSupported(Exception):
    pass


class XmlImporter:
    """
    XmlImporter class contains methods to import an XML PROGRAM
        document into IDA.
    """
    def __init__(self, plugin, arg):
        """
        Initializes the XmlImporter attributes

        Args:
            arg: Integer, non-zero value enables auto-run feature for
                IDA batch (no gui) processing mode. Default is 0.
        """
        self.plugin = plugin
        self.autorun = False
        self.debug = False # set this to True for debug prints
        self.Elements = {}
        self.Counters = []
        self.Tags = []
        self.Timers = dict()
        self.xmlfile = 0
        self.addr_mode = 1
        self.create = True
        self.dataseg = None
        self.deferred = []
        self.callbacks = {
            'start' : {
                BOOKMARKS           : self.update_import,
                CODE                : self.update_import,
                COMMENTS            : self.update_import,
                COMPILER            : self.import_compiler,
                DATA                : self.update_import,
                DATATYPES           : self.update_import,
                EQUATES             : self.update_import,
                FUNCTIONS           : self.update_import,
                INFO_SOURCE         : self.import_info_source,
                MARKUP              : self.update_import,
                MEMORY_MAP          : self.import_memory_map,
                PROCESSOR           : self.import_processor,
                PROGRAM             : self.import_program,
                PROGRAM_ENTRY_POINTS: self.update_import,
                REGISTER_VALUES     : self.update_import,
                SYMBOL_TABLE        : self.update_import },
            'end' : {
                BOOKMARK            : self.import_bookmark,
                CODE_BLOCK          : self.import_codeblock,
                COMMENT             : self.import_comment,
                DEFINED_DATA        : self.import_defined_data,
                ENUM                : self.import_enum,
                EQUATE_GROUP        : self.import_equate_group,
                EQUATE_REFERENCE    : self.import_equate_reference,
                FUNCTION            : self.import_function,
                FUNCTION_DEF        : self.import_function_def,
                MANUAL_INSTRUCTION  : self.import_manual_instruction,
                MANUAL_OPERAND      : self.import_manual_operand,
                MEMORY_REFERENCE    : self.import_memory_reference,
                MEMORY_SECTION      : self.import_memory_section,
                PROGRAM_ENTRY_POINT : self.import_program_entry_point,
                REGISTER_VALUE_RANGE: self.import_register_value_range,
                STACK_REFERENCE     : self.import_stack_reference,
                STRUCTURE           : self.import_structure,
                SYMBOL              : self.import_symbol,
                TYPE_DEF            : self.import_typedef,
                UNION               : self.import_union,
                # end element for elapse time
                BOOKMARKS           : self.display_timer,
                CODE                : self.display_timer,
                COMMENTS            : self.display_timer,
                DATA                : self.display_timer,
                DATATYPES           : self.process_deferred,
                EQUATES             : self.display_timer,
                FUNCTIONS           : self.display_timer,
                MARKUP              : self.display_timer,
                MEMORY_MAP          : self.display_timer,
                PROGRAM             : self.display_total_time,
                PROGRAM_ENTRY_POINTS: self.display_timer,
                REGISTER_VALUES     : self.display_timer,
                SYMBOL_TABLE        : self.display_timer }
            }


    def import_xml(self):
        """
        Imports the XML PROGRAM file into the database.
        """
        self.display_xml_importer_version()
        displayMenu = self.plugin == True and self.autorun == False 
        #displayMenu = self.autorun == False 
        self.get_options(displayMenu)
        if self.plugin:
            self.filename=idaapi.askfile_c(0, "*.xml",
                                        "Enter name of xml file:")
        else:
            self.filename = idaapi.get_input_file_path()
        if (len(self.filename) == 0):
            return
        idaapi.msg('\nImporting from: ' + self.filename + '\n') 
        idc.Wait()
        if self.plugin == False: idaapi.hide_wait_box()
        idaapi.show_wait_box("Importing XML PROGRAM document....")
        n = 0
        for event,element in cElementTree.iterparse(self.filename,
                                                events=("start","end")):
            if idaapi.wasBreak() == True:
                raise Cancelled
            
            if self.debug == True and event == 'start':
                print element.tag, element.attrib, element.text
            
            if event in self.callbacks:
                if element.tag in self.callbacks[event]:
                    if event == 'start':
                        self.Timers[element.tag] = time.clock()
                    self.callbacks[event][element.tag](element)
                    if event == 'end':
                        element.clear()
            if event == 'end':
                n += 1
        end = time.clock()
        idaapi.hide_wait_box()
        self.display_summary()
        idaapi.msg('\nXML Elements parsed: ' + str(n) + '\n\n')
        return 1
    

    def get_options(self, display):
        """
        Displays the options menu and retrieves the option settings. 
        """
        fmt =  "HELP\n"
        fmt += "XML PROGRAM loader/importer plugin (Python)\n"
        fmt += "IDA SDK: "+ str(IDA_SDK_VERSION) + "\n\n"
        fmt +=  "The XML PROGRAM loader loads elements from a "
        fmt +=  "XML <PROGRAM> document to create an idb database.\n\n"
        fmt +=  "ENDHELP\n"
        fmt +=  "Import from XML PROGRAM document...."
        fmt += "\n <##Options##Code Blocks:{CodeBlocks}>"
        fmt += "\n <Entry Points:{EntryPoints}>"
        fmt += "\n <Segment Register Value Ranges:{RegisterValues}>"
        fmt += "\n <Data Types:{DataTypes}>"
        fmt += "\n <Data Definitions:{DataDefinitions}>"
        fmt += "\n <Symbols:{Symbols}>"
        fmt += "\n <Comments:{Comments}>"
        fmt += "\n <Bookmarks:{Bookmarks}>"
        fmt += "\n <Functions:{Functions}>"
        fmt += "\n <Memory References:{MemoryReferences}>"
        fmt += "\n <Equate/Enum References:{EquateReferences}>"
        fmt += "\n <Manual Instructions/Operands:{Manual}>{cGroup1}>"
        fmt += "\n\n"
    
        Opts = { 'cGroup1': idaapi.Form.ChkGroupControl(
                ( 
                "CodeBlocks", "EntryPoints", "RegisterValues",
                "DataTypes", "DataDefinitions",
                "Symbols", "Comments", "Bookmarks",
                "Functions", "MemoryReferences", "EquateReferences",
                "Manual")) }
        
        self.Options = idaapi.Form(fmt, Opts)
        self.Options.Compile()
    
        self.Options.CodeBlocks.checked       = True
        self.Options.EntryPoints.checked      = True
        self.Options.RegisterValues.checked   = True
        self.Options.DataTypes.checked        = True
        self.Options.DataDefinitions.checked  = True
        self.Options.Symbols.checked          = True
        self.Options.Functions.checked        = True
        self.Options.Comments.checked         = True
        self.Options.Bookmarks.checked        = True
        self.Options.MemoryReferences.checked = True
        self.Options.EquateReferences.checked = True
        self.Options.Manual.checked           = True
        
        if display == True:
            ok = self.Options.Execute()
            if (ok == 0):
                raise Cancelled
            
    def cleanup(self):
        """
        Frees memory and closes message box at termination.
        """
        if self.plugin:
            self.Options.Free()
        idaapi.hide_wait_box()
    

    def display_timer(self, element):
        """
        Displays the elapsed processing time for XML elements.
        
        Args:
            element: XML element object value containing the element tag.
        """
        if element.tag == MEMORY_MAP and self.plugin:
            return
        if element.tag in self.Timers:
            idaapi.msg('elapsed time: %.4f' %
                    (time.clock()-self.Timers[element.tag]))
    

    def display_total_time(self, element):
        """
        Displays the total processing time.
        
        Args:
            element: XML element object value (not used).
        """
        TOTAL = 'Total '
        idaapi.msg('\n%35selapsed time: %.4f' %
                (TOTAL,time.clock()-self.Timers[PROGRAM]))
    

    def display_summary(self):
        """
        Displays summary of the XML PROGRAM import in IDA output window.
        """
        summary = ''
        total = 0
        for tag in self.Tags:
            count = self.Counters[self.Elements[tag]]
            summary += "\n%s: %d" % (tag, count)
            total += count
        summary += '\n--------------------------------------'
        header = '\n--------------------------------------'
        header += ('\nTotal XML Elements Processed: %d') % total
        summary = header + summary
        idaapi.msg(summary)
        if self.plugin and self.autorun == False:
            frmt  = "TITLE XML Import Successful!\n"
            frmt += "ICON INFO\n"
            frmt += "AUTOHIDE NONE\n"
            frmt += "HIDECANCEL\n"
            firstline = "\n XML IMPORT SUCCESSFUL!"
            fileline = '\n\nFile: %s' % self.filename
            details = '\n\nSee output window for details...'
            idaapi.info("%s" % (frmt + firstline + fileline + details))
    

    def display_xml_importer_version(self):
        """
        Displays XML Importer plugin version info in IDA output window.
        """
        if self.plugin:
            f = idaapi.idadir(idaapi.PLG_SUBDIR) + '/xmlldr.py'
        else:
            f = idaapi.idadir(idaapi.LDR_SUBDIR) + '/xmlldr.py'
        plugintime = time.localtime(os.path.getmtime(f))
        ts = time.strftime('%b %d %Y %H:%M:%S', plugintime)
        version = "\nXML Importer Version " + XML_IMPORTER_VERSION
        version += " : SDK " + str(IDA_SDK_VERSION)
        version += " : Python : "+ ts + '\n'
        idaapi.msg(version)
    

    def get_address(self, element, attr):
        """
        Returns the address value for an element.
        
        Args:
            element: XML element object.
            attr: String containing the address attribute name.
            
        Returns:
            Numeric value representing the address.
        """
        addrstr = element.get(attr)
        if '::' in addrstr:
            # overlayed addresses not currently handled
            return BADADDR
        elif ':' in addrstr:
            [segstr, offset_str] = string.split(addrstr,':')
            offset = int(offset_str,16)
            if self.is_int(segstr) == True:
                sgmt = int(segstr,16)
                addr = (sgmt << 4) + offset
            else:
                # multiple address spaces not currently implemented
                addr = BADADDR
            return addr
        else:
            return int(element.get(attr), 16)
    

    def get_attribute(self, element, attr):
        """
        Returns the attribute value string.
        
        Args:
            element: XML element object.
            attr: String containing the attribute name.
            
        Returns:
            String representing the attribute value.
        """
        return element.get(attr)
    

    def get_attribute_value(self, element, attr):
        """
        Returns the numeric attribute value.
        
        Args:
            element: XML element object.
            attr: String containing the attribute name.
            
        Returns:
            Numeric value representing the attribute value.
        """
        val = element.get(attr)
        try:
            if val.upper().startswith('0X') or val.upper().startswith('-0X'):
                return int(val, 16)
            return int(val)
        except:
            idaapi.msg('\nUnable to decode string as value: ' + val)
            return 0
        

    def get_cbsize(self):
        """
        Returns the size of the addressable codebyte for the processor.
        
        Returns:
            Integer representing the number of 8-bit bytes in an
            addressable codebyte.
        """
        return (idaapi.ph_get_cnbits()+7)/8
    

    def get_datatype_flags(self, datatype, size):
        """
        Returns the flags bitmask for the datatype.
        
        Args:
            datatype: String representing the datatype.
            size: Integer representing the datatype size.

        Returns:
            Integer representing the bitmask.
        """
        if datatype.lower().startswith("byte"):     return idaapi.byteflag()
        if datatype.lower().startswith("word"):     return idaapi.wordflag()
        if datatype.lower().startswith("dword"):    return idaapi.dwrdflag()
        if datatype.lower().startswith("qword"):    return idaapi.qwrdflag()
        if datatype.lower().startswith("oword"):    return idaapi.owrdflag()
        if datatype.lower().startswith("tbyte"):    return idaapi.tbytflag()
        if datatype.lower().startswith("float"):    return idaapi.floatflag()
        if datatype.lower().startswith("double"):   return idaapi.doubleflag()
        if datatype.lower().startswith("packed"):   return idaapi.packrealflag()
        if self.is_string_type(datatype):           return idaapi.asciflag()
        if self.is_enumeration(datatype):           return idaapi.enumflag()
        if self.is_structure(datatype):             return idaapi.struflag()
        #if size == 4:                               return idaapi.dwrdflag()
        return 0
    

    def get_string_type(self, datatype):
        if datatype.lower() == 'mbcstring':
            return idaapi.ASCSTR_UNICODE
        if datatype.lower().find('unicode') != -1:
            if datatype.lower().find('pascal') != -1:
                return idaapi.ASCSTR_ULEN2
            return idaapi.ASCSTR_UNICODE
        if datatype.lower().find('pascal') != -1:
            return idaapi.ASCSTR_LEN2
        return idaapi.ASCSTR_TERMCHR
            
    
    def has_attribute(self, element, attr):
        """
        Returns true if the XML element contains the named attribute.
        
        Args:
            element: XML element object
            attr: String containing name of the attribute
        
        Returns:
            True if the element contains the named attribute, otherwise False.
        """
        return attr in element.attrib
    

    def is_enumeration(self, datatype):
        """
        Returns true if datatype is an existing enumeration in the database.
        
        Args:
            datatype: String representing the datatype.
            
        Returns:
            True if the datatype is an enumeration in the database,
            otherwise False.
        """
        if idaapi.get_enum(datatype) == idaapi.BADNODE:  return False
        return True
    
    
    def is_int(self, s):
        try:
            int(s, 16)
            return True
        except:
            return False
        

    def is_pointer_type(self, dtype):
        """
        Returns true if the datatype represents a pointer.
        
        Args:
            dtype: String representing the datatype.
            
        Returns:
            True if the datatype represents a pointer, otherwise False.
        """
        if dtype.lower().startswith("pointer") or dtype.endswith('*'):
            return True
        return False
    

    def is_string_type(self, datatype):
        """
        Returns true if the datatype represents a string type.
        
        Args:
            datatype: String representing the datatype.
            
        Returns:
            True if the datatype represents a string, otherwise False.
        """
        if datatype.lower().startswith("unicode"):  return True
        if datatype.lower().startswith("string"):   return True
        return False
    

    def is_structure(self, datatype):
        """
        Returns true if the datatype represents a structure in the database.
        
        Args:
            dtype: String representing the datatype.
            
        Returns:
            True if the datatype represents an existing structure,
            otherwise False.
        """
        if idaapi.get_struc_id(datatype) == idaapi.BADNODE:  return False
        return True
    

    def import_address_range(self, address_range):
        """
        Processes ADDRESS_RANGE element.
        
        Args:
            address_range: XML element object containing start and end address
                attributes for the address range.
                
        Returns:
            Tuple containing two integers, the start and end address values.
        """
        start = self.get_address(address_range,START)
        end = self.get_address(address_range, END)
        self.update_counter(ADDRESS_RANGE)
        return (start, end)
    

    def import_bit_mask(self, bitmask, eid):
        """
        Processes a BIT_MASK element as an enum bitmask member.
        
        Args:
            bitmask: XML element object representing the IDA enum bitmask.
            eid: Integer representing the IDA enum id
        """
        if IDA_SDK_VERSION < 640:
            return
        name = self.get_attribute(bitmask,NAME)
        value = self.get_attribute_value(bitmask,VALUE)
        idaapi.set_bmask_name(eid, value, name)
        cid = idaapi.get_const_by_name(name)
        self.update_counter(BIT_MASK)
        regcmt = bitmask.find(REGULAR_CMT)
        if regcmt != None:
            idaapi.set_const_cmt(cid, regcmt.text, False);
            self.update_counter(BIT_MASK + ':' + REGULAR_CMT)
        rptcmt = bitmask.find(REPEATABLE_CMT)
        if rptcmt != None:
            idaapi.set_const_cmt(cid, rptcmt.txt, True);
            self.update_counter(BIT_MASK + ':' + REPEATABLE_CMT)
    

    def import_bookmark(self, bookmark):
        """
        Processes a BOOKMARK element.
        
        Args:
            bookmark: XML element object containing bookmark data.
        """
        if self.Options.Bookmarks.checked == False:
            return
        try:
            addr = self.get_address(bookmark, ADDRESS)
            if self.has_attribute(bookmark, TYPE):
                typ = self.get_attribute(bookmark, TYPE)
            category = ''
            if self.has_attribute(bookmark, CATEGORY):
                category = self.get_attribute(bookmark, CATEGORY)
            description = ''
            if self.has_attribute(bookmark, DESCRIPTION):
                description = self.get_attribute(bookmark, DESCRIPTION)
            if idaapi.isEnabled(addr) == False:
                msg = ("import_bookmark: address %X not enabled in database"
                       % addr)
                print msg
                return
            self.update_counter(BOOKMARK)
            for slot in range(1,1025):
                if IDA_SDK_VERSION != 695:
                    ea = idc.GetMarkedPos(slot)
                else:
                    import ida_moves
                    import ida_pro
                    curloc = ida_moves.curloc()
                    intp = ida_pro.int_pointer()
                    intp.assign(slot)
                    ea = curloc.markedpos(intp)
                if ea == BADADDR:
                    curloc = idaapi.curloc()
                    curloc.ea = addr
                    curloc.mark(slot, category, description)
                    break
        except:
            msg = "** Exception occurred in import_bookmark **"
            print "\n" + msg + "\n", sys.exc_type, sys.exc_value
    

    def import_cmts(self, element, sid, typ):
        """
        Processes REGULAR_CMT and REPEATABLE_CMT elements for structures.
        
        Args:
            element: XML element object containing a REGULAR_CMT or
                REPEATABLE_CMT element
            sid: Integer representing the structure id
            typ: String indicating structure type (STRUCTURE or UNION)
        """
        regcmt = element.find(REGULAR_CMT)
        if regcmt != None:
            idaapi.set_struc_cmt(sid, regcmt.text, False)
            self.update_counter(typ + ':' + REGULAR_CMT)
        rptcmt = element.find(REPEATABLE_CMT)
        if rptcmt != None:
            idaapi.set_struc_cmt(sid, rptcmt.text, True)
            self.update_counter(typ + ':' + REPEATABLE_CMT)
            
        
    def import_codeblock(self, code_block):
        """
        Processes a CODE_BLOCK element by disassembling the address range.
        
        Args:
            code_block: XML element containing codeblock start and end
                addresses.
        """
        if self.Options.CodeBlocks.checked == False:
            return
        start = self.get_address(code_block, START)
        end = self.get_address(code_block, END)
        idaapi.do_unknown_range(start, end-start+1, 3)
        addr = start
        while (addr <= end):
            length = idaapi.create_insn(addr)
            addr += idaapi.get_item_size(addr) * self.get_cbsize()
        self.update_counter(CODE_BLOCK)
    

    def import_comment(self, comment):
        """
        Processes a COMMENT element by creating the comment at the address.
        
        Args:
            comment: XML element containing the comment address, type,
                and text.
        """
        if self.Options.Comments.checked == False:
            return
        addr = self.get_address(comment, ADDRESS)
        ctype = self.get_attribute(comment,TYPE)
        text = comment.text
        if ctype == 'pre':
            idaapi.add_long_cmt(addr, True, text)
        elif ctype == 'end-of-line':
            idaapi.set_cmt(addr, text, False)
        elif ctype == 'repeatable':
            idaapi.set_cmt(addr, text, True)
        elif ctype == 'post':
            idaapi.add_long_cmt(addr, False, text)
        self.update_counter(COMMENT+':' + ctype)
    

    def import_compiler(self, compiler):
        """
        Processes the COMPILER element containing the compiler name.
        
        Args:
            compiler: XML element containing the compiler name.
        """
        name = self.get_attribute(compiler, NAME)
        comp = idaapi.COMP_UNK
        if   name == "Visual C++":      comp = idaapi.COMP_MS
        elif name == "Borland C++":     comp = idaapi.COMP_BC
        elif name == "Watcom C++":      comp = idaapi.COMP_WATCOM
        elif name == "GNU C++":         comp = idaapi.COMP_GNU
        elif name == "Visual Age C++":  comp = idaapi.COMP_VISAGE
        elif name == "Delphi":          comp = idaapi.COMP_BP
        idaapi.cvar.inf.cc.id = comp
        self.update_counter(COMPILER)
    

    def import_defined_data(self, defined_data):
        """
        Processes a DEFINED_DATA element by creating a data item at the
            specified address.
        
        Args:
            defined_data: XML element containing the address and
                datatype information for the data item
        """
        if self.Options.DataDefinitions.checked == False:
            return
        addr = self.get_address(defined_data, ADDRESS)
        datatype = self.get_attribute(defined_data, DATATYPE)
        size = self.get_attribute_value(defined_data, SIZE)
        self.update_counter(DEFINED_DATA)        
        ti = idaapi.opinfo_t()
        if self.is_pointer_type(datatype):
            #idaapi.set_refinfo(ti, 0, 0, 0, REF_OFF32)
            flag = idaapi.dwrdflag() | idaapi.FF_0OFF
            #idaapi.set_typeinfo(addr, 0, flag, ti)
        else:
            flag = self.get_datatype_flags(datatype, size)
        if flag == idaapi.asciflag():
            idaapi.make_ascii_string(addr, size,
                                     self.get_string_type(datatype))
        elif flag == idaapi.struflag():
            idaapi.doStruct(addr, size, idaapi.get_struc_id(datatype))
        else:
            idaapi.do_data_ex(addr, flag, size, idaapi.BADNODE)        
        typecmt = defined_data.find(TYPEINFO_CMT)
        if typecmt != None:
            self.update_counter(DEFINED_DATA + ':' + TYPEINFO_CMT)


    def import_enum(self, enum):
        """
        Processes an ENUM element by creating the enumeration.
        
        Args:
            enum: XML element containing the enumeration name and
                member data.
        """
        if self.Options.DataTypes.checked == False:
            return
        name = self.get_attribute(enum, NAME)
        if self.has_attribute(enum,NAMESPACE):
            namespace = self.get_attribute(enum, NAMESPACE)
        if self.has_attribute(enum,SIZE):
            size = self.get_attribute_value(enum, SIZE)
        eid = idaapi.add_enum(BADADDR, name,
                              idaapi.hexflag() | idaapi.dwrdflag())
        self.update_counter(ENUM)
        regcmt = enum.find(REGULAR_CMT)
        if regcmt != None:
            idaapi.set_enum_cmt(eid, regcmt.text, False)
            self.update_counter(ENUM + ':' + REGULAR_CMT)
        rptcmt = enum.find(REPEATABLE_CMT)
        if rptcmt != None:
            idaapi.set_enum_cmt(eid, rptcmt.text, True)
            self.update_counter(ENUM + ':' + REPEATABLE_CMT)
        display_settings = enum.find(DISPLAY_SETTINGS)
        if display_settings != None:
            self.update_counter(ENUM + ':' + DISPLAY_SETTINGS)
        enum_entries = enum.findall(ENUM_ENTRY)
        for enum_entry in enum_entries:
            self.import_enum_entry(enum_entry, eid)
    

    def import_enum_entry(self, enum_entry, eid):
        """
        Processes an ENUM_ENTRY by creating a member in the enumeration.
        
        Args:
            enum_entry: XML element containing the member name and value.
            eid: Integer representing the id of the enumeration.
        """
        name = self.get_attribute(enum_entry, NAME)
        value = self.get_attribute_value(enum_entry, VALUE)
        if IDA_SDK_VERSION > 630:
            idaapi.add_const(eid, name, value)
        else:
            idaapi.add_enum_member(eid, name, value)
        if IDA_SDK_VERSION > 630:
            cid = idaapi.get_const_by_name(name)
        else:
            cid = idaapi.get_enum_member_by_name(name)
        self.update_counter(ENUM_ENTRY)
        regcmt = enum_entry.find(REGULAR_CMT)
        if regcmt != None:
            if IDA_SDK_VERSION > 630:
                idaapi.set_const_cmt(cid, regcmt.text, False);
            else:
                idaapi.set_enum_member_cmt(cid, regcmt.text, False);
            self.update_counter(ENUM_ENTRY + ':' + REGULAR_CMT)
        rptcmt = enum_entry.find(REPEATABLE_CMT)
        if rptcmt != None:
            if IDA_SDK_VERSION > 630:
                idaapi.set_const_cmt(cid, rptcmt.text, True);
            else:
                idaapi.set_enum_member_cmt(cid, rptcmt.text, True);
            self.update_counter(ENUM_ENTRY + ':' + REPEATABLE_CMT)
    

    def import_equate(self, equate, eid):
        """
        Processes EQUATE element as member of an enumeration.
        
        Args:
            enum_entry: XML element containing the equate name and value.
            eid: Integer representing the id for the enumeration.
        """
        name = self.get_attribute(equate,NAME)
        value = self.get_attribute_value(equate,VALUE)
        bm = -1
        if self.has_attribute(equate, BIT_MASK):
            bm = self.get_attribute_value(equate, BIT_MASK)
        if IDA_SDK_VERSION > 630:
            #idaapi.add_const(eid, name, value, bm)
            idaapi.add_const(eid, name, value)
            cid = idaapi.get_const_by_name(name)
        else:
            #idaapi.add_enum_member(eid, name, value, bm)
            idaapi.add_enum_member(eid, name, value)
            cid = idaapi.get_enum_member_by_name(name)
        self.update_counter(EQUATE)
        regcmt = equate.find(REGULAR_CMT) 
        if regcmt != None:
            if IDA_SDK_VERSION > 630:
                idaapi.set_const_cmt(cid, regcmt.text, False);
            else:
                idaapi.set_enum_member_cmt(cid, regcmt.text, False);
            self.update_counter(EQUATE + ':' + REGULAR_CMT)
        rptcmt = equate.find(REPEATABLE_CMT)
        if rptcmt != None:
            if IDA_SDK_VERSION > 630:
                idaapi.set_const_cmt(cid, rptcmt.text, True);
            else:
                idaapi.set_enum_member_cmt(cid, rptcmt.text, True);
            self.update_counter(EQUATE + ':' + REPEATABLE_CMT)
    

    def import_equate_group(self, equate_group):
        """
        Processes EQUATE_GROUP as IDA enumeration type.
        
        Args:
            equate_group: XML element containing the group name and
                equate definitions.
        """
        if self.Options.DataTypes.checked == False:
            return
        msg = EQUATE_GROUP
        name = ''
        if self.has_attribute(equate_group, NAME):
            name = self.get_attribute(equate_group, NAME)
        bf = ''
        if self.has_attribute(equate_group, BIT_FIELD):
            bf = self.get_attribute(equate_group, BIT_FIELD)
        eid = idaapi.add_enum(BADADDR, name, idaapi.hexflag())
        idaapi.set_enum_bf(eid, (bf == 'yes'))
        self.update_counter(EQUATE_GROUP)
        regcmt = equate_group.find(REGULAR_CMT)
        if regcmt != None:
            idaapi.set_enum_cmt(eid, regcmt.text, False)
            self.update_counter(EQUATE_GROUP + ':' + REGULAR_CMT)
        rptcmt = equate_group.find(REPEATABLE_CMT)
        if rptcmt != None:
            idaapi.set_enum_cmt(eid, rptcmt.text, True)
            self.update_counter(EQUATE_GROUP + ':' + REPEATABLE_CMT)
        equates = equate_group.findall(EQUATE)
        for equate in equates:
            self.import_equate(equate,eid)
        bit_masks = equate_group.findall(BIT_MASK)
        for bit_mask in bit_masks:
            self.import_bit_mask(bit_mask, eid)
    

    def import_equate_reference(self, equate_reference):
        if (self.Options.DataTypes.checked == False or
            self.Options.EquateReferences.checked == False):
            return
        self.update_counter(EQUATE_REFERENCE)
        addr = self.get_address(equate_reference, ADDRESS)
        name = ''
        if self.has_attribute(equate_reference, NAME):
            name = self.get_attribute(equate_reference, NAME)
        if name == '':
            return
        opnd = 0
        if self.has_attribute(equate_reference, OPERAND_INDEX):
            opnd = self.get_attribute_value(equate_reference, OPERAND_INDEX)
        value = None
        if self.has_attribute(equate_reference, VALUE):
            value = self.get_attribute_value(equate_reference, VALUE)
        if IDA_SDK_VERSION < 650:
            cid = idaapi.get_enum_member_by_name(name)
        else:
            cid = idaapi.get_const_by_name(name)
        if cid == idaapi.BADNODE:
            return
        if IDA_SDK_VERSION < 650:
            eid = idaapi.get_enum_member_enum(cid)
        else:
            eid = idaapi.get_const_enum(cid)
        if eid == idaapi.BADNODE:
            return
        idaapi.op_enum(addr, opnd, eid, 0)
        

    def import_function(self, function):
        """
        Creates a function using the FUNCTION attributes.
        
        Args:
            function: XML element containing the function address and
                attributes.
        """
        if self.Options.Functions.checked == False:
            return
        try:
            entry_point = self.get_address(function, ENTRY_POINT)
            name = ''
            if self.has_attribute(function, NAME):
                name = self.get_attribute(function, NAME)
            libfunc = 'n'
            if self.has_attribute(function, LIBRARY_FUNCTION):
                libfunc = self.get_attribute(function, LIBRARY_FUNCTION)
            if idaapi.isEnabled(entry_point) == False:
                msg = ("import_function: address %X not enabled in database"
                       % entry_point)
                print msg
                return
            idaapi.add_func(entry_point, BADADDR)
            self.update_counter(FUNCTION)
            func = idaapi.get_func(entry_point)
            if libfunc == 'y':
                func.flags |= idaapi.FUNC_LIB
            ranges = function.findall(ADDRESS_RANGE)
            for addr_range in ranges:
                (start, end) = self.import_address_range(addr_range)
                idaapi.append_func_tail(func, start, end)
                #idaapi.analyze_area(start, end+1)
            idc.Wait()
            regcmt = function.find(REGULAR_CMT)
            if regcmt != None:
                self.update_counter(FUNCTION + ':' + REGULAR_CMT)
                idaapi.set_func_cmt(func, regcmt.text, False)
            rptcmt = function.find(REPEATABLE_CMT)
            if rptcmt != None:
                self.update_counter(FUNCTION + ':' + REPEATABLE_CMT)
                idaapi.set_func_cmt(func, rptcmt.text, True)
            typecmt = function.find(TYPEINFO_CMT)
            if typecmt != None:
                self.update_counter(FUNCTION + ':' + TYPEINFO_CMT)
                idc.SetType(entry_point, typecmt.text + ';')
            sf = function.find(STACK_FRAME)
            if sf != None:
                self.import_stack_frame(sf, func)
            register_vars = function.findall(REGISTER_VAR)
            for register_var in register_vars:
                self.import_register_var(register_var, func)
        except:
            msg = "** Exception occurred in import_function **"
            print "\n" + msg + "\n", sys.exc_type, sys.exc_value


    def import_function_def(self, function_def):
        # import_function_def: NOT IMPLEMENTED
        if self.Options.DataTypes.checked == False:
            return
        self.update_counter(FUNCTION_DEF)
        

    def import_info_source(self, info_source):
        """
        Processes INFO_SOURCE containing information about the
            source of the XML PROGRAM file.
        
        Args:
            info_source: XML element containing attributes that identify
                the source of the PROGRAM data.
        """
        if self.has_attribute(info_source, TOOL):
            tool = self.get_attribute(info_source, TOOL)
        if self.has_attribute(info_source, USER):
            user = self.get_attribute(info_source, USER)
        if self.has_attribute(info_source, FILE):
            f = self.get_attribute(info_source, FILE)
        if self.has_attribute(info_source, TIMESTAMP):
            ts = self.get_attribute(info_source, TIMESTAMP)
        self.update_counter(INFO_SOURCE)
    

    def import_manual_instruction(self, manual_instruction):
        """
        Creates a manual instruction.
        
        Args:
            manual_instruction: XML element containing MANUAL_INSTRUCTION.
        """
        if self.Options.Manual.checked == False:
            return
        addr = self.get_address(manual_instruction, ADDRESS)
        idaapi.set_manual_insn(addr, manual_instruction.text)
        self.update_counter(MANUAL_INSTRUCTION)
    

    def import_manual_operand(self, manual_operand):
        """
        Creates a manual operand at an address.
        
        Args:
            manual_operand: MANUAL_OPERAND XML element.
        """
        if self.Options.Manual.checked == False:
            return
        addr = self.get_address(manual_operand, ADDRESS)
        op = self.get_attribute_value(manual_operand, OPERAND_INDEX)
        if idaapi.isEnabled(addr):
            idaapi.set_forced_operand(addr, op, manual_operand.text)
            self.update_counter(MANUAL_OPERAND)


    def process_deferred(self, element):
        """
        Processes the list of deferred structure members when the
        DATATYPES end element is encountered.
        
        Args:
            element: XML end element for DATATYPES 
        """
        for (member, sptr) in self.deferred:
            self.import_member(member, sptr, False)
        self.display_timer(element)


    def import_member(self, member, sptr, defer=True):
        """
        Creates a member for a structure.
        
        Args:
            member: MEMBER XML element.
            sptr:
            defer:  boolean indicating if processing a member should be
                    deferred when the type is unknown. A member should
                    only be deferred on the first pass, not when processing
                    the deferred list.
        """
        offset = self.get_attribute_value(member, OFFSET)
        datatype = self.get_attribute(member, DATATYPE)
        if self.has_attribute(member, DATATYPE_NAMESPACE):
            dt_namespace = self.get_attribute(member, DATATYPE_NAMESPACE)
        name = ''
        if self.has_attribute(member, NAME):
            name = self.get_attribute(member, NAME)
        size = 0
        if self.has_attribute(member, SIZE):
            size = self.get_attribute_value(member, SIZE)
        ti = idaapi.opinfo_t()
        if self.is_pointer_type(datatype):
            #idaapi.set_refinfo(ti, 0, BADADDR, 0, REF_OFF32)
            flag = idaapi.dwrdflag() | idaapi.FF_0OFF
        else:
            flag = self.get_datatype_flags(datatype, size)
        if flag == 0:
            if defer:
                self.deferred.append([member, sptr])
                return                
        if flag == idaapi.enumflag():
            t = idaapi.get_enum(datatype)
            ti.ec.tid = t
            ti.ec.serial = idaapi.get_enum_idx(t)
        if flag == idaapi.struflag():
            t = idaapi.get_struc_id(datatype)
            ti.tid = t
        error = idaapi.add_struc_member(sptr, name, offset, flag, ti, size)
        mbr = idaapi.get_member(sptr, offset)
        self.import_member_cmts(member, mbr)
        self.update_counter(MEMBER)
        
        
    def import_member_cmts(self, member, mbr):
        """
        Processes REGULAR_CMT and REPEATABLE_CMT elements for members.
        
        Args:
            element: XML element object containing a REGULAR_CMT or
                REPEATABLE_CMT element
            mbr: Integer representing the member id
        """
        regcmt = member.find(REGULAR_CMT)
        if regcmt != None:
            idaapi.set_member_cmt(mbr, regcmt.text, False)
            self.update_counter(MEMBER + ':' + REGULAR_CMT)
        rptcmt = member.find(REPEATABLE_CMT)
        if rptcmt != None:
            idaapi.set_member_cmt(mbr, rptcmt.text, True)
            self.update_counter(MEMBER + ':' + REPEATABLE_CMT)
        

    def import_members(self, element, sptr):
        """
        Add data members to a structure.
        
        Args:
            element: STRUCTURE XML element containing MEMBER sub-elements.
            sptr:
        """
        members = element.findall(MEMBER)
        for member in members:
            self.import_member(member, sptr)
                

    def import_memory_contents(self, memory_contents, start, size):
        """
        Processes MEMORY_CONTENTS to load data for a memory block.
        
        Args:
            memory_contents: MEMORY_CONTENTS XML element.
        """
        if memory_contents.get(START_ADDR) == None:
            saddr = start
        else:
            saddr = self.get_address(memory_contents, START_ADDR)
        fname = self.get_attribute(memory_contents, FILE_NAME)
        offset = self.get_attribute_value(memory_contents, FILE_OFFSET)
        if memory_contents.get(LENGTH) == None:
            length = size
        else:
            length = self.get_attribute_value(memory_contents, LENGTH)
        #(binfilename, ext) = os.path.splitext(self.filename)
        #binfilename += ".bytes"
        (binfilename, fileext) = os.path.split(self.filename)
        binfilename += '/' + fname
        binfile = idaapi.loader_input_t()
        binfile.open(binfilename)
        binfile.file2base(offset,saddr,saddr+length,False)
        binfile.close()
        self.update_counter(MEMORY_CONTENTS)
    

    def import_memory_map(self, memory_map):
        """
        Processes the MEMORY_MAP element.
        
        Args:
            memory_map: MEMORY_MAP XML element.
            
        MEMORY_MAP is only processed by the IDA loader. It is ignored when
            run as an IDA plugin.
        """
        # import memory sections only when run as loader
        if self.plugin:
            return
        self.update_import(memory_map)
    

    def import_memory_reference(self, memory_reference):
        """
        Processes the MEMORY_REFERENCE element.
        
        Args:
            memory_reference: MEMORY_REFERENCE XML element.
        """
        if self.Options.MemoryReferences.checked == False:
            return        
        addr = self.get_address(memory_reference, ADDRESS)
        if self.has_attribute(memory_reference, OPERAND_INDEX):
            op = self.get_attribute_value(memory_reference, OPERAND_INDEX)
        if self.has_attribute(memory_reference, USER_DEFINED):
            user = self.get_attribute(memory_reference, USER_DEFINED)
        to_addr = self.get_address(memory_reference, TO_ADDRESS)
        if self.has_attribute(memory_reference, BASE_ADDRESS):
            base_addr = self.get_address(memory_reference, BASE_ADDRESS)
        if self.has_attribute(memory_reference, PRIMARY):
            primary = self.get_attribute(memory_reference, PRIMARY)
        self.update_counter(MEMORY_REFERENCE)
        # TODO: import_memory_reference - add code to store reference
 
 
    def import_memory_section(self, memory_section):
        """
        Creates a memory segment in the database.
        
        Args:
            memory_section: MEMORY_SECTION XML element.
            
        MEMORY_SECTION is only processed by the IDA loader. It is ignored
            when run as an IDA plugin.
        """
        # TODO: import_memory_section - handle overlays?
        # import memory sections only when run as loader
        if self.plugin:
            return
        name = self.get_attribute(memory_section, NAME)
        length = self.get_attribute_value(memory_section, LENGTH)

        s = idaapi.segment_t()
        addrstr = self.get_attribute(memory_section, START_ADDR)
        seg_str = ''
        if '::' in addrstr:
            # overlay - skip for now
            print '  ** Overlayed memory block %s skipped **  ' % name
            msg  = 'Overlayed memory block %s skipped!' % name
            msg += "\n\nXML Import does not currently support"
            msg += "\noverlayed memory blocks."
            idaapi.warning(msg)
            return
        elif ':' in addrstr:
            [seg_str, offset_str] = string.split(addrstr,':')
            offset = int(offset_str, 16)
            if self.is_int(seg_str):
                base = int(seg_str, 16)
                sel = idaapi.setup_selector(base)
                start = self.get_address(memory_section, START_ADDR)
            else:
                raise MultipleAddressSpacesNotSupported
                return
        else:
            sel = idaapi.allocate_selector(0)
            start = int(addrstr, 16)

        s.sel = sel
        s.startEA = start
        s.endEA = start+length
        s.bitness = self.addr_mode

        perms = ''
        if self.has_attribute(memory_section, PERMISSIONS):
            perms = self.get_attribute(memory_section, PERMISSIONS)
        s.perm = 0
        if 'r' in perms: s.perm |= idaapi.SEGPERM_READ
        if 'w' in perms: s.perm |= idaapi.SEGPERM_WRITE
        if 'x' in perms: s.perm |= idaapi.SEGPERM_EXEC
        idaapi.add_segm_ex(s, name, "", idaapi.ADDSEG_OR_DIE |
                                            idaapi.ADDSEG_QUIET)
        self.update_counter(MEMORY_SECTION)
        for memory_contents in memory_section.findall(MEMORY_CONTENTS):
            self.import_memory_contents(memory_contents, start, length)
    

    def import_processor(self, processor):
        """
        Processes the PROCESSOR element.
        
        Args:
            processor: PROCESSOR XML element.
        """
        name = self.get_attribute(processor, NAME)
        self.update_counter(PROCESSOR)
        if self.plugin:
            return
        address_model = self.get_attribute(processor, ADDRESS_MODEL)
        if address_model != None:
            flag = idaapi.ph_get_flag()
            if str.lower(address_model) == '16-bit':
                self.addr_mode = 0
            elif str.lower(address_model) == '32-bit':
                self.addr_mode = 1
            elif str.lower(address_model) == '64-bit':
                self.addr_mode = 2
        

    def import_program(self, program):
        """
        Processes the PROGRAM element.
        
        Args:
            program: PROGRAM XML element.
        """
        self.update_status(PROGRAM)
        name = self.get_attribute(program, NAME)
        if self.has_attribute(program, EXE_PATH):
            epath = self.get_attribute(program, EXE_PATH)
            idaapi.set_root_filename(epath)
        else:
            idaapi.set_root_filename(name)
        if self.has_attribute(program, EXE_FORMAT):
            eformat = self.get_attribute(program, EXE_FORMAT)
            RootNode = idaapi.netnode('Root Node')
            if IDA_SDK_VERSION < 650:
                RootNode.supset(1, eformat)
            else:
                RootNode.supset(idaapi.RIDX_FILE_FORMAT_NAME, eformat)
        if self.has_attribute(program, IMAGE_BASE):
            base = self.get_attribute_value(program, IMAGE_BASE)
            idaapi.set_imagebase(base)
        if self.has_attribute(program, INPUT_MD5):
            input_md5 = self.get_attribute(program, INPUT_MD5)
            # store original md5 in a special netnode
            md5 = idaapi.netnode(INPUT_MD5, len(INPUT_MD5), True)
            if (IDA_SDK_VERSION < 650):
                md5.supset(1302, input_md5)
            else:
                md5.supset(idaapi.RIDX_MD5, input_md5)
        self.update_counter(PROGRAM)
        """
        # TODO: this needs to be on "end" event for PROGRAM
        description = program.find(DESCRIPTION)
        if description != None:
            pass  # figure out what to do with it
        """

        
    def import_program_entry_point(self, program_entry_point):
        """
        Defines a program entry point.
        
        Args:
            program_entry_point: PROGRAM_ENTRY_POINT XML element.
                Contains the entry point address.
        """
        if self.Options.EntryPoints.checked == False:
            return
        addr = self.get_address(program_entry_point, ADDRESS)
        idaapi.add_entry(addr, addr, "", True)
        self.update_counter(PROGRAM_ENTRY_POINT)
    

    def import_register_value_range(self, register_value_range):
        """
        Defines the address range for a register value.
        
        Args:
            register_value_range: REGISTER_VALUE_RANGE XML element.
                Contains the register, value, start address and range length.
        """
        if self.Options.RegisterValues.checked == False:
            return
        self.update_counter(REGISTER_VALUE_RANGE)
        reg = self.get_attribute(register_value_range, REGISTER)
        if reg == 'cs': return
        value = self.get_attribute_value(register_value_range, VALUE)
        addr = self.get_address(register_value_range, START_ADDRESS)
        length = self.get_attribute_value(register_value_range, LENGTH)
        r = idaapi.str2reg(reg)
        if r >= idaapi.ph.regFirstSreg and r <= idaapi.ph.regLastSreg:
            if IDA_SDK_VERSION < 670:
                idaapi.splitSRarea1(addr, r, value, 2)
            else:
                idaapi.splitSRarea1(addr, r, value, idaapi.SR_user)
    

    def import_register_var(self, register_var, func):
        """
        Defines a register variable for a function.
        
        Args:
            register_var: REGISTER_VAR XML element.
                Contains register, variable name, and datatype.
            func: IDA function object
        """
        name = self.get_attribute(register_var, NAME)
        reg = self.get_attribute(register_var, REGISTER)
        if self.has_attribute(register_var, DATATYPE):
            datatype = self.get_attribute(register_var, DATATYPE)
        if self.has_attribute(register_var, DATATYPE_NAMESPACE):
            namespace = self.get_attribute(register_var, DATATYPE_NAMESPACE)
        idc.MakeLocal(func.startEA, func.endEA, reg, name)
        self.update_counter(REGISTER_VAR)
    

    def import_stack_frame(self, stack_frame, func):
        """
        Defines a stack frame for a function.
        
        Args:
            stack_frame: STACK_FRAME element with STACK_VAR child elements.
        """
        if self.has_attribute(stack_frame, LOCAL_VAR_SIZE):
            lvsize = self.get_attribute_value(stack_frame, LOCAL_VAR_SIZE)
        if self.has_attribute(stack_frame, PARAM_OFFSET):
            param_offset = self.get_attribute_value(stack_frame, PARAM_OFFSET)
        if self.has_attribute(stack_frame, REGISTER_SAVE_SIZE):
            reg_save_size = self.get_attribute_value(stack_frame,
                                                     REGISTER_SAVE_SIZE)
        if self.has_attribute(stack_frame, RETURN_ADDR_SIZE):
            retaddr_size = self.get_attribute_value(stack_frame,
                                                    RETURN_ADDR_SIZE)
        if self.has_attribute(stack_frame, BYTES_PURGED):
            bytes_purged = self.get_attribute_value(stack_frame, BYTES_PURGED)
        self.update_counter(STACK_FRAME)
        for stack_var in stack_frame.findall(STACK_VAR):
            self.import_stack_var(stack_var, func)
    

    def import_stack_reference(self, stack_reference):
        # import_stack_reference: NOT IMPLEMENTED
        #self.update_counter(STACK_REFERENCE)
        pass
                    

    def import_stack_var(self, stack_var, func):
        """
        Processes STACK_VAR element.
        
        Args:
            stack_var: STACK_VAR XML element.
            
        Stack variables are created by IDA's function analysis. 
        Only the STACK_VAR NAME attribute is used to set the name for
        a stack variable at the specified stack/frame offset. 
        """
        spoffset = self.get_attribute_value(stack_var, STACK_PTR_OFFSET)
        datatype = self.get_attribute(stack_var, DATATYPE)
        offset = spoffset + func.frsize + func.frregs
        if self.has_attribute(stack_var, FRAME_PTR_OFFSET):
            fpoffset = self.get_attribute_value(stack_var, FRAME_PTR_OFFSET)
            offset = fpoffset + func.frsize
        name = ''
        if self.has_attribute(stack_var, NAME):
            name = self.get_attribute(stack_var, NAME)
        if self.has_attribute(stack_var, DATATYPE_NAMESPACE):
            namespace = self.get_attribute(stack_var, DATATYPE_NAMESPACE)
        if self.has_attribute(stack_var, SIZE):
            size = self.get_attribute_value(stack_var, SIZE)
        self.update_counter(STACK_VAR)
        sf = idaapi.get_frame(func)
        if name != '':
            idaapi.set_member_name(sf, offset, name)
    

    def import_structure(self, structure):
        """
        Adds a structure.
        
        Args:
            structure: STRUCTURE XML element.
                Contains the STRUCTURE attributes and child elements.
        """
        if self.Options.DataTypes.checked == False:
            return
        name = self.get_attribute(structure, NAME)
        dtyp = idaapi.get_struc_id(name)
        if dtyp != idaapi.BADNODE:
            # duplicate name, try adding name space
            if self.has_attribute(structure, NAMESPACE) == False:
                return
            namespace = self.get_attribute(structure, NAMESPACE)
            name = namespace + '__' + name
            name.replace('/','_')
            name.replace('.','_')
            dtyp = idaapi.get_struc_id(name)
            # skip if still duplicate (could add sequence #)
            if dtyp != idaapi.BADNODE:
                return
        size = 0
        if self.has_attribute(structure, SIZE):
            size = self.get_attribute_value(structure, SIZE)
        if self.has_attribute(structure, VARIABLE_LENGTH):
            vl = self.get_attribute_value(structure, VARIABLE_LENGTH)
            isVariableLength = vl == 'y'
        sid = idaapi.add_struc(BADADDR, name, False)
        sptr = idaapi.get_struc(sid)
        self.update_counter(STRUCTURE)
        self.import_cmts(structure, sid, STRUCTURE)
        self.import_members(structure, sptr)
        t = idaapi.opinfo_t()
        if idaapi.get_struc_size(sptr) < size:
            idaapi.add_struc_member(sptr,"",size-1,idaapi.byteflag(),t,1)
        

    def import_symbol(self, symbol):
        """
        Adds a symbol name at the specified address.
        
        Args:
            symbol: SYMBOL XML element.
                Contains symbol name and address. Optionally includes
                type and mangled symbol.
        """
        if self.Options.Symbols.checked == False:
            return
        addr = self.get_address(symbol, ADDRESS)
        name = self.get_attribute(symbol, NAME)
        if self.has_attribute(symbol, MANGLED):
            name = self.get_attribute(symbol, MANGLED)
        flag = idaapi.SN_NOWARN
        if self.has_attribute(symbol, TYPE):
            typ = self.get_attribute(symbol, TYPE)
            if  typ == 'local': flag |= idaapi.SN_LOCAL
        idaapi.set_name(addr, name, flag)
        self.update_counter(SYMBOL)
    

    def import_typedef(self, type_def):
        # import_typedef: NOT IMPLEMENTED
        if self.Options.DataTypes.checked == False:
            return
        self.update_counter(TYPE_DEF)
        

    def import_union(self, union):
        """
        Adds a union datatype.
        
        Args:
            union: UNION XML element.
                Contains UNION attributes and child elements.
        """
        if self.Options.DataTypes.checked == False:
            return
        name = self.get_attribute(union, NAME)
        dtyp = idaapi.get_struc_id(name)
        if dtyp != idaapi.BADNODE:
            # duplicate name, try adding name space
            if self.has_attribute(union, NAMESPACE) == False:
                return
            namespace = self.get_attribute(union, NAMESPACE)
            name = namespace + '__' + name
            name.replace('/','_')
            name.replace('.','_')
            dtyp = idaapi.get_struc_id(name)
            # skip if still duplicate (could add sequence #)
            if dtyp != idaapi.BADNODE:
                return
        size = 0
        if self.has_attribute(union, SIZE):
            size = self.get_attribute_value(union, SIZE)
        sid = idaapi.add_struc(BADADDR, name, True)
        sptr = idaapi.get_struc(sid)
        self.update_counter(UNION)
        self.import_cmts(union, sid, UNION)
        self.import_members(union, sptr)
        t = idaapi.opinfo_t()
        if idaapi.get_struc_size(sptr) < size:
            idaapi.add_struc_member(sptr,"", size-1, idaapi.byteflag(), None, 1)
            

    def open_file(self, filename, mode):
        """
        Opens filename to specified mode.
        
        Args:
            filename: String representing absolute filepath.
            mode: String representing mode for open.
            
        Returns
            File handle.
        
        Exceptions:
            Displays a warning and raises FileError exception
            if open fails.
        """
        try:
            f = open(filename, mode)
            return f
        except:
            fmt = "TITLE ERROR!\n"
            fmt += "ICON ERROR\n"
            fmt += "AUTOHIDE NONE\n"
            fmt += "HIDECANCEL\n"
            fmt += "Error opening file" + filename + "!\n"
            idaapi.warning(fmt)
            raise FileError
    

    def update_counter(self, tag):
        """
        Updates the counter for the element tag.
        
        Args:
            tag: String representing element tag.
        """
        if tag in self.Elements:
            self.Counters[self.Elements[tag]] += 1
        else:
            self.Elements[tag] = len(self.Elements)
            self.Counters.append(1)
            self.Tags.append(tag)
    

    def update_import(self, element):
        """
        Update the element counter and processing status.
        
        Args:
            element: XML element
            
        This function is used to process certain high-level elements
        (such as COMMENTS, CODE_BLOCKS, SYMBOL_TABLE, FUNCTIONS, etc.)
        that are used to group sub-elements.
        """
        self.update_counter(element.tag)
        self.update_status(element.tag)
    

    def update_status(self, tag):
        """
        Displays the processing status in the IDA window.
        
        Args:
            tag: String representing XML element tag
        """
        status = 'Importing ' + tag
        idaapi.msg('\n%-35s' % status)
        idaapi.hide_wait_box()
        idaapi.show_wait_box(status)
    
        
# Global XML string constants for elements and attributes
ADDRESS = 'ADDRESS'
ADDRESS_MODEL = 'ADDRESS_MODEL'
ADDRESS_RANGE = 'ADDRESS_RANGE'
BASE_ADDRESS = 'BASE_ADDRESS'
BIT_FIELD = 'BIT_FIELD'
BIT_MAPPED = 'BIT_MAPPED'
BIT_MASK = 'BIT_MASK'
BOOKMARK = 'BOOKMARK'
BOOKMARKS = 'BOOKMARKS'
BYTES = 'BYTES'
BYTES_PURGED = 'BYTES_PURGED'
CATEGORY = 'CATEGORY'
CODE = 'CODE'
CODE_BLOCK = 'CODE_BLOCK'
COMMENT = 'COMMENT'
COMMENTS = 'COMMENTS'
COMPILER = 'COMPILER'
DATA = 'DATA'
DATATYPE = 'DATATYPE'
DATATYPES = 'DATATYPES'
DATATYPE_NAMESPACE = 'DATATYPE_NAMESPACE'
DEFINED_DATA = 'DEFINED_DATA'
DESCRIPTION = 'DESCRIPTION'
DISPLAY_SETTINGS = 'DISPLAY_SETTINGS'
END = 'END'
ENDIAN = 'ENDIAN'
ENTRY_POINT = 'ENTRY_POINT'
ENUM = 'ENUM'
ENUM_ENTRY = 'ENUM_ENTRY'
EQUATE = 'EQUATE'
EQUATES = 'EQUATES'
EQUATE_GROUP = 'EQUATE_GROUP'
EQUATE_REFERENCE = 'EQUATE_REFERENCE'
EXE_FORMAT = 'EXE_FORMAT'
EXE_PATH = 'EXE_PATH'
EXT_LIBRARY = 'EXT_LIBRARY'
EXT_LIBRARY_REFERENCE = 'EXT_LIBRARY_REFERENCE'
EXT_LIBRARY_TABLE = 'EXT_LIBRARY_TABLE'
FAMILY = 'FAMILY'
FILE = 'FILE'
FILE_NAME = 'FILE_NAME'
FILE_OFFSET = 'FILE_OFFSET'
FOLDER = 'FOLDER'
FORMAT = 'FORMAT'
FRAGMENT = 'FRAGMENT'
FRAME_PTR_OFFSET = 'FRAME_PTR_OFFSET'
FUNCTION = 'FUNCTION'
FUNCTIONS = 'FUNCTIONS'
FUNCTION_DEF = 'FUNCTION_DEF'
IMAGE_BASE = 'IMAGE_BASE'
INPUT_MD5 = 'INPUT_MD5'
INFO_SOURCE = 'INFO_SOURCE'
LANGUAGE_PROVIDER = 'LANGUAGE_PROVIDER'
LENGTH = 'LENGTH'
LIB_ADDR = 'LIB_ADDR'
LIB_LABEL = 'LIB_LABEL'
LIB_ORDINAL = 'LIB_ORDINAL'
LIB_PROG_NAME = 'LIB_PROG_NAME'
LIBRARY_FUNCTION = 'LIBRARY_FUNCTION'
LOCAL_VAR_SIZE = 'LOCAL_VAR_SIZE'
MANGLED = 'MANGLED'
MANUAL_INSTRUCTION = 'MANUAL_INSTRUCTION'
MANUAL_OPERAND = 'MANUAL_OPERAND'
MARKUP = 'MARKUP'
MEMBER = 'MEMBER'
MEMORY_CONTENTS = 'MEMORY_CONTENTS'
MEMORY_MAP = 'MEMORY_MAP'
MEMORY_REFERENCE = 'MEMORY_REFERENCE'
MEMORY_SECTION = 'MEMORY_SECTION'
NAME = 'NAME'
NAMESPACE = 'NAMESPACE'
OFFSET = 'OFFSET'
OPERAND_INDEX = 'OPERAND_INDEX'
PARAM_OFFSET = 'PARAM_OFFSET'
PATH = 'PATH'
PERMISSIONS = 'PERMISSIONS'
PRIMARY = 'PRIMARY'
PROCESSOR = 'PROCESSOR'
PROGRAM = 'PROGRAM'
PROGRAM_ENTRY_POINT = 'PROGRAM_ENTRY_POINT'
PROGRAM_ENTRY_POINTS = 'PROGRAM_ENTRY_POINTS'
PROGRAM_TREES = 'PROGRAM_TREES'
PROPERTIES = 'PROPERTIES'
PROPERTY = 'PROPERTY'
REGISTER = 'REGISTER'
REGISTER_SAVE_SIZE = 'REGISTER_SAVE_SIZE'
REGISTER_VALUES = 'REGISTER_VALUES'
REGISTER_VALUE_RANGE = 'REGISTER_VALUE_RANGE'
REGISTER_VAR = 'REGISTER_VAR'
REGULAR_CMT = 'REGULAR_CMT'
RELOCATION = 'RELOCATION'
RELOCATION_TABLE = 'RELOCATION_TABLE'
REPEATABLE_CMT = 'REPEATABLE_CMT'
RETURN_ADDR_SIZE = 'RETURN_ADDR_SIZE'
RETURN_TYPE = 'RETURN_TYPE'
SHOW_TERMINATOR = 'SHOW_TERMINATOR'
SIGNED = 'SIGNED'
SIZE = 'SIZE'
SOURCE_ADDRESS = 'SOURCE_ADDRESS'
SOURCE_TYPE = 'SOURCE_TYPE'
STACK_FRAME = 'STACK_FRAME'
STACK_PTR_OFFSET = 'STACK_PTR_OFFSET'
STACK_REFERENCE = 'STACK_REFERENCE'
STACK_VAR = 'STACK_VAR'
START = 'START'
START_ADDR = 'START_ADDR'
START_ADDRESS = 'START_ADDRESS'
STRUCTURE = 'STRUCTURE'
SYMBOL = 'SYMBOL'
SYMBOL_TABLE = 'SYMBOL_TABLE'
TIMESTAMP = 'TIMESTAMP'
TOOL = 'TOOL'
TO_ADDRESS = 'TO_ADDRESS'
TREE = 'TREE'
TYPE = 'TYPE'
TYPEINFO_CMT = 'TYPEINFO_CMT'
TYPE_DEF = 'TYPE_DEF'
UNION = 'UNION'
USER = 'USER'
USER_DEFINED = 'USER_DEFINED'
VALUE = 'VALUE'
VARIABLE_LENGTH = 'VARIABLE_LENGTH'
ZERO_PAD = 'ZERO_PAD'


repository: MassinissaDjellouli/CTFStuff
file: tools/ghidra_11.2.1_PUBLIC/Extensions/IDAPro/Python/6xx/loaders/xmlldr.py
wanted_name: XML Importer
comment: Import XML PROGRAM file
created_at: 2024-01-16 04:53:55+00:00
pushed_at: 2025-05-13 04:26:11+00:00
url:
https://github.com/MassinissaDjellouli/CTFStuff/blob/8be397cdd7b2b33bde74f33314fdcaf7244d36cb/tools/ghidra_11.2.1_PUBLIC/Extensions/IDAPro/Python/6xx/loaders/xmlldr.py

content:
## ###
#  IP: GHIDRA
# 
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#  
#       http://www.apache.org/licenses/LICENSE-2.0
#  
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
##
#---------------------------------------------------------------------
# xmlldr.py - IDA XML Import loader and plugin
#---------------------------------------------------------------------
"""
Loader and plugin for IDA to import a XML PROGRAM file to a database.
"""

import idaapi
import sys
import time
from xml.etree import cElementTree


XML_IMPORTER_VERSION = "2.1.1"
BASELINE_IDA_VERSION = 620
IDA_SDK_VERSION = idaapi.IDA_SDK_VERSION
BADADDR = idaapi.BADADDR


"""
Loader functions
"""
def accept_file(li, n):
    """
    Check if the file is of supported format

    @param li: a file-like object which can be used to access the input data
    @param n : format number. The function will be called with incrementing
                number until it returns zero
    @return: 0 - no more supported formats
                string "name" - format name to display in the chooser dialog
                dictionary { 'format': "name", 'options': integer }
                options: should be 1, possibly ORed with ACCEPT_FIRST (0x8000)
                to indicate preferred format
    """
    if IDA_SDK_VERSION < BASELINE_IDA_VERSION:
        return 0
    # we support only one format per file
    if n > 0: return 0
    # read 16K bytes to allow for the DTD
    data = li.read(0x4000)
    # look for start of <PROGRAM> element
    start = data.find("<PROGRAM")
    if start >= 0: return "XML PROGRAM file"
    return 0


def load_file(li, neflags, format):
    """
    Load the file into database

    @param li: a file-like object which can be used to access the input data
    @param neflags: options selected by the user, see loader.hpp
    @return: 0-failure, 1-ok
    """
    global event, element
    plugin = False
    status = 0
    st = idaapi.setStat(idaapi.st_Work)
    xml = XmlImporter(False, 0)
    try:
        status = xml.import_xml()
    except Cancelled:
        msg = "XML PROGRAM import cancelled!"
        print "\n" + msg
        idaapi.warning(msg)
    except MultipleAddressSpacesNotSupported:
        msg  = "XML Import cancelled!"
        msg += "\n\nXML Import does not currently support"
        msg += "\nimporting multiple address spaces."
        print "\n" + msg
        idaapi.warning(msg)
    except:
        print "\nHouston, we have a problem!"
        msg = "***** Exception occurred: XML loader failed! *****"
        print "\n" + msg + "\n", sys.exc_type, sys.exc_value
        print event, element.tag, element.attrib
        idaapi.warning(msg)
    finally:
        idaapi.setStat(st)
        xml.cleanup()
        return status


class XmlImporterPlugin(idaapi.plugin_t):
    """
    XML Importer plugin class
    """
    flags = 0
    comment = "Import XML PROGRAM file"
    help = "Import XML <PROGRAM> document to database"
    wanted_name = "XML Importer"
    wanted_hotkey = "Ctrl-Alt-l"

    def init(self):
        """
        init function for XML Importer plugin.
        
        Returns:
            Constant PLUGIN_OK if this IDA version supports the plugin,
            else returns PLUGIN_SKIP if this IDA is older than the supported
            baseline version.
        """
        if IDA_SDK_VERSION < BASELINE_IDA_VERSION:
            idaapi.msg('\nXML Importer plugin (xmlldr.py) not supported ' +
                        'by this version of IDA\n')
            return idaapi.PLUGIN_SKIP
        return idaapi.PLUGIN_OK

    def run(self, arg):
        """
        run function for XML Importer plugin.
        
        Args:
            arg: Integer, a non-zero value enables auto-run feature for
                 IDA batch (no gui) processing mode. Default is 0.
        """
        st = idaapi.setStat(idaapi.st_Work)
        xml = XmlImporter(True, arg)
        try:
            try:
                xml.import_xml()
            except Cancelled:
                msg = "XML Import cancelled!"
                print "\n" + msg
                idaapi.warning(msg)
            except MultipleAddressSpacesNotSupported:
                msg  = "XML Import cancelled!"
                msg += "\n\nXML Import does not currently support"
                msg += "\nimporting multiple address spaces."
                print "\n" + msg
                idaapi.warning(msg)
            except:
                msg = "***** Exception occurred: XML Importer failed! *****"
                print "\n" + msg + "\n", sys.exc_type, sys.exc_value
                idaapi.warning(msg)
        finally:
            xml.cleanup()
            idaapi.setStat(st)


    def term(self):
        pass


def PLUGIN_ENTRY():
    return XmlImporterPlugin()


class Cancelled(Exception):
    pass


class FileError(Exception):
    pass


class MultipleAddressSpacesNotSupported(Exception):
    pass


class XmlImporter:
    """
    XmlImporter class contains methods to import an XML PROGRAM
        document into IDA.
    """
    def __init__(self, plugin, arg):
        """
        Initializes the XmlImporter attributes

        Args:
            arg: Integer, non-zero value enables auto-run feature for
                IDA batch (no gui) processing mode. Default is 0.
        """
        self.plugin = plugin
        self.autorun = False
        self.debug = False # set this to True for debug prints
        self.Elements = {}
        self.Counters = []
        self.Tags = []
        self.Timers = dict()
        self.xmlfile = 0
        self.addr_mode = 1
        self.create = True
        self.dataseg = None
        self.deferred = []
        self.callbacks = {
            'start' : {
                BOOKMARKS           : self.update_import,
                CODE                : self.update_import,
                COMMENTS            : self.update_import,
                COMPILER            : self.import_compiler,
                DATA                : self.update_import,
                DATATYPES           : self.update_import,
                EQUATES             : self.update_import,
                FUNCTIONS           : self.update_import,
                INFO_SOURCE         : self.import_info_source,
                MARKUP              : self.update_import,
                MEMORY_MAP          : self.import_memory_map,
                PROCESSOR           : self.import_processor,
                PROGRAM             : self.import_program,
                PROGRAM_ENTRY_POINTS: self.update_import,
                REGISTER_VALUES     : self.update_import,
                SYMBOL_TABLE        : self.update_import },
            'end' : {
                BOOKMARK            : self.import_bookmark,
                CODE_BLOCK          : self.import_codeblock,
                COMMENT             : self.import_comment,
                DEFINED_DATA        : self.import_defined_data,
                ENUM                : self.import_enum,
                EQUATE_GROUP        : self.import_equate_group,
                EQUATE_REFERENCE    : self.import_equate_reference,
                FUNCTION            : self.import_function,
                FUNCTION_DEF        : self.import_function_def,
                MANUAL_INSTRUCTION  : self.import_manual_instruction,
                MANUAL_OPERAND      : self.import_manual_operand,
                MEMORY_REFERENCE    : self.import_memory_reference,
                MEMORY_SECTION      : self.import_memory_section,
                PROGRAM_ENTRY_POINT : self.import_program_entry_point,
                REGISTER_VALUE_RANGE: self.import_register_value_range,
                STACK_REFERENCE     : self.import_stack_reference,
                STRUCTURE           : self.import_structure,
                SYMBOL              : self.import_symbol,
                TYPE_DEF            : self.import_typedef,
                UNION               : self.import_union,
                # end element for elapse time
                BOOKMARKS           : self.display_timer,
                CODE                : self.display_timer,
                COMMENTS            : self.display_timer,
                DATA                : self.display_timer,
                DATATYPES           : self.process_deferred,
                EQUATES             : self.display_timer,
                FUNCTIONS           : self.display_timer,
                MARKUP              : self.display_timer,
                MEMORY_MAP          : self.display_timer,
                PROGRAM             : self.display_total_time,
                PROGRAM_ENTRY_POINTS: self.display_timer,
                REGISTER_VALUES     : self.display_timer,
                SYMBOL_TABLE        : self.display_timer }
            }


    def import_xml(self):
        """
        Imports the XML PROGRAM file into the database.
        """
        self.display_xml_importer_version()
        displayMenu = self.plugin == True and self.autorun == False 
        #displayMenu = self.autorun == False 
        self.get_options(displayMenu)
        if self.plugin:
            self.filename=idaapi.askfile_c(0, "*.xml",
                                        "Enter name of xml file:")
        else:
            self.filename = idaapi.get_input_file_path()
        if (len(self.filename) == 0):
            return
        idaapi.msg('\nImporting from: ' + self.filename + '\n') 
        idc.Wait()
        if self.plugin == False: idaapi.hide_wait_box()
        idaapi.show_wait_box("Importing XML PROGRAM document....")
        n = 0
        for event,element in cElementTree.iterparse(self.filename,
                                                events=("start","end")):
            if idaapi.wasBreak() == True:
                raise Cancelled
            
            if self.debug == True and event == 'start':
                print element.tag, element.attrib, element.text
            
            if event in self.callbacks:
                if element.tag in self.callbacks[event]:
                    if event == 'start':
                        self.Timers[element.tag] = time.clock()
                    self.callbacks[event][element.tag](element)
                    if event == 'end':
                        element.clear()
            if event == 'end':
                n += 1
        end = time.clock()
        idaapi.hide_wait_box()
        self.display_summary()
        idaapi.msg('\nXML Elements parsed: ' + str(n) + '\n\n')
        return 1
    

    def get_options(self, display):
        """
        Displays the options menu and retrieves the option settings. 
        """
        fmt =  "HELP\n"
        fmt += "XML PROGRAM loader/importer plugin (Python)\n"
        fmt += "IDA SDK: "+ str(IDA_SDK_VERSION) + "\n\n"
        fmt +=  "The XML PROGRAM loader loads elements from a "
        fmt +=  "XML <PROGRAM> document to create an idb database.\n\n"
        fmt +=  "ENDHELP\n"
        fmt +=  "Import from XML PROGRAM document...."
        fmt += "\n <##Options##Code Blocks:{CodeBlocks}>"
        fmt += "\n <Entry Points:{EntryPoints}>"
        fmt += "\n <Segment Register Value Ranges:{RegisterValues}>"
        fmt += "\n <Data Types:{DataTypes}>"
        fmt += "\n <Data Definitions:{DataDefinitions}>"
        fmt += "\n <Symbols:{Symbols}>"
        fmt += "\n <Comments:{Comments}>"
        fmt += "\n <Bookmarks:{Bookmarks}>"
        fmt += "\n <Functions:{Functions}>"
        fmt += "\n <Memory References:{MemoryReferences}>"
        fmt += "\n <Equate/Enum References:{EquateReferences}>"
        fmt += "\n <Manual Instructions/Operands:{Manual}>{cGroup1}>"
        fmt += "\n\n"
    
        Opts = { 'cGroup1': idaapi.Form.ChkGroupControl(
                ( 
                "CodeBlocks", "EntryPoints", "RegisterValues",
                "DataTypes", "DataDefinitions",
                "Symbols", "Comments", "Bookmarks",
                "Functions", "MemoryReferences", "EquateReferences",
                "Manual")) }
        
        self.Options = idaapi.Form(fmt, Opts)
        self.Options.Compile()
    
        self.Options.CodeBlocks.checked       = True
        self.Options.EntryPoints.checked      = True
        self.Options.RegisterValues.checked   = True
        self.Options.DataTypes.checked        = True
        self.Options.DataDefinitions.checked  = True
        self.Options.Symbols.checked          = True
        self.Options.Functions.checked        = True
        self.Options.Comments.checked         = True
        self.Options.Bookmarks.checked        = True
        self.Options.MemoryReferences.checked = True
        self.Options.EquateReferences.checked = True
        self.Options.Manual.checked           = True
        
        if display == True:
            ok = self.Options.Execute()
            if (ok == 0):
                raise Cancelled
            
    def cleanup(self):
        """
        Frees memory and closes message box at termination.
        """
        if self.plugin:
            self.Options.Free()
        idaapi.hide_wait_box()
    

    def display_timer(self, element):
        """
        Displays the elapsed processing time for XML elements.
        
        Args:
            element: XML element object value containing the element tag.
        """
        if element.tag == MEMORY_MAP and self.plugin:
            return
        if element.tag in self.Timers:
            idaapi.msg('elapsed time: %.4f' %
                    (time.clock()-self.Timers[element.tag]))
    

    def display_total_time(self, element):
        """
        Displays the total processing time.
        
        Args:
            element: XML element object value (not used).
        """
        TOTAL = 'Total '
        idaapi.msg('\n%35selapsed time: %.4f' %
                (TOTAL,time.clock()-self.Timers[PROGRAM]))
    

    def display_summary(self):
        """
        Displays summary of the XML PROGRAM import in IDA output window.
        """
        summary = ''
        total = 0
        for tag in self.Tags:
            count = self.Counters[self.Elements[tag]]
            summary += "\n%s: %d" % (tag, count)
            total += count
        summary += '\n--------------------------------------'
        header = '\n--------------------------------------'
        header += ('\nTotal XML Elements Processed: %d') % total
        summary = header + summary
        idaapi.msg(summary)
        if self.plugin and self.autorun == False:
            frmt  = "TITLE XML Import Successful!\n"
            frmt += "ICON INFO\n"
            frmt += "AUTOHIDE NONE\n"
            frmt += "HIDECANCEL\n"
            firstline = "\n XML IMPORT SUCCESSFUL!"
            fileline = '\n\nFile: %s' % self.filename
            details = '\n\nSee output window for details...'
            idaapi.info("%s" % (frmt + firstline + fileline + details))
    

    def display_xml_importer_version(self):
        """
        Displays XML Importer plugin version info in IDA output window.
        """
        if self.plugin:
            f = idaapi.idadir(idaapi.PLG_SUBDIR) + '/xmlldr.py'
        else:
            f = idaapi.idadir(idaapi.LDR_SUBDIR) + '/xmlldr.py'
        plugintime = time.localtime(os.path.getmtime(f))
        ts = time.strftime('%b %d %Y %H:%M:%S', plugintime)
        version = "\nXML Importer Version " + XML_IMPORTER_VERSION
        version += " : SDK " + str(IDA_SDK_VERSION)
        version += " : Python : "+ ts + '\n'
        idaapi.msg(version)
    

    def get_address(self, element, attr):
        """
        Returns the address value for an element.
        
        Args:
            element: XML element object.
            attr: String containing the address attribute name.
            
        Returns:
            Numeric value representing the address.
        """
        addrstr = element.get(attr)
        if '::' in addrstr:
            # overlayed addresses not currently handled
            return BADADDR
        elif ':' in addrstr:
            [segstr, offset_str] = string.split(addrstr,':')
            offset = int(offset_str,16)
            if self.is_int(segstr) == True:
                sgmt = int(segstr,16)
                addr = (sgmt << 4) + offset
            else:
                # multiple address spaces not currently implemented
                addr = BADADDR
            return addr
        else:
            return int(element.get(attr), 16)
    

    def get_attribute(self, element, attr):
        """
        Returns the attribute value string.
        
        Args:
            element: XML element object.
            attr: String containing the attribute name.
            
        Returns:
            String representing the attribute value.
        """
        return element.get(attr)
    

    def get_attribute_value(self, element, attr):
        """
        Returns the numeric attribute value.
        
        Args:
            element: XML element object.
            attr: String containing the attribute name.
            
        Returns:
            Numeric value representing the attribute value.
        """
        val = element.get(attr)
        try:
            if val.upper().startswith('0X') or val.upper().startswith('-0X'):
                return int(val, 16)
            return int(val)
        except:
            idaapi.msg('\nUnable to decode string as value: ' + val)
            return 0
        

    def get_cbsize(self):
        """
        Returns the size of the addressable codebyte for the processor.
        
        Returns:
            Integer representing the number of 8-bit bytes in an
            addressable codebyte.
        """
        return (idaapi.ph_get_cnbits()+7)/8
    

    def get_datatype_flags(self, datatype, size):
        """
        Returns the flags bitmask for the datatype.
        
        Args:
            datatype: String representing the datatype.
            size: Integer representing the datatype size.

        Returns:
            Integer representing the bitmask.
        """
        if datatype.lower().startswith("byte"):     return idaapi.byteflag()
        if datatype.lower().startswith("word"):     return idaapi.wordflag()
        if datatype.lower().startswith("dword"):    return idaapi.dwrdflag()
        if datatype.lower().startswith("qword"):    return idaapi.qwrdflag()
        if datatype.lower().startswith("oword"):    return idaapi.owrdflag()
        if datatype.lower().startswith("tbyte"):    return idaapi.tbytflag()
        if datatype.lower().startswith("float"):    return idaapi.floatflag()
        if datatype.lower().startswith("double"):   return idaapi.doubleflag()
        if datatype.lower().startswith("packed"):   return idaapi.packrealflag()
        if self.is_string_type(datatype):           return idaapi.asciflag()
        if self.is_enumeration(datatype):           return idaapi.enumflag()
        if self.is_structure(datatype):             return idaapi.struflag()
        #if size == 4:                               return idaapi.dwrdflag()
        return 0
    

    def get_string_type(self, datatype):
        if datatype.lower() == 'mbcstring':
            return idaapi.ASCSTR_UNICODE
        if datatype.lower().find('unicode') != -1:
            if datatype.lower().find('pascal') != -1:
                return idaapi.ASCSTR_ULEN2
            return idaapi.ASCSTR_UNICODE
        if datatype.lower().find('pascal') != -1:
            return idaapi.ASCSTR_LEN2
        return idaapi.ASCSTR_TERMCHR
            
    
    def has_attribute(self, element, attr):
        """
        Returns true if the XML element contains the named attribute.
        
        Args:
            element: XML element object
            attr: String containing name of the attribute
        
        Returns:
            True if the element contains the named attribute, otherwise False.
        """
        return attr in element.attrib
    

    def is_enumeration(self, datatype):
        """
        Returns true if datatype is an existing enumeration in the database.
        
        Args:
            datatype: String representing the datatype.
            
        Returns:
            True if the datatype is an enumeration in the database,
            otherwise False.
        """
        if idaapi.get_enum(datatype) == idaapi.BADNODE:  return False
        return True
    
    
    def is_int(self, s):
        try:
            int(s, 16)
            return True
        except:
            return False
        

    def is_pointer_type(self, dtype):
        """
        Returns true if the datatype represents a pointer.
        
        Args:
            dtype: String representing the datatype.
            
        Returns:
            True if the datatype represents a pointer, otherwise False.
        """
        if dtype.lower().startswith("pointer") or dtype.endswith('*'):
            return True
        return False
    

    def is_string_type(self, datatype):
        """
        Returns true if the datatype represents a string type.
        
        Args:
            datatype: String representing the datatype.
            
        Returns:
            True if the datatype represents a string, otherwise False.
        """
        if datatype.lower().startswith("unicode"):  return True
        if datatype.lower().startswith("string"):   return True
        return False
    

    def is_structure(self, datatype):
        """
        Returns true if the datatype represents a structure in the database.
        
        Args:
            dtype: String representing the datatype.
            
        Returns:
            True if the datatype represents an existing structure,
            otherwise False.
        """
        if idaapi.get_struc_id(datatype) == idaapi.BADNODE:  return False
        return True
    

    def import_address_range(self, address_range):
        """
        Processes ADDRESS_RANGE element.
        
        Args:
            address_range: XML element object containing start and end address
                attributes for the address range.
                
        Returns:
            Tuple containing two integers, the start and end address values.
        """
        start = self.get_address(address_range,START)
        end = self.get_address(address_range, END)
        self.update_counter(ADDRESS_RANGE)
        return (start, end)
    

    def import_bit_mask(self, bitmask, eid):
        """
        Processes a BIT_MASK element as an enum bitmask member.
        
        Args:
            bitmask: XML element object representing the IDA enum bitmask.
            eid: Integer representing the IDA enum id
        """
        if IDA_SDK_VERSION < 640:
            return
        name = self.get_attribute(bitmask,NAME)
        value = self.get_attribute_value(bitmask,VALUE)
        idaapi.set_bmask_name(eid, value, name)
        cid = idaapi.get_const_by_name(name)
        self.update_counter(BIT_MASK)
        regcmt = bitmask.find(REGULAR_CMT)
        if regcmt != None:
            idaapi.set_const_cmt(cid, regcmt.text, False);
            self.update_counter(BIT_MASK + ':' + REGULAR_CMT)
        rptcmt = bitmask.find(REPEATABLE_CMT)
        if rptcmt != None:
            idaapi.set_const_cmt(cid, rptcmt.txt, True);
            self.update_counter(BIT_MASK + ':' + REPEATABLE_CMT)
    

    def import_bookmark(self, bookmark):
        """
        Processes a BOOKMARK element.
        
        Args:
            bookmark: XML element object containing bookmark data.
        """
        if self.Options.Bookmarks.checked == False:
            return
        try:
            addr = self.get_address(bookmark, ADDRESS)
            if self.has_attribute(bookmark, TYPE):
                typ = self.get_attribute(bookmark, TYPE)
            category = ''
            if self.has_attribute(bookmark, CATEGORY):
                category = self.get_attribute(bookmark, CATEGORY)
            description = ''
            if self.has_attribute(bookmark, DESCRIPTION):
                description = self.get_attribute(bookmark, DESCRIPTION)
            if idaapi.isEnabled(addr) == False:
                msg = ("import_bookmark: address %X not enabled in database"
                       % addr)
                print msg
                return
            self.update_counter(BOOKMARK)
            for slot in range(1,1025):
                if IDA_SDK_VERSION != 695:
                    ea = idc.GetMarkedPos(slot)
                else:
                    import ida_moves
                    import ida_pro
                    curloc = ida_moves.curloc()
                    intp = ida_pro.int_pointer()
                    intp.assign(slot)
                    ea = curloc.markedpos(intp)
                if ea == BADADDR:
                    curloc = idaapi.curloc()
                    curloc.ea = addr
                    curloc.mark(slot, category, description)
                    break
        except:
            msg = "** Exception occurred in import_bookmark **"
            print "\n" + msg + "\n", sys.exc_type, sys.exc_value
    

    def import_cmts(self, element, sid, typ):
        """
        Processes REGULAR_CMT and REPEATABLE_CMT elements for structures.
        
        Args:
            element: XML element object containing a REGULAR_CMT or
                REPEATABLE_CMT element
            sid: Integer representing the structure id
            typ: String indicating structure type (STRUCTURE or UNION)
        """
        regcmt = element.find(REGULAR_CMT)
        if regcmt != None:
            idaapi.set_struc_cmt(sid, regcmt.text, False)
            self.update_counter(typ + ':' + REGULAR_CMT)
        rptcmt = element.find(REPEATABLE_CMT)
        if rptcmt != None:
            idaapi.set_struc_cmt(sid, rptcmt.text, True)
            self.update_counter(typ + ':' + REPEATABLE_CMT)
            
        
    def import_codeblock(self, code_block):
        """
        Processes a CODE_BLOCK element by disassembling the address range.
        
        Args:
            code_block: XML element containing codeblock start and end
                addresses.
        """
        if self.Options.CodeBlocks.checked == False:
            return
        start = self.get_address(code_block, START)
        end = self.get_address(code_block, END)
        idaapi.do_unknown_range(start, end-start+1, 3)
        addr = start
        while (addr <= end):
            length = idaapi.create_insn(addr)
            addr += idaapi.get_item_size(addr) * self.get_cbsize()
        self.update_counter(CODE_BLOCK)
    

    def import_comment(self, comment):
        """
        Processes a COMMENT element by creating the comment at the address.
        
        Args:
            comment: XML element containing the comment address, type,
                and text.
        """
        if self.Options.Comments.checked == False:
            return
        addr = self.get_address(comment, ADDRESS)
        ctype = self.get_attribute(comment,TYPE)
        text = comment.text
        if ctype == 'pre':
            idaapi.add_long_cmt(addr, True, text)
        elif ctype == 'end-of-line':
            idaapi.set_cmt(addr, text, False)
        elif ctype == 'repeatable':
            idaapi.set_cmt(addr, text, True)
        elif ctype == 'post':
            idaapi.add_long_cmt(addr, False, text)
        self.update_counter(COMMENT+':' + ctype)
    

    def import_compiler(self, compiler):
        """
        Processes the COMPILER element containing the compiler name.
        
        Args:
            compiler: XML element containing the compiler name.
        """
        name = self.get_attribute(compiler, NAME)
        comp = idaapi.COMP_UNK
        if   name == "Visual C++":      comp = idaapi.COMP_MS
        elif name == "Borland C++":     comp = idaapi.COMP_BC
        elif name == "Watcom C++":      comp = idaapi.COMP_WATCOM
        elif name == "GNU C++":         comp = idaapi.COMP_GNU
        elif name == "Visual Age C++":  comp = idaapi.COMP_VISAGE
        elif name == "Delphi":          comp = idaapi.COMP_BP
        idaapi.cvar.inf.cc.id = comp
        self.update_counter(COMPILER)
    

    def import_defined_data(self, defined_data):
        """
        Processes a DEFINED_DATA element by creating a data item at the
            specified address.
        
        Args:
            defined_data: XML element containing the address and
                datatype information for the data item
        """
        if self.Options.DataDefinitions.checked == False:
            return
        addr = self.get_address(defined_data, ADDRESS)
        datatype = self.get_attribute(defined_data, DATATYPE)
        size = self.get_attribute_value(defined_data, SIZE)
        self.update_counter(DEFINED_DATA)        
        ti = idaapi.opinfo_t()
        if self.is_pointer_type(datatype):
            #idaapi.set_refinfo(ti, 0, 0, 0, REF_OFF32)
            flag = idaapi.dwrdflag() | idaapi.FF_0OFF
            #idaapi.set_typeinfo(addr, 0, flag, ti)
        else:
            flag = self.get_datatype_flags(datatype, size)
        if flag == idaapi.asciflag():
            idaapi.make_ascii_string(addr, size,
                                     self.get_string_type(datatype))
        elif flag == idaapi.struflag():
            idaapi.doStruct(addr, size, idaapi.get_struc_id(datatype))
        else:
            idaapi.do_data_ex(addr, flag, size, idaapi.BADNODE)        
        typecmt = defined_data.find(TYPEINFO_CMT)
        if typecmt != None:
            self.update_counter(DEFINED_DATA + ':' + TYPEINFO_CMT)


    def import_enum(self, enum):
        """
        Processes an ENUM element by creating the enumeration.
        
        Args:
            enum: XML element containing the enumeration name and
                member data.
        """
        if self.Options.DataTypes.checked == False:
            return
        name = self.get_attribute(enum, NAME)
        if self.has_attribute(enum,NAMESPACE):
            namespace = self.get_attribute(enum, NAMESPACE)
        if self.has_attribute(enum,SIZE):
            size = self.get_attribute_value(enum, SIZE)
        eid = idaapi.add_enum(BADADDR, name,
                              idaapi.hexflag() | idaapi.dwrdflag())
        self.update_counter(ENUM)
        regcmt = enum.find(REGULAR_CMT)
        if regcmt != None:
            idaapi.set_enum_cmt(eid, regcmt.text, False)
            self.update_counter(ENUM + ':' + REGULAR_CMT)
        rptcmt = enum.find(REPEATABLE_CMT)
        if rptcmt != None:
            idaapi.set_enum_cmt(eid, rptcmt.text, True)
            self.update_counter(ENUM + ':' + REPEATABLE_CMT)
        display_settings = enum.find(DISPLAY_SETTINGS)
        if display_settings != None:
            self.update_counter(ENUM + ':' + DISPLAY_SETTINGS)
        enum_entries = enum.findall(ENUM_ENTRY)
        for enum_entry in enum_entries:
            self.import_enum_entry(enum_entry, eid)
    

    def import_enum_entry(self, enum_entry, eid):
        """
        Processes an ENUM_ENTRY by creating a member in the enumeration.
        
        Args:
            enum_entry: XML element containing the member name and value.
            eid: Integer representing the id of the enumeration.
        """
        name = self.get_attribute(enum_entry, NAME)
        value = self.get_attribute_value(enum_entry, VALUE)
        if IDA_SDK_VERSION > 630:
            idaapi.add_const(eid, name, value)
        else:
            idaapi.add_enum_member(eid, name, value)
        if IDA_SDK_VERSION > 630:
            cid = idaapi.get_const_by_name(name)
        else:
            cid = idaapi.get_enum_member_by_name(name)
        self.update_counter(ENUM_ENTRY)
        regcmt = enum_entry.find(REGULAR_CMT)
        if regcmt != None:
            if IDA_SDK_VERSION > 630:
                idaapi.set_const_cmt(cid, regcmt.text, False);
            else:
                idaapi.set_enum_member_cmt(cid, regcmt.text, False);
            self.update_counter(ENUM_ENTRY + ':' + REGULAR_CMT)
        rptcmt = enum_entry.find(REPEATABLE_CMT)
        if rptcmt != None:
            if IDA_SDK_VERSION > 630:
                idaapi.set_const_cmt(cid, rptcmt.text, True);
            else:
                idaapi.set_enum_member_cmt(cid, rptcmt.text, True);
            self.update_counter(ENUM_ENTRY + ':' + REPEATABLE_CMT)
    

    def import_equate(self, equate, eid):
        """
        Processes EQUATE element as member of an enumeration.
        
        Args:
            enum_entry: XML element containing the equate name and value.
            eid: Integer representing the id for the enumeration.
        """
        name = self.get_attribute(equate,NAME)
        value = self.get_attribute_value(equate,VALUE)
        bm = -1
        if self.has_attribute(equate, BIT_MASK):
            bm = self.get_attribute_value(equate, BIT_MASK)
        if IDA_SDK_VERSION > 630:
            #idaapi.add_const(eid, name, value, bm)
            idaapi.add_const(eid, name, value)
            cid = idaapi.get_const_by_name(name)
        else:
            #idaapi.add_enum_member(eid, name, value, bm)
            idaapi.add_enum_member(eid, name, value)
            cid = idaapi.get_enum_member_by_name(name)
        self.update_counter(EQUATE)
        regcmt = equate.find(REGULAR_CMT) 
        if regcmt != None:
            if IDA_SDK_VERSION > 630:
                idaapi.set_const_cmt(cid, regcmt.text, False);
            else:
                idaapi.set_enum_member_cmt(cid, regcmt.text, False);
            self.update_counter(EQUATE + ':' + REGULAR_CMT)
        rptcmt = equate.find(REPEATABLE_CMT)
        if rptcmt != None:
            if IDA_SDK_VERSION > 630:
                idaapi.set_const_cmt(cid, rptcmt.text, True);
            else:
                idaapi.set_enum_member_cmt(cid, rptcmt.text, True);
            self.update_counter(EQUATE + ':' + REPEATABLE_CMT)
    

    def import_equate_group(self, equate_group):
        """
        Processes EQUATE_GROUP as IDA enumeration type.
        
        Args:
            equate_group: XML element containing the group name and
                equate definitions.
        """
        if self.Options.DataTypes.checked == False:
            return
        msg = EQUATE_GROUP
        name = ''
        if self.has_attribute(equate_group, NAME):
            name = self.get_attribute(equate_group, NAME)
        bf = ''
        if self.has_attribute(equate_group, BIT_FIELD):
            bf = self.get_attribute(equate_group, BIT_FIELD)
        eid = idaapi.add_enum(BADADDR, name, idaapi.hexflag())
        idaapi.set_enum_bf(eid, (bf == 'yes'))
        self.update_counter(EQUATE_GROUP)
        regcmt = equate_group.find(REGULAR_CMT)
        if regcmt != None:
            idaapi.set_enum_cmt(eid, regcmt.text, False)
            self.update_counter(EQUATE_GROUP + ':' + REGULAR_CMT)
        rptcmt = equate_group.find(REPEATABLE_CMT)
        if rptcmt != None:
            idaapi.set_enum_cmt(eid, rptcmt.text, True)
            self.update_counter(EQUATE_GROUP + ':' + REPEATABLE_CMT)
        equates = equate_group.findall(EQUATE)
        for equate in equates:
            self.import_equate(equate,eid)
        bit_masks = equate_group.findall(BIT_MASK)
        for bit_mask in bit_masks:
            self.import_bit_mask(bit_mask, eid)
    

    def import_equate_reference(self, equate_reference):
        if (self.Options.DataTypes.checked == False or
            self.Options.EquateReferences.checked == False):
            return
        self.update_counter(EQUATE_REFERENCE)
        addr = self.get_address(equate_reference, ADDRESS)
        name = ''
        if self.has_attribute(equate_reference, NAME):
            name = self.get_attribute(equate_reference, NAME)
        if name == '':
            return
        opnd = 0
        if self.has_attribute(equate_reference, OPERAND_INDEX):
            opnd = self.get_attribute_value(equate_reference, OPERAND_INDEX)
        value = None
        if self.has_attribute(equate_reference, VALUE):
            value = self.get_attribute_value(equate_reference, VALUE)
        if IDA_SDK_VERSION < 650:
            cid = idaapi.get_enum_member_by_name(name)
        else:
            cid = idaapi.get_const_by_name(name)
        if cid == idaapi.BADNODE:
            return
        if IDA_SDK_VERSION < 650:
            eid = idaapi.get_enum_member_enum(cid)
        else:
            eid = idaapi.get_const_enum(cid)
        if eid == idaapi.BADNODE:
            return
        idaapi.op_enum(addr, opnd, eid, 0)
        

    def import_function(self, function):
        """
        Creates a function using the FUNCTION attributes.
        
        Args:
            function: XML element containing the function address and
                attributes.
        """
        if self.Options.Functions.checked == False:
            return
        try:
            entry_point = self.get_address(function, ENTRY_POINT)
            name = ''
            if self.has_attribute(function, NAME):
                name = self.get_attribute(function, NAME)
            libfunc = 'n'
            if self.has_attribute(function, LIBRARY_FUNCTION):
                libfunc = self.get_attribute(function, LIBRARY_FUNCTION)
            if idaapi.isEnabled(entry_point) == False:
                msg = ("import_function: address %X not enabled in database"
                       % entry_point)
                print msg
                return
            idaapi.add_func(entry_point, BADADDR)
            self.update_counter(FUNCTION)
            func = idaapi.get_func(entry_point)
            if libfunc == 'y':
                func.flags |= idaapi.FUNC_LIB
            ranges = function.findall(ADDRESS_RANGE)
            for addr_range in ranges:
                (start, end) = self.import_address_range(addr_range)
                idaapi.append_func_tail(func, start, end)
                #idaapi.analyze_area(start, end+1)
            idc.Wait()
            regcmt = function.find(REGULAR_CMT)
            if regcmt != None:
                self.update_counter(FUNCTION + ':' + REGULAR_CMT)
                idaapi.set_func_cmt(func, regcmt.text, False)
            rptcmt = function.find(REPEATABLE_CMT)
            if rptcmt != None:
                self.update_counter(FUNCTION + ':' + REPEATABLE_CMT)
                idaapi.set_func_cmt(func, rptcmt.text, True)
            typecmt = function.find(TYPEINFO_CMT)
            if typecmt != None:
                self.update_counter(FUNCTION + ':' + TYPEINFO_CMT)
                idc.SetType(entry_point, typecmt.text + ';')
            sf = function.find(STACK_FRAME)
            if sf != None:
                self.import_stack_frame(sf, func)
            register_vars = function.findall(REGISTER_VAR)
            for register_var in register_vars:
                self.import_register_var(register_var, func)
        except:
            msg = "** Exception occurred in import_function **"
            print "\n" + msg + "\n", sys.exc_type, sys.exc_value


    def import_function_def(self, function_def):
        # import_function_def: NOT IMPLEMENTED
        if self.Options.DataTypes.checked == False:
            return
        self.update_counter(FUNCTION_DEF)
        

    def import_info_source(self, info_source):
        """
        Processes INFO_SOURCE containing information about the
            source of the XML PROGRAM file.
        
        Args:
            info_source: XML element containing attributes that identify
                the source of the PROGRAM data.
        """
        if self.has_attribute(info_source, TOOL):
            tool = self.get_attribute(info_source, TOOL)
        if self.has_attribute(info_source, USER):
            user = self.get_attribute(info_source, USER)
        if self.has_attribute(info_source, FILE):
            f = self.get_attribute(info_source, FILE)
        if self.has_attribute(info_source, TIMESTAMP):
            ts = self.get_attribute(info_source, TIMESTAMP)
        self.update_counter(INFO_SOURCE)
    

    def import_manual_instruction(self, manual_instruction):
        """
        Creates a manual instruction.
        
        Args:
            manual_instruction: XML element containing MANUAL_INSTRUCTION.
        """
        if self.Options.Manual.checked == False:
            return
        addr = self.get_address(manual_instruction, ADDRESS)
        idaapi.set_manual_insn(addr, manual_instruction.text)
        self.update_counter(MANUAL_INSTRUCTION)
    

    def import_manual_operand(self, manual_operand):
        """
        Creates a manual operand at an address.
        
        Args:
            manual_operand: MANUAL_OPERAND XML element.
        """
        if self.Options.Manual.checked == False:
            return
        addr = self.get_address(manual_operand, ADDRESS)
        op = self.get_attribute_value(manual_operand, OPERAND_INDEX)
        if idaapi.isEnabled(addr):
            idaapi.set_forced_operand(addr, op, manual_operand.text)
            self.update_counter(MANUAL_OPERAND)


    def process_deferred(self, element):
        """
        Processes the list of deferred structure members when the
        DATATYPES end element is encountered.
        
        Args:
            element: XML end element for DATATYPES 
        """
        for (member, sptr) in self.deferred:
            self.import_member(member, sptr, False)
        self.display_timer(element)


    def import_member(self, member, sptr, defer=True):
        """
        Creates a member for a structure.
        
        Args:
            member: MEMBER XML element.
            sptr:
            defer:  boolean indicating if processing a member should be
                    deferred when the type is unknown. A member should
                    only be deferred on the first pass, not when processing
                    the deferred list.
        """
        offset = self.get_attribute_value(member, OFFSET)
        datatype = self.get_attribute(member, DATATYPE)
        if self.has_attribute(member, DATATYPE_NAMESPACE):
            dt_namespace = self.get_attribute(member, DATATYPE_NAMESPACE)
        name = ''
        if self.has_attribute(member, NAME):
            name = self.get_attribute(member, NAME)
        size = 0
        if self.has_attribute(member, SIZE):
            size = self.get_attribute_value(member, SIZE)
        ti = idaapi.opinfo_t()
        if self.is_pointer_type(datatype):
            #idaapi.set_refinfo(ti, 0, BADADDR, 0, REF_OFF32)
            flag = idaapi.dwrdflag() | idaapi.FF_0OFF
        else:
            flag = self.get_datatype_flags(datatype, size)
        if flag == 0:
            if defer:
                self.deferred.append([member, sptr])
                return                
        if flag == idaapi.enumflag():
            t = idaapi.get_enum(datatype)
            ti.ec.tid = t
            ti.ec.serial = idaapi.get_enum_idx(t)
        if flag == idaapi.struflag():
            t = idaapi.get_struc_id(datatype)
            ti.tid = t
        error = idaapi.add_struc_member(sptr, name, offset, flag, ti, size)
        mbr = idaapi.get_member(sptr, offset)
        self.import_member_cmts(member, mbr)
        self.update_counter(MEMBER)
        
        
    def import_member_cmts(self, member, mbr):
        """
        Processes REGULAR_CMT and REPEATABLE_CMT elements for members.
        
        Args:
            element: XML element object containing a REGULAR_CMT or
                REPEATABLE_CMT element
            mbr: Integer representing the member id
        """
        regcmt = member.find(REGULAR_CMT)
        if regcmt != None:
            idaapi.set_member_cmt(mbr, regcmt.text, False)
            self.update_counter(MEMBER + ':' + REGULAR_CMT)
        rptcmt = member.find(REPEATABLE_CMT)
        if rptcmt != None:
            idaapi.set_member_cmt(mbr, rptcmt.text, True)
            self.update_counter(MEMBER + ':' + REPEATABLE_CMT)
        

    def import_members(self, element, sptr):
        """
        Add data members to a structure.
        
        Args:
            element: STRUCTURE XML element containing MEMBER sub-elements.
            sptr:
        """
        members = element.findall(MEMBER)
        for member in members:
            self.import_member(member, sptr)
                

    def import_memory_contents(self, memory_contents, start, size):
        """
        Processes MEMORY_CONTENTS to load data for a memory block.
        
        Args:
            memory_contents: MEMORY_CONTENTS XML element.
        """
        if memory_contents.get(START_ADDR) == None:
            saddr = start
        else:
            saddr = self.get_address(memory_contents, START_ADDR)
        fname = self.get_attribute(memory_contents, FILE_NAME)
        offset = self.get_attribute_value(memory_contents, FILE_OFFSET)
        if memory_contents.get(LENGTH) == None:
            length = size
        else:
            length = self.get_attribute_value(memory_contents, LENGTH)
        #(binfilename, ext) = os.path.splitext(self.filename)
        #binfilename += ".bytes"
        (binfilename, fileext) = os.path.split(self.filename)
        binfilename += '/' + fname
        binfile = idaapi.loader_input_t()
        binfile.open(binfilename)
        binfile.file2base(offset,saddr,saddr+length,False)
        binfile.close()
        self.update_counter(MEMORY_CONTENTS)
    

    def import_memory_map(self, memory_map):
        """
        Processes the MEMORY_MAP element.
        
        Args:
            memory_map: MEMORY_MAP XML element.
            
        MEMORY_MAP is only processed by the IDA loader. It is ignored when
            run as an IDA plugin.
        """
        # import memory sections only when run as loader
        if self.plugin:
            return
        self.update_import(memory_map)
    

    def import_memory_reference(self, memory_reference):
        """
        Processes the MEMORY_REFERENCE element.
        
        Args:
            memory_reference: MEMORY_REFERENCE XML element.
        """
        if self.Options.MemoryReferences.checked == False:
            return        
        addr = self.get_address(memory_reference, ADDRESS)
        if self.has_attribute(memory_reference, OPERAND_INDEX):
            op = self.get_attribute_value(memory_reference, OPERAND_INDEX)
        if self.has_attribute(memory_reference, USER_DEFINED):
            user = self.get_attribute(memory_reference, USER_DEFINED)
        to_addr = self.get_address(memory_reference, TO_ADDRESS)
        if self.has_attribute(memory_reference, BASE_ADDRESS):
            base_addr = self.get_address(memory_reference, BASE_ADDRESS)
        if self.has_attribute(memory_reference, PRIMARY):
            primary = self.get_attribute(memory_reference, PRIMARY)
        self.update_counter(MEMORY_REFERENCE)
        # TODO: import_memory_reference - add code to store reference
 
 
    def import_memory_section(self, memory_section):
        """
        Creates a memory segment in the database.
        
        Args:
            memory_section: MEMORY_SECTION XML element.
            
        MEMORY_SECTION is only processed by the IDA loader. It is ignored
            when run as an IDA plugin.
        """
        # TODO: import_memory_section - handle overlays?
        # import memory sections only when run as loader
        if self.plugin:
            return
        name = self.get_attribute(memory_section, NAME)
        length = self.get_attribute_value(memory_section, LENGTH)

        s = idaapi.segment_t()
        addrstr = self.get_attribute(memory_section, START_ADDR)
        seg_str = ''
        if '::' in addrstr:
            # overlay - skip for now
            print '  ** Overlayed memory block %s skipped **  ' % name
            msg  = 'Overlayed memory block %s skipped!' % name
            msg += "\n\nXML Import does not currently support"
            msg += "\noverlayed memory blocks."
            idaapi.warning(msg)
            return
        elif ':' in addrstr:
            [seg_str, offset_str] = string.split(addrstr,':')
            offset = int(offset_str, 16)
            if self.is_int(seg_str):
                base = int(seg_str, 16)
                sel = idaapi.setup_selector(base)
                start = self.get_address(memory_section, START_ADDR)
            else:
                raise MultipleAddressSpacesNotSupported
                return
        else:
            sel = idaapi.allocate_selector(0)
            start = int(addrstr, 16)

        s.sel = sel
        s.startEA = start
        s.endEA = start+length
        s.bitness = self.addr_mode

        perms = ''
        if self.has_attribute(memory_section, PERMISSIONS):
            perms = self.get_attribute(memory_section, PERMISSIONS)
        s.perm = 0
        if 'r' in perms: s.perm |= idaapi.SEGPERM_READ
        if 'w' in perms: s.perm |= idaapi.SEGPERM_WRITE
        if 'x' in perms: s.perm |= idaapi.SEGPERM_EXEC
        idaapi.add_segm_ex(s, name, "", idaapi.ADDSEG_OR_DIE |
                                            idaapi.ADDSEG_QUIET)
        self.update_counter(MEMORY_SECTION)
        for memory_contents in memory_section.findall(MEMORY_CONTENTS):
            self.import_memory_contents(memory_contents, start, length)
    

    def import_processor(self, processor):
        """
        Processes the PROCESSOR element.
        
        Args:
            processor: PROCESSOR XML element.
        """
        name = self.get_attribute(processor, NAME)
        self.update_counter(PROCESSOR)
        if self.plugin:
            return
        address_model = self.get_attribute(processor, ADDRESS_MODEL)
        if address_model != None:
            flag = idaapi.ph_get_flag()
            if str.lower(address_model) == '16-bit':
                self.addr_mode = 0
            elif str.lower(address_model) == '32-bit':
                self.addr_mode = 1
            elif str.lower(address_model) == '64-bit':
                self.addr_mode = 2
        

    def import_program(self, program):
        """
        Processes the PROGRAM element.
        
        Args:
            program: PROGRAM XML element.
        """
        self.update_status(PROGRAM)
        name = self.get_attribute(program, NAME)
        if self.has_attribute(program, EXE_PATH):
            epath = self.get_attribute(program, EXE_PATH)
            idaapi.set_root_filename(epath)
        else:
            idaapi.set_root_filename(name)
        if self.has_attribute(program, EXE_FORMAT):
            eformat = self.get_attribute(program, EXE_FORMAT)
            RootNode = idaapi.netnode('Root Node')
            if IDA_SDK_VERSION < 650:
                RootNode.supset(1, eformat)
            else:
                RootNode.supset(idaapi.RIDX_FILE_FORMAT_NAME, eformat)
        if self.has_attribute(program, IMAGE_BASE):
            base = self.get_attribute_value(program, IMAGE_BASE)
            idaapi.set_imagebase(base)
        if self.has_attribute(program, INPUT_MD5):
            input_md5 = self.get_attribute(program, INPUT_MD5)
            # store original md5 in a special netnode
            md5 = idaapi.netnode(INPUT_MD5, len(INPUT_MD5), True)
            if (IDA_SDK_VERSION < 650):
                md5.supset(1302, input_md5)
            else:
                md5.supset(idaapi.RIDX_MD5, input_md5)
        self.update_counter(PROGRAM)
        """
        # TODO: this needs to be on "end" event for PROGRAM
        description = program.find(DESCRIPTION)
        if description != None:
            pass  # figure out what to do with it
        """

        
    def import_program_entry_point(self, program_entry_point):
        """
        Defines a program entry point.
        
        Args:
            program_entry_point: PROGRAM_ENTRY_POINT XML element.
                Contains the entry point address.
        """
        if self.Options.EntryPoints.checked == False:
            return
        addr = self.get_address(program_entry_point, ADDRESS)
        idaapi.add_entry(addr, addr, "", True)
        self.update_counter(PROGRAM_ENTRY_POINT)
    

    def import_register_value_range(self, register_value_range):
        """
        Defines the address range for a register value.
        
        Args:
            register_value_range: REGISTER_VALUE_RANGE XML element.
                Contains the register, value, start address and range length.
        """
        if self.Options.RegisterValues.checked == False:
            return
        self.update_counter(REGISTER_VALUE_RANGE)
        reg = self.get_attribute(register_value_range, REGISTER)
        if reg == 'cs': return
        value = self.get_attribute_value(register_value_range, VALUE)
        addr = self.get_address(register_value_range, START_ADDRESS)
        length = self.get_attribute_value(register_value_range, LENGTH)
        r = idaapi.str2reg(reg)
        if r >= idaapi.ph.regFirstSreg and r <= idaapi.ph.regLastSreg:
            if IDA_SDK_VERSION < 670:
                idaapi.splitSRarea1(addr, r, value, 2)
            else:
                idaapi.splitSRarea1(addr, r, value, idaapi.SR_user)
    

    def import_register_var(self, register_var, func):
        """
        Defines a register variable for a function.
        
        Args:
            register_var: REGISTER_VAR XML element.
                Contains register, variable name, and datatype.
            func: IDA function object
        """
        name = self.get_attribute(register_var, NAME)
        reg = self.get_attribute(register_var, REGISTER)
        if self.has_attribute(register_var, DATATYPE):
            datatype = self.get_attribute(register_var, DATATYPE)
        if self.has_attribute(register_var, DATATYPE_NAMESPACE):
            namespace = self.get_attribute(register_var, DATATYPE_NAMESPACE)
        idc.MakeLocal(func.startEA, func.endEA, reg, name)
        self.update_counter(REGISTER_VAR)
    

    def import_stack_frame(self, stack_frame, func):
        """
        Defines a stack frame for a function.
        
        Args:
            stack_frame: STACK_FRAME element with STACK_VAR child elements.
        """
        if self.has_attribute(stack_frame, LOCAL_VAR_SIZE):
            lvsize = self.get_attribute_value(stack_frame, LOCAL_VAR_SIZE)
        if self.has_attribute(stack_frame, PARAM_OFFSET):
            param_offset = self.get_attribute_value(stack_frame, PARAM_OFFSET)
        if self.has_attribute(stack_frame, REGISTER_SAVE_SIZE):
            reg_save_size = self.get_attribute_value(stack_frame,
                                                     REGISTER_SAVE_SIZE)
        if self.has_attribute(stack_frame, RETURN_ADDR_SIZE):
            retaddr_size = self.get_attribute_value(stack_frame,
                                                    RETURN_ADDR_SIZE)
        if self.has_attribute(stack_frame, BYTES_PURGED):
            bytes_purged = self.get_attribute_value(stack_frame, BYTES_PURGED)
        self.update_counter(STACK_FRAME)
        for stack_var in stack_frame.findall(STACK_VAR):
            self.import_stack_var(stack_var, func)
    

    def import_stack_reference(self, stack_reference):
        # import_stack_reference: NOT IMPLEMENTED
        #self.update_counter(STACK_REFERENCE)
        pass
                    

    def import_stack_var(self, stack_var, func):
        """
        Processes STACK_VAR element.
        
        Args:
            stack_var: STACK_VAR XML element.
            
        Stack variables are created by IDA's function analysis. 
        Only the STACK_VAR NAME attribute is used to set the name for
        a stack variable at the specified stack/frame offset. 
        """
        spoffset = self.get_attribute_value(stack_var, STACK_PTR_OFFSET)
        datatype = self.get_attribute(stack_var, DATATYPE)
        offset = spoffset + func.frsize + func.frregs
        if self.has_attribute(stack_var, FRAME_PTR_OFFSET):
            fpoffset = self.get_attribute_value(stack_var, FRAME_PTR_OFFSET)
            offset = fpoffset + func.frsize
        name = ''
        if self.has_attribute(stack_var, NAME):
            name = self.get_attribute(stack_var, NAME)
        if self.has_attribute(stack_var, DATATYPE_NAMESPACE):
            namespace = self.get_attribute(stack_var, DATATYPE_NAMESPACE)
        if self.has_attribute(stack_var, SIZE):
            size = self.get_attribute_value(stack_var, SIZE)
        self.update_counter(STACK_VAR)
        sf = idaapi.get_frame(func)
        if name != '':
            idaapi.set_member_name(sf, offset, name)
    

    def import_structure(self, structure):
        """
        Adds a structure.
        
        Args:
            structure: STRUCTURE XML element.
                Contains the STRUCTURE attributes and child elements.
        """
        if self.Options.DataTypes.checked == False:
            return
        name = self.get_attribute(structure, NAME)
        dtyp = idaapi.get_struc_id(name)
        if dtyp != idaapi.BADNODE:
            # duplicate name, try adding name space
            if self.has_attribute(structure, NAMESPACE) == False:
                return
            namespace = self.get_attribute(structure, NAMESPACE)
            name = namespace + '__' + name
            name.replace('/','_')
            name.replace('.','_')
            dtyp = idaapi.get_struc_id(name)
            # skip if still duplicate (could add sequence #)
            if dtyp != idaapi.BADNODE:
                return
        size = 0
        if self.has_attribute(structure, SIZE):
            size = self.get_attribute_value(structure, SIZE)
        if self.has_attribute(structure, VARIABLE_LENGTH):
            vl = self.get_attribute_value(structure, VARIABLE_LENGTH)
            isVariableLength = vl == 'y'
        sid = idaapi.add_struc(BADADDR, name, False)
        sptr = idaapi.get_struc(sid)
        self.update_counter(STRUCTURE)
        self.import_cmts(structure, sid, STRUCTURE)
        self.import_members(structure, sptr)
        t = idaapi.opinfo_t()
        if idaapi.get_struc_size(sptr) < size:
            idaapi.add_struc_member(sptr,"",size-1,idaapi.byteflag(),t,1)
        

    def import_symbol(self, symbol):
        """
        Adds a symbol name at the specified address.
        
        Args:
            symbol: SYMBOL XML element.
                Contains symbol name and address. Optionally includes
                type and mangled symbol.
        """
        if self.Options.Symbols.checked == False:
            return
        addr = self.get_address(symbol, ADDRESS)
        name = self.get_attribute(symbol, NAME)
        if self.has_attribute(symbol, MANGLED):
            name = self.get_attribute(symbol, MANGLED)
        flag = idaapi.SN_NOWARN
        if self.has_attribute(symbol, TYPE):
            typ = self.get_attribute(symbol, TYPE)
            if  typ == 'local': flag |= idaapi.SN_LOCAL
        idaapi.set_name(addr, name, flag)
        self.update_counter(SYMBOL)
    

    def import_typedef(self, type_def):
        # import_typedef: NOT IMPLEMENTED
        if self.Options.DataTypes.checked == False:
            return
        self.update_counter(TYPE_DEF)
        

    def import_union(self, union):
        """
        Adds a union datatype.
        
        Args:
            union: UNION XML element.
                Contains UNION attributes and child elements.
        """
        if self.Options.DataTypes.checked == False:
            return
        name = self.get_attribute(union, NAME)
        dtyp = idaapi.get_struc_id(name)
        if dtyp != idaapi.BADNODE:
            # duplicate name, try adding name space
            if self.has_attribute(union, NAMESPACE) == False:
                return
            namespace = self.get_attribute(union, NAMESPACE)
            name = namespace + '__' + name
            name.replace('/','_')
            name.replace('.','_')
            dtyp = idaapi.get_struc_id(name)
            # skip if still duplicate (could add sequence #)
            if dtyp != idaapi.BADNODE:
                return
        size = 0
        if self.has_attribute(union, SIZE):
            size = self.get_attribute_value(union, SIZE)
        sid = idaapi.add_struc(BADADDR, name, True)
        sptr = idaapi.get_struc(sid)
        self.update_counter(UNION)
        self.import_cmts(union, sid, UNION)
        self.import_members(union, sptr)
        t = idaapi.opinfo_t()
        if idaapi.get_struc_size(sptr) < size:
            idaapi.add_struc_member(sptr,"", size-1, idaapi.byteflag(), None, 1)
            

    def open_file(self, filename, mode):
        """
        Opens filename to specified mode.
        
        Args:
            filename: String representing absolute filepath.
            mode: String representing mode for open.
            
        Returns
            File handle.
        
        Exceptions:
            Displays a warning and raises FileError exception
            if open fails.
        """
        try:
            f = open(filename, mode)
            return f
        except:
            fmt = "TITLE ERROR!\n"
            fmt += "ICON ERROR\n"
            fmt += "AUTOHIDE NONE\n"
            fmt += "HIDECANCEL\n"
            fmt += "Error opening file" + filename + "!\n"
            idaapi.warning(fmt)
            raise FileError
    

    def update_counter(self, tag):
        """
        Updates the counter for the element tag.
        
        Args:
            tag: String representing element tag.
        """
        if tag in self.Elements:
            self.Counters[self.Elements[tag]] += 1
        else:
            self.Elements[tag] = len(self.Elements)
            self.Counters.append(1)
            self.Tags.append(tag)
    

    def update_import(self, element):
        """
        Update the element counter and processing status.
        
        Args:
            element: XML element
            
        This function is used to process certain high-level elements
        (such as COMMENTS, CODE_BLOCKS, SYMBOL_TABLE, FUNCTIONS, etc.)
        that are used to group sub-elements.
        """
        self.update_counter(element.tag)
        self.update_status(element.tag)
    

    def update_status(self, tag):
        """
        Displays the processing status in the IDA window.
        
        Args:
            tag: String representing XML element tag
        """
        status = 'Importing ' + tag
        idaapi.msg('\n%-35s' % status)
        idaapi.hide_wait_box()
        idaapi.show_wait_box(status)
    
        
# Global XML string constants for elements and attributes
ADDRESS = 'ADDRESS'
ADDRESS_MODEL = 'ADDRESS_MODEL'
ADDRESS_RANGE = 'ADDRESS_RANGE'
BASE_ADDRESS = 'BASE_ADDRESS'
BIT_FIELD = 'BIT_FIELD'
BIT_MAPPED = 'BIT_MAPPED'
BIT_MASK = 'BIT_MASK'
BOOKMARK = 'BOOKMARK'
BOOKMARKS = 'BOOKMARKS'
BYTES = 'BYTES'
BYTES_PURGED = 'BYTES_PURGED'
CATEGORY = 'CATEGORY'
CODE = 'CODE'
CODE_BLOCK = 'CODE_BLOCK'
COMMENT = 'COMMENT'
COMMENTS = 'COMMENTS'
COMPILER = 'COMPILER'
DATA = 'DATA'
DATATYPE = 'DATATYPE'
DATATYPES = 'DATATYPES'
DATATYPE_NAMESPACE = 'DATATYPE_NAMESPACE'
DEFINED_DATA = 'DEFINED_DATA'
DESCRIPTION = 'DESCRIPTION'
DISPLAY_SETTINGS = 'DISPLAY_SETTINGS'
END = 'END'
ENDIAN = 'ENDIAN'
ENTRY_POINT = 'ENTRY_POINT'
ENUM = 'ENUM'
ENUM_ENTRY = 'ENUM_ENTRY'
EQUATE = 'EQUATE'
EQUATES = 'EQUATES'
EQUATE_GROUP = 'EQUATE_GROUP'
EQUATE_REFERENCE = 'EQUATE_REFERENCE'
EXE_FORMAT = 'EXE_FORMAT'
EXE_PATH = 'EXE_PATH'
EXT_LIBRARY = 'EXT_LIBRARY'
EXT_LIBRARY_REFERENCE = 'EXT_LIBRARY_REFERENCE'
EXT_LIBRARY_TABLE = 'EXT_LIBRARY_TABLE'
FAMILY = 'FAMILY'
FILE = 'FILE'
FILE_NAME = 'FILE_NAME'
FILE_OFFSET = 'FILE_OFFSET'
FOLDER = 'FOLDER'
FORMAT = 'FORMAT'
FRAGMENT = 'FRAGMENT'
FRAME_PTR_OFFSET = 'FRAME_PTR_OFFSET'
FUNCTION = 'FUNCTION'
FUNCTIONS = 'FUNCTIONS'
FUNCTION_DEF = 'FUNCTION_DEF'
IMAGE_BASE = 'IMAGE_BASE'
INPUT_MD5 = 'INPUT_MD5'
INFO_SOURCE = 'INFO_SOURCE'
LANGUAGE_PROVIDER = 'LANGUAGE_PROVIDER'
LENGTH = 'LENGTH'
LIB_ADDR = 'LIB_ADDR'
LIB_LABEL = 'LIB_LABEL'
LIB_ORDINAL = 'LIB_ORDINAL'
LIB_PROG_NAME = 'LIB_PROG_NAME'
LIBRARY_FUNCTION = 'LIBRARY_FUNCTION'
LOCAL_VAR_SIZE = 'LOCAL_VAR_SIZE'
MANGLED = 'MANGLED'
MANUAL_INSTRUCTION = 'MANUAL_INSTRUCTION'
MANUAL_OPERAND = 'MANUAL_OPERAND'
MARKUP = 'MARKUP'
MEMBER = 'MEMBER'
MEMORY_CONTENTS = 'MEMORY_CONTENTS'
MEMORY_MAP = 'MEMORY_MAP'
MEMORY_REFERENCE = 'MEMORY_REFERENCE'
MEMORY_SECTION = 'MEMORY_SECTION'
NAME = 'NAME'
NAMESPACE = 'NAMESPACE'
OFFSET = 'OFFSET'
OPERAND_INDEX = 'OPERAND_INDEX'
PARAM_OFFSET = 'PARAM_OFFSET'
PATH = 'PATH'
PERMISSIONS = 'PERMISSIONS'
PRIMARY = 'PRIMARY'
PROCESSOR = 'PROCESSOR'
PROGRAM = 'PROGRAM'
PROGRAM_ENTRY_POINT = 'PROGRAM_ENTRY_POINT'
PROGRAM_ENTRY_POINTS = 'PROGRAM_ENTRY_POINTS'
PROGRAM_TREES = 'PROGRAM_TREES'
PROPERTIES = 'PROPERTIES'
PROPERTY = 'PROPERTY'
REGISTER = 'REGISTER'
REGISTER_SAVE_SIZE = 'REGISTER_SAVE_SIZE'
REGISTER_VALUES = 'REGISTER_VALUES'
REGISTER_VALUE_RANGE = 'REGISTER_VALUE_RANGE'
REGISTER_VAR = 'REGISTER_VAR'
REGULAR_CMT = 'REGULAR_CMT'
RELOCATION = 'RELOCATION'
RELOCATION_TABLE = 'RELOCATION_TABLE'
REPEATABLE_CMT = 'REPEATABLE_CMT'
RETURN_ADDR_SIZE = 'RETURN_ADDR_SIZE'
RETURN_TYPE = 'RETURN_TYPE'
SHOW_TERMINATOR = 'SHOW_TERMINATOR'
SIGNED = 'SIGNED'
SIZE = 'SIZE'
SOURCE_ADDRESS = 'SOURCE_ADDRESS'
SOURCE_TYPE = 'SOURCE_TYPE'
STACK_FRAME = 'STACK_FRAME'
STACK_PTR_OFFSET = 'STACK_PTR_OFFSET'
STACK_REFERENCE = 'STACK_REFERENCE'
STACK_VAR = 'STACK_VAR'
START = 'START'
START_ADDR = 'START_ADDR'
START_ADDRESS = 'START_ADDRESS'
STRUCTURE = 'STRUCTURE'
SYMBOL = 'SYMBOL'
SYMBOL_TABLE = 'SYMBOL_TABLE'
TIMESTAMP = 'TIMESTAMP'
TOOL = 'TOOL'
TO_ADDRESS = 'TO_ADDRESS'
TREE = 'TREE'
TYPE = 'TYPE'
TYPEINFO_CMT = 'TYPEINFO_CMT'
TYPE_DEF = 'TYPE_DEF'
UNION = 'UNION'
USER = 'USER'
USER_DEFINED = 'USER_DEFINED'
VALUE = 'VALUE'
VARIABLE_LENGTH = 'VARIABLE_LENGTH'
ZERO_PAD = 'ZERO_PAD'


