# plugins

repository: patois/xray
file: xray.py
url: https://github.com/patois/xray/blob/10fda7eaa330e6c3450643c5b3829cdf759c8b20/xray.py
created_at: 2019-06-21 16:18:10+00:00
pushed_at: 2023-12-08 16:20:53+00:00
description:
Hexrays decompiler plugin that colorizes and filters the decompiler's output based on regular expressions

content:
import shutil, re, os, configparser, errno
from PyQt5.QtWidgets import QWidget, QSplitter

import ida_hexrays
import ida_idaapi
import ida_kernwin as kw
import ida_lines as il
import ida_diskio
from ida_pro import IDA_SDK_VERSION

__author__ = "Dennis Elser"

PLUGIN_NAME = "xray"

XRAY_FILTER_ACTION_ID = "%s:filter" % PLUGIN_NAME
XRAY_LOADCFG_ACTION_ID = "%s:loadcfg" % PLUGIN_NAME
XRAY_QUERY_ACTION_ID = "%s:query" % PLUGIN_NAME
PATTERN_LIST = []
HIGH_CONTRAST = False

IS_ACTIVATED = False
TEXT_INPUT_FORMS = {}

CFG_FILENAME = "%s.cfg" % PLUGIN_NAME
DEFAULT_CFG = """# configuration file for xray.py

[global]
# set to 1 for better contrast
high_contrast=0
# enable/disable xray when loading a database 
auto_enable=0

# each group contains a list of regular
# expressions, a background color in
# RRGGBB format and an optional hint field.
# priority is determined by order of
# appearance, first group gets assigned
# lowest priority.
# check out https://regex101.com/r and
# https://www.debuggex.com/

[group_01]
hint=loop
bgcolor=4c0037

expr_01=^while\(
expr_02=^for\(

[group_02]
hint=function name
bgcolor=00374c

expr_01=recv\(
expr_02=malloc\(
expr_03=realloc\(
expr_04=free\(
expr_05=memcpy\(
expr_06=memmove\(
expr_07=strcpy\(
expr_08=sscanf\(
expr_09=sprintf\(
expr_10=recvfrom\(

[group_03]
hint=format strings
bgcolor=4c1500

expr_01=sscanf\(.*,.*%s.*,.*\)
expr_02=sprintf\(.*,.*%s.*,.*\)

[group_04]
hint=arithmetic
bgcolor=4c1500

expr_01=malloc\(.*[\*\+\-\/%][^>].*?\)
expr_02=realloc\(([^,]+,){1}(.*[^,][\+\-\*\/%][^>].*[^,])
expr_03=memcpy\(([^,]+,){2}(.*[^,][\+\-\*\/%][^>].*[^,])
expr_04=memmove\(([^,]+,){2}(.*[^,][\+\-\*\/%][^>].*[^,])
expr_05=recv\(([^,]+,){2}(.*[^,][\+\-\*\/%][^>].*[^,])
expr_06=recvfrom\(([^,]+,){2}(.*[^,][\+\-\*\/%][^>].*[^,])"""

# -----------------------------------------------------------------------------
def is_plugin():
    """returns True if this script is executed from within an IDA plugins
    directory, False otherwise."""
    return "__plugins__" in __name__

# -----------------------------------------------------------------------------
def get_dest_filename():
    """returns destination path for plugin installation."""
    return os.path.join(
        ida_diskio.get_user_idadir(),
        "plugins",
        "%s%s" % (PLUGIN_NAME, ".py"))

# -----------------------------------------------------------------------------
def is_installed():
    """checks whether script is present in designated plugins directory."""
    return os.path.isfile(get_dest_filename())

# -----------------------------------------------------------------------------
def get_cfg_filename():
    """returns full path for config file."""
    return os.path.join(
        ida_diskio.get_user_idadir(),
        "plugins",
        "%s" % CFG_FILENAME)

# -----------------------------------------------------------------------
def is_min_sdk_ver(min_ver_required):
    return IDA_SDK_VERSION >= min_ver_required

# -----------------------------------------------------------------------------
def is_compatible():
    """Checks whether script is compatible with current IDA and
    decompiler versions."""
    min_ida_ver = 720
    return is_min_sdk_ver(min_ida_ver) and ida_hexrays.init_hexrays_plugin()

# -----------------------------------------------------------------------------
SELF = __file__
def install_plugin():
    """Installs script to IDA userdir as a plugin."""
    dst = get_dest_filename()
    src = SELF
    if is_installed():
        btnid = kw.ask_yn(kw.ASKBTN_NO,
            "File exists:\n\n%s\n\nReplace?" % dst)
        if btnid is not kw.ASKBTN_YES:
            return False
    else:
        btnid = kw.ask_yn(kw.ASKBTN_NO,
            "This plugin is about to be installed to:\n\n%s\n\nInstall now?" % dst)
        if btnid is not kw.ASKBTN_YES:
            return False

    usrdir = os.path.dirname(dst)
    kw.msg("%s: copying script from \"%s\" to \"%s\" ..." % (PLUGIN_NAME, src, usrdir))
    if not os.path.exists(usrdir):
        try:
            os.makedirs(usrdir)
        except OSError as e:
            if e.errno != errno.EEXIST:
                kw.msg("failed (mkdir)!\n")
                return False
    try:
        shutil.copy(src, dst)
    except:
        kw.msg("failed (copy)!\n")
        return False
    kw.msg(("done\n"
        "Plugin installed - please restart this instance of IDA.\n"))
    return True

# -----------------------------------------------------------------------------
def swapcol(x):
    """converts between RRGGBB and BBGGRR color encodings."""
    return (((x & 0x000000FF) << 16) |
             (x & 0x0000FF00) |
            ((x & 0x00FF0000) >> 16))

# -----------------------------------------------------------------------------
def load_cfg(reload=False):
    """loads xray configuration from file. Creates and loads default config
    if none is present."""
    global PATTERN_LIST
    global HIGH_CONTRAST
    global IS_ACTIVATED

    cfg_file = get_cfg_filename()
    kw.msg("%s: %sloading %s...\n" % (PLUGIN_NAME,
        "re" if reload else "",
        cfg_file))
    if not os.path.isfile(cfg_file):
        kw.msg("%s: %s does not exist! creating default config... " % (PLUGIN_NAME, cfg_file))
        try:
            with open(cfg_file, "w") as f:
                f.write(DEFAULT_CFG)
                kw.msg("success!\n")
        except:
            kw.msg("failed!\n")
            return False
        return load_cfg(reload=True)

    PATTERN_LIST = []

    config = configparser.RawConfigParser()
    config.readfp(open(cfg_file))

    # read all sections
    for section in config.sections():
        expr_list = []
        if section.startswith("group_"):
            for k,v in config.items(section):
                if k.startswith("expr_"):
                    expr_list.append(v)
            try:
                bgcolor = swapcol(int(config.get(section, "bgcolor"), 16))
            except:
                bgcolor = swapcol(0x000000)
            try:
                hint = config.get(section, "hint")
            except:
                hint = None
            PATTERN_LIST.append(ConfigGroupSection(expr_list, bgcolor, hint))
        elif section == "global":
            try:
                HIGH_CONTRAST = config.getboolean(section, "high_contrast")
            except:
                HIGH_CONTRAST = False
            if not reload:
                try:
                    IS_ACTIVATED = config.getboolean(section, "auto_enable")
                except:
                    IS_ACTIVATED = False

    if not len(PATTERN_LIST):
        kw.warning("Config file does not contain any regular expressions.")
    return True

# -----------------------------------------------------------------------------
class TextInputForm(kw.Form):
    """Input form for regex search queries."""

    # flags
    SO_FIND_TEXT = 1
    SO_FIND_REGEX = 2
    SO_FILTER_TEXT = 4
    SO_FILTER_COLOR = 8
    SO_FIND_CASE = 16

    def __init__(self, parent_widget):
        self.parent_widget = parent_widget
        self.parent_title = kw.get_widget_title(self.parent_widget)
        i=1
        while kw.find_widget("%s-%d" % (PLUGIN_NAME, i)):
            i+=1
        self.idx = i
        __title = "%s-%s" % (PLUGIN_NAME, self.idx)
        self.options = (TEXT_INPUT_FORMS[self.parent_title].options
            if self.parent_title in TEXT_INPUT_FORMS.keys()
            else TextInputForm.SO_FILTER_TEXT | TextInputForm.SO_FIND_TEXT)
        self.query = (TEXT_INPUT_FORMS[self.parent_title].query
            if self.parent_title in TEXT_INPUT_FORMS.keys()
            else "")
        kw.Form.__init__(self,
("BUTTON YES NONE\n"
"BUTTON NO NONE\n"
"BUTTON CANCEL NONE\n"
"%s\n\n"
"{FormChangeCb}\n"
"<##Enter text##Filter:{cbEditable}>"
"|<##Search##Text:{rText}><Regex:{rRegex}>{cSearchMethod}>"
"|<##Filter##Lines:{rFilterLines}><Colors:{rFilterColors}>{cFilterType}>"
"|<##Options##Case sensitive:{rCase}>{cOptions}>\n") % (__title), {
    'FormChangeCb': kw.Form.FormChangeCb(self.OnFormChange),
    'cbEditable': kw.Form.StringInput(value = self.query),
    'cSearchMethod': kw.Form.RadGroupControl(("rText", "rRegex")),
    'cFilterType': kw.Form.RadGroupControl(("rFilterLines", "rFilterColors")),
    'cOptions': kw.Form.ChkGroupControl(("rCase",))})

    def init_controls(self):
        self.SetControlValue(self.cbEditable, self.query)
        self.SetControlValue(self.cSearchMethod, 0 if self.options & TextInputForm.SO_FIND_TEXT else 1)
        self.SetControlValue(self.cFilterType, 0 if self.options & TextInputForm.SO_FILTER_TEXT else 1)
        self.SetControlValue(self.cOptions, 1 if self.options & TextInputForm.SO_FIND_CASE else 0)
        self.SetFocusedField(self.cbEditable)
        return

    def _commit_changes(self):
        vu = ida_hexrays.get_widget_vdui(self.parent_widget)
        if vu:
            vu.refresh_ctext()
            # "refresh_ctext()" took away the focus, take it back
            kw.activate_widget(kw.find_widget(self.title), True)
            self.SetFocusedField(self.cbEditable)
            return True
        return False

    def OnFormChange(self, fid):
        if fid == self.cbEditable.id:
            self.query = self.GetControlValue(self.cbEditable)
        elif fid == self.rCase.id:
            if self.GetControlValue(self.cOptions):
                self.options |= TextInputForm.SO_FIND_CASE
            else:
                self.options &= ~TextInputForm.SO_FIND_CASE & 0xFFFFFFFF
        elif fid in [self.rFilterLines.id, self.rFilterColors.id]:
            filter_text = fid == self.rFilterLines.id
            filter_color = fid == self.rFilterColors.id

            if filter_text:
                self.options |= TextInputForm.SO_FILTER_TEXT
            else:
                self.options &= ~TextInputForm.SO_FILTER_TEXT & 0xFFFFFFFF

            if filter_color:
                self.options |= TextInputForm.SO_FILTER_COLOR
            else:
                self.options &= ~TextInputForm.SO_FILTER_COLOR & 0xFFFFFFFF
        elif fid in [self.rText.id, self.rRegex.id]:
            find_text = fid == self.rText.id
            find_regex = fid == self.rRegex.id

            if find_text:
                self.options |= TextInputForm.SO_FIND_TEXT
            else:
                self.options &= ~TextInputForm.SO_FIND_TEXT & 0xFFFFFFFF

            if find_regex:
                self.options |= TextInputForm.SO_FIND_REGEX
            else:
                self.options &= ~TextInputForm.SO_FIND_REGEX & 0xFFFFFFFF
        
        self._commit_changes()
        return 1

# -----------------------------------------------------------------------------
class ConfigGroupSection():
    """class that represents a config file's "group" section."""
    def __init__(self, expr_list, bgcolor, hint):
        self.expr_list = expr_list
        self.bgcolor = bgcolor
        self.hint = hint

# -----------------------------------------------------------------------------
class xray_hooks_t(ida_hexrays.Hexrays_Hooks):
    """class for handling decompiler events."""

    def _remove_color_tags(self, l):
        """removes all color tags from a tagged simple_line_t object
        but preserves COLOR_ADDR tags."""
        line = ""
        i = 0
        while i<len(l):
            if l[i] is il.COLOR_ON:
                n = il.tag_skipcode(l[i:])
                if l[i:].find(chr(il.COLOR_ADDR)) == 1:
                    line += l[i:i+n]
                i += n
            elif l[i] in [il.COLOR_OFF, il.COLOR_ESC, il.COLOR_INV]:
                n = il.tag_skipcode(l[i:])
                i += n
            else:
                line += l[i]
                i += 1
        return line

    def _search(self, regexp, sl, case_sensitive = False):
        line = il.tag_remove(sl.line).lstrip().rstrip()
        return re.search(regexp, line, flags=re.I if not case_sensitive else 0) is not None

    def _apply_xray_filter(self, vu, pc):
        if IS_ACTIVATED and pc:
            #col = il.calc_bg_color(ida_idaapi.get_inf_structure().min_ea)
            #col = pc[0].bgcolor
            for sl in pc:
                match=False
                for group in PATTERN_LIST:
                    for expr in group.expr_list:
                        if self._search(expr, sl):
                            #sl.bgcolor = (col & 0xfefefe) >> 1
                            sl.bgcolor = group.bgcolor
                            match=True
                            break
                if not match and HIGH_CONTRAST:
                    sl.line = self._remove_color_tags(sl.line)
        return

    def _apply_query_filter(self, vu, pc):
        new_pc = []
        title = kw.get_widget_title(vu.ct)
        if title in TEXT_INPUT_FORMS.keys() and pc:
            sq = TEXT_INPUT_FORMS[title]
            query = sq.query
            if not len(query):
                kw.set_highlight(vu.ct, None, HL_FLAGS)
                return
            options = sq.options
            case_sensitive = options & TextInputForm.SO_FIND_CASE

            # TODO
            if options & TextInputForm.SO_FIND_TEXT:
                kw.set_highlight(vu.ct, query, HL_FLAGS)
                tmpquery = query.lower() if not case_sensitive else query
                for sl in pc:
                    haystack = il.tag_remove(sl.line).lstrip().rstrip()
                    haystack = haystack.lower() if not case_sensitive else haystack
                    if tmpquery in haystack:
                        new_pc.append(sl.line)
                    else:
                        if options & TextInputForm.SO_FILTER_COLOR:
                            # add line but remove color
                            new_pc.append(self._remove_color_tags(sl.line))
                        elif options & TextInputForm.SO_FILTER_TEXT:
                            # do not add non-matching text
                            pass
            elif options & TextInputForm.SO_FIND_REGEX:
                kw.set_highlight(vu.ct, None, 0)
                for sl in pc:
                    try:
                        if self._search(query, sl, case_sensitive):
                            new_pc.append(sl.line)
                        else:
                            if options & TextInputForm.SO_FILTER_COLOR:
                                new_pc.append(self._remove_color_tags(sl.line))
                            elif options & TextInputForm.SO_FILTER_TEXT:
                                # do not add non-matching text
                                pass
                    except re.error as error:
                        kw.msg("%s: %s: \"%s\"" %
                            (PLUGIN_NAME, error, query))
                        return
            pc.clear()
            sl = kw.simpleline_t()
            for line in new_pc:
                sl.line = line
                pc.push_back(sl)
        return

    def _build_hint(self, vu):
        if IS_ACTIVATED and vu.refresh_cpos(ida_hexrays.USE_MOUSE):
            sl = vu.cfunc.get_pseudocode()[vu.cpos.lnnum]
            hint_lines = ["%s pattern(s):" % PLUGIN_NAME]
            delim_s = "%s" % "="*len(hint_lines[0])
            delim_e = "\n%s\n" % ("-"*len(hint_lines[0]))
            hint_lines.append(delim_s)
            hint = ""
            hint_created = False
            for group in PATTERN_LIST:
                for expr in group.expr_list:
                    if self._search(expr, sl):
                        tmp = (" (%s)" % group.hint) if group.hint else ""
                        hint_lines.append("> \"%s\"%s" % (expr, tmp))
                        hint_created = True
            hint_lines.append(delim_e)
            hint = "\n".join(hint_lines)
            if hint_created:
                return (hint, len(hint_lines)+1)
        return None

    def text_ready(self, vu):
        pc = vu.cfunc.get_pseudocode()
        self._apply_query_filter(vu, pc)
        self._apply_xray_filter(vu, pc)
        return 0

    def populating_popup(self, widget, phandle, vu):
        kw.attach_action_to_popup(vu.ct, None, XRAY_FILTER_ACTION_ID, PLUGIN_NAME+"/")
        kw.attach_action_to_popup(vu.ct, None, XRAY_LOADCFG_ACTION_ID, PLUGIN_NAME+"/")
        kw.attach_action_to_popup(vu.ct, None, XRAY_QUERY_ACTION_ID, PLUGIN_NAME+"/")
        return 0

    def create_hint(self, vu):
        result = self._build_hint(vu)
        if result:
            hint, count = result
            return (2, hint, count)
        return (0, None)

# -----------------------------------------------------------------------------
class xray_action_handler_t(kw.action_handler_t):
    """action handler for turning xray on and off."""
    def __init__(self):
        kw.action_handler_t.__init__(self)

    def activate(self, ctx):
        global IS_ACTIVATED
        IS_ACTIVATED = not IS_ACTIVATED
        vu = ida_hexrays.get_widget_vdui(ctx.widget)
        if vu:
            vu.refresh_ctext()
        kw.msg("[%s] %sactivated.\n" % (PLUGIN_NAME, "" if IS_ACTIVATED else "de"))
        return 1

    def update(self, ctx):
        return kw.AST_ENABLE_FOR_WIDGET if \
            ctx.widget_type == kw.BWN_PSEUDOCODE else \
            kw.AST_DISABLE_FOR_WIDGET

# -----------------------------------------------------------------------------
class loadcfg_action_handler_t(kw.action_handler_t):
    """action handler for reloading xray cfg file."""
    def __init__(self):
        kw.action_handler_t.__init__(self)

    def activate(self, ctx):
        if load_cfg(reload=True):
            vu = ida_hexrays.get_widget_vdui(ctx.widget)
            if vu:
                vu.refresh_ctext()
        return 1

    def update(self, ctx):
        return kw.AST_ENABLE_FOR_WIDGET if \
            ctx.widget_type == kw.BWN_PSEUDOCODE else \
            kw.AST_DISABLE_FOR_WIDGET

# -----------------------------------------------------------------------------
class regexfilter_action_handler_t(kw.action_handler_t):
    """action handler for search queries."""
    def __init__(self):
        kw.action_handler_t.__init__(self)

    # TODO
    def _dirty_resize_hack(self, w, form):
        title = form.title
        widget = kw.find_widget(title)
        if not widget:
            return

        w1 = kw.PluginForm.TWidgetToPyQtWidget(widget)
        w2 = kw.PluginForm.TWidgetToPyQtWidget(w)
        if not w1 or not w2:
            return

        p1 = w1.parentWidget()
        p2 = w2.parentWidget()
        if not p1 or not p2:
            return

        splitter = p1.parentWidget()
        hr = p2.parentWidget()
        if not splitter or not hr:
            return

        if not type(splitter) is QSplitter or not type(hr) is QWidget:
            return

        sizes = splitter.sizes()
        if len(sizes) != 2:
            return

        idx = splitter.indexOf(p1)
        _min, _max = splitter.getRange(idx)
        sizes[idx] = _min

        idx = splitter.indexOf(hr)
        _min, _max = splitter.getRange(idx)
        sizes[idx] = _max

        splitter.setSizes(sizes)
        return

    def _open_search_form(self, widget):
        global TEXT_INPUT_FORMS

        title = kw.get_widget_title(widget)
        if title not in TEXT_INPUT_FORMS.keys():
            search_form = TextInputForm(widget)
            search_form.modal = False
            search_form.openform_flags = (kw.PluginForm.WOPN_DP_BOTTOM |
                kw.PluginForm.WOPN_PERSIST)
            search_form, _ = search_form.Compile()
            search_form.Open()
            TEXT_INPUT_FORMS[title] = search_form
            self._dirty_resize_hack(widget, search_form)
        else:
            search_form = TEXT_INPUT_FORMS[title]
            search_form.Open()
            search_form.init_controls()
            self._dirty_resize_hack(widget, search_form)
        return

    def activate(self, ctx):
        self._open_search_form(ctx.widget)
        return 1

    def update(self, ctx):
        return kw.AST_ENABLE_FOR_WIDGET if \
            ctx.widget_type == kw.BWN_PSEUDOCODE else \
            kw.AST_DISABLE_FOR_WIDGET

# -----------------------------------------------------------------------------
class xray_plugin_t(ida_idaapi.plugin_t):
    """plugin class."""
    flags = ida_idaapi.PLUGIN_HIDE
    comment = PLUGIN_NAME
    help = PLUGIN_NAME
    wanted_name = PLUGIN_NAME
    wanted_hotkey = ""

    def init(self):
        self.xray_hooks = None
        if not is_compatible():
            kw.msg("%s: decompiler not available, skipping." % PLUGIN_NAME)
            return ida_idaapi.PLUGIN_SKIP

        load_cfg()

        kw.register_action(
            kw.action_desc_t(
                XRAY_LOADCFG_ACTION_ID,
                "%s: reload config" % PLUGIN_NAME,
                loadcfg_action_handler_t(),
                "Ctrl-R"))

        kw.register_action(
            kw.action_desc_t(
                XRAY_FILTER_ACTION_ID,
                "%s: toggle" % PLUGIN_NAME,
                xray_action_handler_t(),
                "F3"))

        kw.register_action(
            kw.action_desc_t(
                XRAY_QUERY_ACTION_ID,
                "%s: search" % PLUGIN_NAME,
                regexfilter_action_handler_t(),
                "Ctrl-F"))

        self.xray_hooks = xray_hooks_t()
        self.xray_hooks.hook()
        kw.msg("[%s] F3: toggle xray, Ctrl-F: filter, Ctrl-R: reload config." % PLUGIN_NAME)
        return ida_idaapi.PLUGIN_KEEP

    def run(self, arg):
        return

    def term(self):
        if self.xray_hooks:
            self.xray_hooks.unhook()
            kw.unregister_action(XRAY_FILTER_ACTION_ID)
            kw.unregister_action(XRAY_LOADCFG_ACTION_ID)
            kw.unregister_action(XRAY_QUERY_ACTION_ID)
        return

# -----------------------------------------------------------------------------
def PLUGIN_ENTRY():
    """plugin entry point."""
    return xray_plugin_t()

# -----------------------------------------------------------------------------
def SCRIPT_ENTRY():
    """script entry point."""
    if not is_plugin():
        (kw.info("Success!") if install_plugin() else
            kw.warning("Error! Plugin could not be installed!"))
    return

# -----------------------------------------------------------------------------

HL_FLAGS = kw.HIF_LOCKED
if is_min_sdk_ver(740):
    HL_FLAGS |= kw.HIF_NOCASE
if is_min_sdk_ver(770):
    HL_FLAGS |= kw.HIF_USE_SLOT

SCRIPT_ENTRY()


repository: iah3/binary-analysis
file: static_data_dependence.py
description: Reverse engineering binary executables using IDA and Pin
wanted_name: Lab 4 plugin
comment: This is a comment
created_at: 2019-11-15 20:45:07+00:00
pushed_at: 2019-11-15 21:08:54+00:00
url:
https://github.com/iah3/binary-analysis/blob/473259631e0611c46ba6c039dac6140f6d2ee20c/static_data_dependence.py

content:
import idaapi
import string
from sets import Set
from idautils import *
from idc import *
from ida_funcs import *
import re, string
from idaapi import *

import os
from sets import Set

from enum import Enum

out_dir = 'outputs'
if not os.path.exists(out_dir):
    os.makedirs(out_dir)

class optype(Enum):
	NONE = 0
	REGISTER = 1
	MEMORY = 2
	BASE_INDEX = 3
	BASE_INDEX_DISP = 4
	IMM = 5
	IMM_FAR = 6
	IMM_NEAR = 7

reg_list = ['al', 'ah', 'ax', 'eax', 'bl', 'bh', 'bx', 'ebx', 'cl', 'ch', 'cx', 'ecx', 'dl', 'dh', 'dx', 'edx', 'si', 'esi', 'di', 'edi', 'sp', 'esp', 'bp', 'ebp']

def write_file(final_string, func_name):
	with open("outputs/" + hex(func_name) + ".dot", "w") as text_file:
		text_file.write(final_string)

def append_node(string_node, node, label, tail):
	return string_node + "	" + node + "	" + "[label = \"" + label + "\"];\n" #+ "; " + tail + "\"];\n"

def append_edge(string_edge, node1, node2):
	return string_edge + "	" + node1 + " -> " + node2 + ";\n"

def final_string_parse(string_node, string_edge):
	return "digraph controlflow {\n" + string_node + "\n\n" + string_edge + "\n}"

def parse_output(nodes, edges, func_name):
	string_edge = ""
	string_node = ""
	d = "def"
	u = "use"
	node_counter = 0
	node_ids = {}

	node_id = 'n' + str(node_counter)
	node_ids['START'] = node_id
	node_counter += 1
	string_node = append_node(string_node, node_id, 'START', '')

	for node in nodes:
		if node == 'START':
			pass
		elif hex(node) not in node_ids.keys():
			node_id = 'n' + str(node_counter)
			node_ids[hex(node)] = node_id
			node_counter += 1
			tail, a, b = def_use_parse(node)
			string_node = append_node(string_node, node_id, hex(node), tail)

	for edge in edges:
		if edge[0] == 'START':
			node1 = node_ids['START']
		elif hex(edge[0]) in node_ids.keys():
			node1 = node_ids[hex(edge[0])]
		else:
			node1 = hex(edge[0])
		if edge[1] == 'START':
			node2 = node_ids['START']
		elif hex(edge[1]) in node_ids.keys():
			node2 = node_ids[hex(edge[1])]
		else:
			node2 = hex(edge[1])
		string_edge = append_edge(string_edge, node1, node2)
		
	final_string = final_string_parse(string_node, string_edge)
	write_file(final_string, func_name)

def parse_indexed_operand(opnd_str):
	ret_list = re.findall('\[(.*)\]', opnd_str)
	split_str = ret_list[0].split('+')
	ret_str = ""

	ret_list2 = list()

	for s in split_str:

		if (s in reg_list):
			if (ret_str == ""):
				ret_str = s
			else:
				ret_str = ret_str + ", " + s
			ret_list2.append(s);

	return ret_list2, ret_str

def parse_use_operand(ea, idx):
	opnd_str = GetOpnd(ea, idx)
	opnd = GetOpType(ea, idx)

	ret_list = list()

	opnd_str = opnd_str.replace('dword ptr ', '')
	opnd_str = opnd_str.replace('byte ptr ', '')

	if (opnd == optype.REGISTER.value):
		ret_str = opnd_str
		ret_list.append(opnd_str)

	if (opnd == optype.IMM.value) or (opnd == optype.IMM_FAR.value) or (opnd == optype.IMM_NEAR.value):
		ret_str = ""

	if (opnd == optype.BASE_INDEX.value) or (opnd == optype.BASE_INDEX_DISP.value):
		ret_list2, ret_str = parse_indexed_operand(opnd_str)
		if (ret_str == ""):
			ret_str = opnd_str
		else:
			ret_str = ret_str + ", " + opnd_str
			ret_list.append(opnd_str)
		ret_list.extend(ret_list2)

	if (opnd  == optype.MEMORY.value):
		if "ds:" in opnd_str:
			ret_str = ""
		else:
			ret_str = opnd_str
			ret_list.append(opnd_str)

	return ret_list, ret_str

def parse_use_of_def_operand(ea, idx):
	opnd_str = GetOpnd(ea, idx)
	opnd = GetOpType(ea, idx)

	ret_list = list()

	opnd_str = opnd_str.replace('dword ptr ', '')
	opnd_str = opnd_str.replace('byte ptr ', '')

	if (opnd == optype.BASE_INDEX.value) or (opnd == optype.BASE_INDEX_DISP.value):
		ret_list2, ret_str = parse_indexed_operand(opnd_str)
		ret_list.extend(ret_list2)
	else:
		ret_str = ""

	return ret_list, ret_str

def parse_def_operand(ea):
	opnd_str = GetOpnd(ea, 0)

	opnd_str = opnd_str.replace('dword ptr ', '')
	opnd_str = opnd_str.replace('byte ptr ', '')

	ret_list = list()
	ret_list.append(opnd_str)

	return ret_list, opnd_str

def push_du(ea):
	def_string = "esp, [esp]"
	def_list = ["esp", "[esp]"]

	use_list, use_string = parse_use_operand(ea, 0) # + ", esp"
	use_string = use_string + ", esp"
	use_list.append("esp")

	du_str = "D: " + def_string + " U: " + use_string
	return du_str, def_list, use_list

def pop_du(ea):
	def_list, def_string = parse_def_operand(ea) # + ", esp"
	def_list.append("esp")
	def_string = def_string + ", esp"

	use_string = "esp, [esp]"
	use_list = ["esp", "[esp]"]

	du_str = "D: " + def_string + " U: " + use_string

	return du_str, def_list, use_list

def mov_du(ea):
	def_list, def_string = parse_def_operand(ea)

	use_list1, use_string1 = parse_use_of_def_operand(ea, 0)
	use_list2, use_string2 = parse_use_operand(ea, 1)

	if (use_string1 != "" and use_string2 != ""):
		use_string = use_string1 + ", " + use_string2
		use_list = use_list1
		use_list.extend(use_list2)
	elif (use_string2 == ""):
		use_string = use_string1
		use_list = use_list1
	elif (use_string1 == ""):
		use_string =  use_string2
		use_list = use_list2
	else:
		use_string = ""
		use_list = list()

	du_str = "D: " + def_string + " U: " + use_string
	return du_str, def_list, use_list
	
def binary_alu_du(ea):
	def_list, def_string = parse_def_operand(ea) # + ", eflags"
	def_string = def_string + ", eflags"
	def_list.append("eflags")

	use_list1, use_string1 = parse_use_operand(ea, 0)
	use_list2, use_string2 = parse_use_operand(ea, 1)

	if (use_string1 != "" and use_string2 != ""):
		use_string = use_string1 + ", " + use_string2
		use_list = use_list1
		use_list.extend(use_list2)
	elif (use_string2 == ""):
		use_string = use_string1
		use_list = use_list1
	elif (use_string1 == ""):
		use_string =  use_string2
		use_list = use_list2
	else:
		use_string = ""
		use_list = list()

	du_str = "D: " + def_string + " U: " + use_string
	return du_str, def_list, use_list

def unary_alu_du(ea):
	def_list, def_string = parse_def_operand(ea) # + ", eflags"
	def_string = def_string + ", eflags"
	def_list.append("eflags")

	use_list, use_string = parse_use_operand(ea, 0)

	du_str = "D: " + def_string + " U: " + use_string
	return du_str, def_list, use_list

def cmp_du(ea):
	def_string = "eflags"
	def_list = ["eflags"]

	use_list1, use_string1 = parse_use_operand(ea, 0)
	use_list2, use_string2 = parse_use_operand(ea, 1)

	if (use_string1 != "" and use_string2 != ""):
		use_string = use_string1 + ", " + use_string2
		use_list = use_list1
		use_list.extend(use_list2)
	elif (use_string2 == ""):
		use_string = use_string1
		use_list = use_list1
	elif (use_string1 == ""):
		use_string =  use_string2
		use_list = use_list2
	else:
		use_string = ""
		use_list = list()

	du_str = "D: " + def_string + " U: " + use_string
	return du_str, def_list, use_list

def lea_du(ea):
	def_list, def_string = parse_def_operand(ea)

	use_list1, use_string1 = parse_use_of_def_operand(ea, 0)
	use_list2, use_string2 = parse_use_of_def_operand(ea, 1)

	if (use_string1 != "" and use_string2 != ""):
		use_string = use_string1 + ", " + use_string2
		use_list = use_list1
		use_list.extend(use_list2)
	elif (use_string2 == ""):
		use_string = use_string1
		use_list = use_list1
	elif (use_string1 == ""):
		use_string =  use_string2
		use_list = use_list2
	else:
		use_string = ""
		use_list = list()

	du_str = "D: " + def_string + " U: " + use_string
	return du_str, def_list, use_list

def leave_du(ea):
	def_string = "esp, ebp"
	def_list = ["esp", "ebp"]

	use_string = "ebp, [ebp]"
	use_list = ["ebp", "[ebp]"]

	du_str = "D: " + def_string + " U: " + use_string
	return du_str, def_list, use_list

def call_du(ea):
	def_string = ""
	def_list = list()

	use_list, use_string = parse_use_operand(ea, 0)

	du_str = "D: " + def_string + " U: " + use_string
	return du_str, def_list, use_list

def ret_du(ea):
	def_string = ""
	def_list = list()

	ret_opnd = GetOpType(ea, 0)
	use_string = ""
	use_list = list()

	if (ret_opnd == optype.IMM.value) or (ret_opnd == optype.IMM_FAR.value) or (ret_opnd == optype.IMM_NEAR.value):
		use_string = "esp"
		use_list.append("esp")
		def_string = "esp"
		def_list.append("esp")

	du_str = "D: " + def_string + " U: " + use_string
	return du_str, def_list, use_list

def jump_du(ea):
	def_string = ""
	def_list = list()

	use_string = ""
	use_list = list()

	du_str = "D: " + def_string + " U: " + use_string
	return du_str, def_list, use_list

def jump_cond_du(ea):
	def_string = ""
	def_list = list()

	use_string = "eflags"
	use_list = ["eflags"]

	du_str = "D: " + def_string + " U: " + use_string
	return du_str, def_list, use_list

def stosd_du(ea):
	def_string = "edi, [edi], ecx"
	def_list = ["edi", "[edi]", "ecx"]

	use_string = "eax, edi, eflags, ecx"
	use_list = ["eax", "edi", "eflags", "ecx"]

	du_str = "D: " + def_string + " U: " + use_string
	return du_str, def_list, use_list

def setnz_du(ea):
	def_list, def_string = parse_def_operand(ea)

	use_string = "eflags"
	use_list = ["eflags"]

	du_str = "D: " + def_string + " U: " + use_string
	return du_str, def_list, use_list

def default_du(ea):
	def_string = "UNDEF"
	def_list = list()

	use_string = "UNDEF"
	use_list = list()

	du_str = "D: " + def_string + " U: " + use_string

	return du_str, def_list, use_list

def mnem_to_func(mnem, ea):
	if (mnem == "push"):
		du_str, def_list, use_list = push_du(ea)
	elif (mnem == "pop"):
		du_str, def_list, use_list = pop_du(ea)
	elif (mnem == "mov") or (mnem == 'movsx'):
		du_str, def_list, use_list = mov_du(ea)
	elif (mnem == "lea"):
		du_str, def_list, use_list = lea_du(ea)
	elif (mnem == "cmp") or (mnem == "test"):
		du_str, def_list, use_list = cmp_du(ea)
	elif (mnem == "leave"):
		du_str, def_list, use_list = leave_du(ea)
	elif (mnem == "ret") or (mnem == "retn"):
		du_str, def_list, use_list = ret_du(ea)
	elif (mnem == "call"):
		du_str, def_list, use_list = call_du(ea)
	elif (mnem == "jmp"):
		du_str, def_list, use_list = jump_du(ea)
	elif (mnem == 'ja') or (mnem == 'jb') or (mnem == 'jbe') or (mnem == 'jg') or (mnem == 'jl') or  (mnem == 'jle') or (mnem == 'jnb') or (mnem == 'jnz') or (mnem == 'jz'):
		du_str, def_list, use_list = jump_cond_du(ea)
	elif (mnem == 'add') or (mnem == 'and') or (mnem == 'sub') or (mnem == 'or') or (mnem == 'xor') or (mnem == 'shl') or (mnem == 'sal') or (mnem == 'shr') or (mnem == 'sar'):
		du_str, def_list, use_list = binary_alu_du(ea)
	elif (mnem == 'inc') or (mnem == 'dec'):
		du_str, def_list, use_list = unary_alu_du(ea)
	elif (mnem == 'imul') or (mnem == 'mul') or (mnem == 'idiv') or (mnem == 'div'):
		du_str, def_list, use_list = binary_alu_du(ea)
	elif (mnem == "stosd"):
		du_str, def_list, use_list = stosd_du(ea)
	elif (mnem == "setnz"):
		du_str, def_list, use_list = setnz_du(ea)
	else:
		du_str, def_list, use_list = default_du(ea)

	return du_str, def_list, use_list

def remove_rep(du_str):
	split_str = du_str.split('U:')
	def_str = split_str[0]
	use_str = split_str[1]

	use_split = use_str.split(',')
	use_list = []
	[use_list.append(x) for x in use_split if x not in use_list]
	use_str = ','.join(use_list)
	
	def_use_str = def_str + "U:" + use_str
	return def_use_str

#creates Def-use string and lists for a given instruction (ea)
def def_use_parse(ea):
	mnem_str = GetMnem(ea)
	du_str, def_list, use_list = mnem_to_func(mnem_str, ea)
	du_str = remove_rep(du_str)
	return du_str, def_list, use_list

#returns def list for a given instruction (ea)
def get_def_operand(ea):
	du_str, def_list, use_list = def_use_parse(ea)
	return def_list

#returns use list for a given instruction (ea)
def get_use_operand(ea):
	du_str, def_list, use_list = def_use_parse(ea)
	return use_list

def takeFirst(elem):
	return elem[0]

#function to find instructions that are arguments for a call
def find_call_data_dependence_list(instr, num_arg, f_start, f_end):

	#prepare list of instructions that flow to "instr"
	to_refs = CodeRefsTo(instr, 1)
	to_refs = list(filter(lambda x: x>=f_start and x<=f_end, to_refs))

	refs = list()
	dd_list = list()

	#prepare list of pairs where (tr, n) refer to flow instruction and number of arguments pushed
	for tr in to_refs:
		refs.append((tr, num_arg))

	for r, n in refs:
		mnem = GetMnem(r);

		#if 'n' is 1 and 'r' is a push instruction, this push adds the last argument for call
		if mnem == "push" and n == 1:
			dd_list.append(r)
			n = 0
		else:
			#Append argument corresponding to push to dd_list and decrement number of arguments to handle
			if mnem == "push":
				n -= 1
				dd_list.append(r)
			#Any pop it uses some memory which has been "pushed" but is not an argument to call
			elif mnem == "pop":
				n += 1

			#find references to 'r' and add (r, n) to refs list
			#to handle push arguments that span branches and jumps
			to_refs = CodeRefsTo(r, 1)
			to_refs = list(filter(lambda x: x>=f_start and x<=f_end, to_refs))

			for tr in to_refs:
				tmp = [x[0] for x in refs]
				if tr not in tmp:
					refs.append((tr, n))

	dd_list = list(set(dd_list))
	return dd_list

#function to find instruction that writes to memory for a pop
def find_pop_data_dependence_list(instr, f_start, f_end):

	#prepare list of instructions that flow to "instr"
	to_refs = CodeRefsTo(instr, 1)
	to_refs = list(filter(lambda x: x>=f_start and x<=f_end, to_refs))

	refs = list()
	dd_list = list()

	#prepare list of pairs where (tr, n) refer to flow instruction and number of pops to handle
	for tr in to_refs:
		refs.append((tr, 1))

	for r, n in refs:
		mnem = GetMnem(r);

		#if 'n' is 1 and 'r' is a push instruction, this instruction writes to memory that is popped
		if mnem == "push" and n == 1:
			dd_list.append(r)
			n = 0
		else:
			#Since 'n' is not 1, current push writes to memory which is popped by a different pop instruction
			if mnem == "push":
				n -= 1
			#If a pop instruction is seen, a corresponding push should be ignored for current 'pop'
			#Increase 'n' (number of pops to handle) to keep track of main 'pop' instruction memory
			elif mnem == "pop":
				n += 1
			#Few subroutines clear stack by calling "retn N"; each N corresponds to N/4 pops
			elif mnem == "call":
				call_str = GetOpnd(r, 0)
				call_addr = LocByName(call_str)
				call_end = FindFuncEnd(call_addr)
				call_end = PrevHead(call_end)
				ret_mnem = GetMnem(call_end)
				if ret_mnem == "ret" or ret_mnem == "retn":
					b = GetOpnd(call_end, 0)
					if b != '':
						if b[-1] == 'h':
							b = int(b[:-1], 16)
						else:
							b = int(b)
						n = n + (b/4)
			#'add esp, N' instructions is equivalent to N/4 pops
			elif mnem == "add":
				opnd = GetOpnd(r, 0)
				if opnd == "esp":
					b = GetOpnd(r, 1)
					if b[-1] == 'h':
						b = int(b[:-1], 16)
					else:
						b = int(b)
					n = n + (b/4);

			#find references to 'r' and add (r, n) to refs list
			#to handle push arguments that span branches and jumps
			to_refs = CodeRefsTo(r, 1)
			to_refs = list(filter(lambda x: x>=f_start and x<=f_end, to_refs))

			for tr in to_refs:
				tmp = [x[0] for x in refs]
				if tr not in tmp:
					refs.append((tr, n))

	dd_list = list(set(dd_list))
	return dd_list

#function to find data dependence for a given instr with function bounds f_start and f_end
def find_data_dependence_list(instr, f_start, f_end):

	dd_list = list()

	#starting instruction depends on START
	if instr == f_start:
		dd_list.append('START')

	mnem = GetMnem(instr)

	#find argument list for call instructions
	if mnem == "call":
		addr_list = get_arg_addrs(instr)
		if addr_list:
			num_arg = len(addr_list)
			addr_list = find_call_data_dependence_list(instr, num_arg, f_start, f_end)
			dd_list.extend(addr_list)

	#find push instructions which pop is dependent on
	if mnem == "pop":
		pop_list = find_pop_data_dependence_list(instr, f_start, f_end)
		if pop_list:
			dd_list.extend(pop_list)

	#find data dependencies for elements in the used list (BFS)
	use_list = get_use_operand(instr)

	if use_list:
		for u in use_list:
			#[esp] is used by pop instructions handled in find_pop_data_dependence_list()
			if u == "[esp]":
				continue

			#prepare list of instructions that flow to "instr"
			refs = CodeRefsTo(instr, 1)
			refs = list(filter(lambda x: x>=f_start and x<=f_end, refs))

			for r in refs:
				def_list = get_def_operand(r)

				#if 'u' is defined by instruction 'r', add it to dd_list and end current iteration
				#else find all references to 'r' and add it to 'refs'
				if def_list:
					if u in def_list:
						dd_list.append(r)
					else:
						if r == f_start:
							dd_list.append("START")
							continue
						to_refs = CodeRefsTo(r, 1)
						to_refs = list(filter(lambda x: x>=f_start and x<=f_end, to_refs))
						for tr in to_refs:
							if tr in refs:
								pass
							else:
								refs.append(tr)
				else:
					if r == f_start:
						dd_list.append("START")
						continue
					to_refs = CodeRefsTo(r, 1)
					to_refs = list(filter(lambda x: x>=f_start and x<=f_end, to_refs))
					for tr in to_refs:
						if tr in refs:
							pass
						else:
							refs.append(tr)
	return dd_list

#function traverses all instructions in a subroutine calculating data dependencies for each instruction
def cyclomatic_complexity(function_ea):

	f_start = function_ea
	f_end = FindFuncEnd(function_ea)
	edges = Set()
	nodes = Set((f_start,))

	# For each defined element in the function.
	for instr in Heads(f_start, f_end):

		# If the element is an instruction
		if isCode(GetFlags(instr)):

			#du_str, def_list, use_list = def_use_parse(instr)
			data_dep = find_data_dependence_list(instr, f_start, f_end)
			if data_dep:
				for elem in data_dep:
					edges.add((elem, instr))

			# Get the references made from the current instruction and keep only the ones local to
			# the function.
			refs = CodeRefsFrom(instr, 1)
			refs = Set(filter(lambda x: x>=f_start and x<=f_end, refs))

			if refs:
				# Update the nodes found so far.
				nodes.union_update(refs)

				# For each of the references found, and edge is created.
				for r in refs:
					# CHANGE 1
					pass
					# edges.add((instr, r))

	sortededges = sorted(edges, key = takeFirst);

	sortednodes = sorted(nodes);

	print GetFunctionName(instr) + ' ' + str(len(sortededges)) + ' ' + str(len(sortednodes))

	return sortededges, sortednodes

def do_functions(ea):
	edges = dict()
	nodes = dict()
	func_start = dict()

	# For each of the segments
	for seg_ea in Segments():
		# For each of the functions
		for function_ea in Functions(seg_ea, SegEnd(seg_ea)):
			func_start[GetFunctionName(function_ea)] = function_ea
			edges[GetFunctionName(function_ea)], nodes[GetFunctionName(function_ea)] = cyclomatic_complexity(function_ea)
	

	return edges, nodes, func_start

def do_one_function(ea):
	edges = dict()
	nodes = dict()
	func_start = dict()

	pfn = func_t()
	pfn = get_func(ea)
	func_rangeset = rangeset_t()
	func_range = range_t()
	get_func_ranges(func_rangeset, pfn)
	func_range = func_rangeset.getrange(0)

	ea = func_range.start_ea
	edges[GetFunctionName(ea)], nodes[GetFunctionName(ea)] = cyclomatic_complexity(ea)

	return edges, nodes, func_start



class testplu_t(idaapi.plugin_t):
	flags = idaapi.PLUGIN_UNL
	comment = "This is a comment"
	help = "This is help"
	wanted_name = "Lab 4 plugin"
	wanted_hotkey = "Alt-F8"

	def init(self):
		return idaapi.PLUGIN_OK

	def run(self, arg):
		ea = ScreenEA()

		# Collect data
		#edges_all, nodes_all, func_start = do_one_function(ea)
		edges_all, nodes_all, func_start = do_functions(ea)

		# Get the list of functions and sort it.
		functions = edges_all.keys()
		functions.sort()

		# Print the cyclomatic complexity for each of the functions.
		for f in functions:
			edges = []
			nodes = []

			edges = edges + list(edges_all[f])
			nodes = nodes + list(nodes_all[f])

			parse_output(nodes, edges, func_start[f])

	def term(self):
		pass

def PLUGIN_ENTRY():
	return testplu_t()


repository: techbliss/IDA_PRO_Stylesheets
file: plugins/Style_Paster_Plugin.py
description: Style sheets for ida pro 7.2 & 7.3
wanted_name: StyleSheet Paster
comment: Run me
created_at: 2019-08-11 18:39:24+00:00
pushed_at: 2023-04-09 14:52:12+00:00
url:
https://github.com/techbliss/IDA_PRO_Stylesheets/blob/321c667fa396c345d32b81bd0cae397fdd142c31/plugins/Style_Paster_Plugin.py

content:
# Created by: Storm Shadow http://www.techbliss.org
import os
import ida_idaapi, ida_kernwin
from ida_diskio import *
import idc
from idc import *
from ida_idaapi import *
import sys
sys.path.insert(0 , ida_diskio.idadir("plugins\\Style_Paster")) #change path to the icon
import ico
from ico import *

PLUGIN_VERSION = "1.0" #change Plugin Version
IDAVERISONS    = "IDA PRO 7.2+" #change Ida version
AUTHORS        = "Storm Shadow" #change author
DATE           = "2019" #change date
TWITTER        = "Twitter @zadow28" #change social media

def banner():
    banner_options = (PLUGIN_VERSION, AUTHORS, DATE, TWITTER, IDAVERISONS)
    banner_titles = "StyleSheet Paster v%s - (c) %s - %s - %s - %s" % banner_options #change Python editor

# print plugin banner
    print("---[" + banner_titles + "]---\n")

banner()

# 1) Create the handler class
class MyEditorHandlerPaste(ida_kernwin.action_handler_t): # the class name must be unique for each plugin nb same as line (53)
    def __init__(self):
        ida_kernwin.action_handler_t.__init__(self)

    # Run editor when invoked.
    def activate(self, ctx):
        g = globals()
        idahome = ida_diskio.idadir("plugins\\Style_Paster") #change to set plguin path where the main plguin script is
        ida_idaapi.IDAPython_ExecScript(idahome + "\\Style_Paste_main.py", globals()) #change for the main plugin script

    def update(self, ctx):
        return ida_kernwin.AST_ENABLE_ALWAYS

class ripeyePaste(ida_idaapi.plugin_t): #change class to unique name for each plugin , also change line (103)
    flags = ida_idaapi.PLUGIN_FIX #different flags this is for plugin visible at startup
    comment = "Run me" #help me text
    help = "StyleSheet Paster" #help text
    wanted_name = "StyleSheet Paster" #change the plugins name
    wanted_hotkey = "" #the tooltip goes away in menu when setting it here ,DONT DO it! and only is shown in File/Plugins menu


    def editor_menuactionPaste(self): #change for something unique.also change line (90)
        action_desc = ida_idaapi.action_desc_t(
            'my:editoractionPaste',  # The action name. This acts like an ID and must be unique same as line (64), (68)
            'StyleSheet Paster',  # The action text.
            MyEditorHandlerPaste(),  # The action handler must be unique , also change line (99)
            'Ctrl+P',  # Optional: the action shortcut DO IT  HERE!
            'Paste and test your custom styles',  # Optional: the action tooltip (available in menus/toolbar)
            ida_idaapi.load_custom_icon(idaapi.idadir("plugins\\Style_Paster")+"\\iconss.png")  # hackish load action icon , if no custom icon use number from 1-150 from internal ida
        )

        # 3) Register the action
        ida_idaapi.register_action(action_desc)

        ida_idaapi.attach_action_to_menu(
            'File/StyleSheet Paster...',  # The relative path of where to add the action
            'my:editoractionPaste',  # The action ID (see line(51))
            ida_idaapi.SETMENU_APP)  # We want to append the action after the 'Manual instruction...

        form = idaapi.get_current_tform()
        ida_idaapi.attach_action_to_popup(form, None, "my:editoractionPaste", None) # The action ID (see line(51)

    def init(self):
        """
        This is called by IDA when it is loading the plugin.
        """
        #self._icon_id_file = idaapi.BADADDR
        # attempt plugin initialization
        try:
            self._install_plugin()

        # failed to initialize or integrate the plugin, log and skip loading
        except Exception as e:
            form = ida_kernwin.get_current_tform()
            pass

        return PLUGIN_KEEP

    def _install_plugin(self):
        """
        Initialize & integrate the plugin into IDA.
        """
        self.editor_menuactionPaste() #same as line (49)
        self._init()

    def term(self):
        pass

    def run(self, arg = 0):
        #we need the calls again if we wanna load it via File/Plugins/editor
        ida_kernwin.msg("StyleSheet Paster Loaded to menu \n use Alt+P hot key to quick load ")
        hackish = MyEditorHandlerPaste() #must be the same as line (53) change also hackish =  must be unique also line (100)
        hackish.activate(self)  #hackish must the same as line hackish = (99)

def PLUGIN_ENTRY():
    return ripeyePaste() #must be unique for each plugin and same as line 43


repository: thezdi/PoC
file: VMware/patchgap-blog/rizzo/rizzo_string.py
description: Proofs-of-concept
wanted_name: Rizzo
comment: Function signature
created_at: 2018-06-06 22:30:00+00:00
pushed_at: 2024-09-03 15:37:21+00:00
url:
https://github.com/thezdi/PoC/blob/4a3b646c853ecfa083031487001282eef2a9f66c/VMware/patchgap-blog/rizzo/rizzo_string.py

content:
##########################################################################
# An IDAPython plugin that generates "fuzzy" function signatures that can
# be shared and applied amongst different IDBs.
#
# There are multiple sets of signatures that are generated:
#
#   o "Formal" signatures, where functions must match exactly
#   o "Fuzzy" signatures, where functions must only resemble each other
#     in terms of data/call references.
#   o String-based signatures, where functions are identified based on
#     unique string references.
#   o Immediate-based signatures, where functions are identified based
#     on immediate value references.
#
# These signatures are applied based on accuracy, that is, formal
# signatures are applied first, then string and immediate based
# signatures, and finally fuzzy signatures.
#
# Further, functions are identified based on call references. Consider,
# for example, two functions, one named 'foo', the other named 'bar'.
# The 'foo' function is fairly unique and a reliable signature is easily
# generated for it, but the 'bar' function is more difficult to reliably
# identify. However, 'foo' calls 'bar', and thus once 'foo' is identified,
# 'bar' can also be identified by association.
#
# Craig Heffner
# @devttys0
##########################################################################
import idc
import idaapi
import idautils

import os
import time
import pickle       # http://natashenka.ca/pickle/
import collections

import ida_shims

class RizzoSignatures(object):
    """
    Simple wrapper class for storing signature info.
    """
    SHOW = []

    def __init__(self):
        self.fuzzy = {}
        self.formal = {}
        self.strings = {}
        self.functions = {}
        self.immediates = {}

        self.fuzzydups = set()
        self.formaldups = set()
        self.stringdups = set()
        self.immediatedups = set()

    def show(self):
        if not self.SHOW:
            return

        print "\n\nGENERATED FORMAL SIGNATURES FOR:"
        for (key, ea) in self.formal.iteritems():
            func = RizzoFunctionDescriptor(self.formal, self.functions, key)
            if func.name in self.SHOW:
                print func.name

        print "\n\nGENERATRED FUZZY SIGNATURES FOR:"
        for (key, ea) in self.fuzzy.iteritems():
            func = RizzoFunctionDescriptor(self.fuzzy, self.functions, key)
            if func.name in self.SHOW:
                print func.name


class RizzoStringDescriptor(object):
    """
    Wrapper class for easily accessing necessary string information.
    """

    def __init__(self, string):
        self.ea = string.ea
        self.value = str(string)
        self.xrefs = [x.frm for x in idautils.XrefsTo(self.ea)]


class RizzoBlockDescriptor(object):
    """
    Code block info is stored in tuples, which minimize pickle storage space.
    This class provides more Pythonic (and sane) access to values of interest
    for a given block.
    """

    def __init__(self, block):
        self.formal = block[0]
        self.fuzzy = block[1]
        self.immediates = block[2]
        self.functions = block[3]

    def match(self, nblock, fuzzy=False):
        # TODO: Fuzzy matching at the block level gets close, but produces a
        #  higher number of false positives; for example, it confuses
        #  hmac_md5 with hmac_sha1.
        return (self.formal == nblock.formal and
                len(self.immediates) == len(nblock.immediates) and
                len(self.functions) == len(nblock.functions))


class RizzoFunctionDescriptor(object):
    """
    Function signature info is stored in dicts and tuples, which minimize pickle
    storage space. This class provides more Pythonic (and sane) access to
    values of interest for a given function.
    """

    def __init__(self, signatures, functions, key):
        self.ea = signatures[key]
        self.name = functions[self.ea][0]
        self.blocks = functions[self.ea][1]


class Rizzo(object):
    """
    Workhorse class which performs the primary logic and functionality.
    """

    DEFAULT_SIGNATURE_FILE = "rizzo.sig"

    def __init__(self, sigfile=None):
        if sigfile:
            self.sigfile = sigfile
        else:
            self.sigfile = self.DEFAULT_SIGNATURE_FILE

        # Useful for identifying string xrefs from individual instructions
        self.strings = {}
        for string in idautils.Strings():
            self.strings[string.ea] = RizzoStringDescriptor(string)

        start = time.time()
        self.signatures = self.generate()
        end = time.time()

        print "Generated %d formal signatures and %d fuzzy signatures for %d " \
              "functions in %.2f seconds." % (len(self.signatures.formal),
                                              len(self.signatures.fuzzy),
                                              len(self.signatures.functions),
                                              (end-start))

    def save(self):
        print ("Saving signatures to %s..." % self.sigfile),
        fp = open(self.sigfile, "wb")
        pickle.dump(self.signatures, fp)
        fp.close()
        print "done."

    def load(self):
        print ("Loading signatures from %s..." % self.sigfile),
        #fp = open(self.sigfile, "rb")
        fp=open(self.sigfile,"rb").readlines()
        fp=map(lambda x:x.replace(b"\r\n","\n"),fp)
        with open(self.sigfile,"wb") as j:
            for i in fp:
                j.write(i)
                
        sigs = pickle.load(open(self.sigfile,"rb"))
        #fp.close()
        print ("done.")
        return sigs
		
    def sighash(self, value):
        return hash(str(value)) & 0xFFFFFFFF

    def block(self, block):
        """
        Returns a tuple:
        ([formal, block, signatures], [fuzzy, block, signatures],
        set([unique, immediate, values]), [called, function, names])
        """
        formal = []
        fuzzy = []
        functions = []
        immediates = []

        ea = ida_shims.start_ea(block)
        while ea < ida_shims.end_ea(block):
            insn = ida_shims.decode_insn(ea)

            # Get a list of all data/code refs from the current instruction
            drefs = [x for x in idautils.DataRefsFrom(ea)]
            crefs = [x for x in idautils.CodeRefsFrom(ea, False)]

            # Add all instruction mnemonics to the formal block hash
            formal.append(ida_shims.print_insn_mnem(ea))

            # If this is a call instruction, be sure to note the name of the
            # function being called. This is used to apply call-based
            # signatures to functions.
            #
            # For fuzzy signatures, we can't use the actual name or EA of the
            # function, but rather just want to note that a function call was
            # made.
            #
            # Formal signatures already have the call instruction mnemonic,
            # which is more specific than just saying that a call was made.
            if idaapi.is_call_insn(ea):
                for cref in crefs:
                    func_name = ida_shims.get_name(cref)
                    if func_name:
                        functions.append(func_name)
                        fuzzy.append("funcref")
            # If there are data references from the instruction, check to see
            # if any of them are strings. These are looked up in the
            # pre-generated strings dictionary.
            #
            # String values are easily identifiable, and are used as part of
            # both the fuzzy and the formal signatures.
            #
            # It is more difficult to determine if non-string values are
            # constants or not; for both fuzzy and formal signatures, just use
            # "data" to indicate that some data was referenced.
            elif drefs:
                for dref in drefs:
                    if self.strings.has_key(dref):
                        formal.append(self.strings[dref].value)
                        fuzzy.append(self.strings[dref].value)
                    else:
                        formal.append("dataref")
                        fuzzy.append("dataref")
            # If there are no data or code references from the instruction, use
            # every operand as part of the formal signature.
            #
            # Fuzzy signatures are only concerned with interesting immediate
            # values, that is, values that are greater than 65,535, are not
            # memory addresses, and are not displayed as negative values.
            elif not drefs and not crefs:
                ops = ida_shims.get_operands(insn)
                for n in range(0, len(ops)):
                    opnd_text = ida_shims.print_operand(ea, n)
                    formal.append(opnd_text)
                    if ops[n].type == idaapi.o_imm and \
                            not opnd_text.startswith('-'):
                        if ops[n].value >= 0xFFFF:
                            if ida_shims.get_full_flags(ops[n].value) == 0:
                                fuzzy.append(str(ops[n].value))
                                immediates.append(ops[n].value)

            ea = ida_shims.next_head(ea)

        return (self.sighash(''.join(formal)),
                self.sighash(''.join(fuzzy)),
                immediates, functions)

    def function(self, func):
        """
        Returns a list of blocks.
        """
        blocks = []

        for block in idaapi.FlowChart(func):
            blocks.append(self.block(block))

        return blocks

    def generate(self):
        signatures = RizzoSignatures()

        # Generate unique string-based function signatures
        for (ea, string) in self.strings.iteritems():
            # Only generate signatures on reasonably long strings with one xref
            if len(string.value) >= 8 and len(string.xrefs) == 1:
                func = idaapi.get_func(string.xrefs[0])
                if func:
                    str_hash = self.sighash(string.value)

                    # Check for and remove string duplicate signatures (the same
                    # string can appear more than once in an IDB).
                    # If no duplicates, add this to the string signature dict.
                    if str_hash in signatures.strings:
                        del signatures.strings[str_hash]
                        signatures.stringdups.add(str_hash)
                    elif str_hash not in signatures.stringdups:
                        signatures.strings[str_hash] = ida_shims.start_ea(func)

        # Generate formal, fuzzy, and immediate-based function signatures
        for ea in idautils.Functions():
            func = idaapi.get_func(ea)
            if func:
                start_ea = ida_shims.start_ea(func)
                signatures.functions[start_ea] = (ida_shims.get_name(start_ea), [])

        # These need not be maintained across function calls,
        # and only add to the size of the saved signature file.
        signatures.fuzzydups = set()
        signatures.formaldups = set()
        signatures.stringdups = set()
        signatures.immediatedups = set()

        # DEBUG
        signatures.show()

        return signatures

    def match(self, extsigs):
        fuzzy = {}
        formal = {}
        strings = {}
        immediates = {}

        # Match string based function signatures
        start = time.time()
        for (extsig, ext_func_ea) in extsigs.strings.iteritems():
            if extsig in self.signatures.strings:
                curr_fun = RizzoFunctionDescriptor(
                    self.signatures.strings, self.signatures.functions, extsig)
                new_fun = RizzoFunctionDescriptor(
                    extsigs.strings, extsigs.functions, extsig)
                strings[curr_fun] = new_fun
        end = time.time()
        print "Found %d string matches in %.2f seconds." % (len(strings),
                                                            (end-start))

        # Return signature matches in the order we want them applied
        # The second tuple of each match is set to True if it is a fuzzy match,
        # e.g.:   ((match, fuzzy), (match, fuzzy), ...)
        return ((formal, False),
                (strings, False),
                (immediates, False),
                (fuzzy, True))

    def rename(self, ea, name):
        # Don't rely on the name in curfunc, it could have already been renamed
        curname = ida_shims.get_name(ea)
        # Don't rename if the name is a special identifier, or if the ea has
        # already been named
        if curname.startswith('sub_') and \
                name.split('_')[0] not in \
                ['sub', 'loc', 'unk', 'dword', 'word', 'byte']:
            # Don't rename if the name already exists in the IDB
            if ida_shims.get_name_ea_simple(name) == idc.BADADDR:
                if ida_shims.set_name(ea, name):
                    ida_shims.set_func_flags(
                        ea, (idc.get_func_flags(ea) | idc.FUNC_LIB))
                    return 1
        return 0

    def apply(self, extsigs):
        count = 0

        start = time.time()

        # This applies formal matches first, then fuzzy matches
        for (match, fuzzy) in self.match(extsigs):
            # Keeps track of all function names that we've identified candidate
            # functions for
            rename = {}

            for (curfunc, newfunc) in match.iteritems():
                if newfunc not in rename:
                    rename[newfunc.name] = []

                # Attempt to rename this function
                rename[newfunc.name].append(curfunc.ea)

                bm = {}
                duplicates = set()

                # Search for unique matching code blocks inside this function
                for nblock in newfunc.blocks:
                    nblock = RizzoBlockDescriptor(nblock)
                    for cblock in curfunc.blocks:
                        cblock = RizzoBlockDescriptor(cblock)

                        if cblock.match(nblock, fuzzy):
                            if bm.has_key(cblock):
                                del bm[cblock]
                                duplicates.add(cblock)
                            elif cblock not in duplicates:
                                bm[cblock] = nblock

                # Rename known function calls from each unique code block
                for (cblock, nblock) in bm.iteritems():
                    for n in range(0, len(cblock.functions)):
                        ea = idc.get_name_ea_simple(cblock.functions[n])
                        if ea != idc.BADADDR:
                            if rename.has_key(nblock.functions[n]):
                                rename[nblock.functions[n]].append(ea)
                            else:
                                rename[nblock.functions[n]] = [ea]

                # Rename the identified functions
                for (name, candidates) in rename.iteritems():
                    if candidates:
                        winner = \
                            collections.Counter(candidates).most_common(1)[0][0]
                        count += self.rename(winner, name)

        end = time.time()
        print "Renamed %d functions in %.2f seconds." % (count, (end-start))


def RizzoBuild(sigfile=None):
    print "Building Rizzo signatures, this may take a few minutes..."
    start = time.time()
    r = Rizzo(sigfile)
    r.save()
    end = time.time()
    print "Built signatures in %.2f seconds" % (end-start)


def RizzoApply(sigfile=None):
    print "Applying Rizzo signatures, this may take a few minutes..."
    start = time.time()
    r = Rizzo(sigfile)
    s = r.load()
    r.apply(s)
    end = time.time()
    print "Signatures applied in %.2f seconds" % (end-start)


def rizzo_produce(arg=None):
    fname = ida_shims.ask_file(1, "*.riz", "Save signature file as")
    if fname:
        if '.' not in fname:
            fname += ".riz"
        RizzoBuild(fname)


def rizzo_load(arg=None):
    fname = ida_shims.ask_file(0, "*.riz", "Load signature file")
    if fname:
        RizzoApply(fname)

try:
    class ProduceRizzoAction(idaapi.action_handler_t):
        def __init__(self):
            idaapi.action_handler_t.__init__(self)

        def activate(self, ctx):
            rizzo_produce()
            return 1

        def update(self, ctx):
            return idaapi.AST_ENABLE_ALWAYS


    class LoadRizzoAction(idaapi.action_handler_t):
        def __init__(self):
            idaapi.action_handler_t.__init__(self)

        def activate(self, ctx):
            rizzo_load()
            return 1

        def update(self, ctx):
            return idaapi.AST_ENABLE_ALWAYS
except AttributeError:
    pass


class RizzoPlugin(idaapi.plugin_t):
    flags = 0
    comment = "Function signature"
    help = ""
    wanted_name = "Rizzo"
    wanted_hotkey = ""
    produce_action_name = 'producerizzo:action'
    load_action_name = 'loadrizzo:action'
    menu_name = "Rizzo signature file..."
    produce_tooltip = "Produce rizzo signature file."
    load_tooltip = "Load rizzo signature file."
    menu_tab = 'File/'
    menu_context = []

    def init(self):
        if idaapi.IDA_SDK_VERSION >= 700:
            produce_desc = idaapi.action_desc_t(self.produce_action_name,
                                                self.menu_name,
                                                ProduceRizzoAction(),
                                                self.wanted_hotkey,
                                                self.produce_tooltip,
                                                199)

            load_desc = idaapi.action_desc_t(self.load_action_name,
                                             self.menu_name,
                                             LoadRizzoAction(),
                                             self.wanted_hotkey,
                                             self.load_tooltip,
                                             199)

            idaapi.register_action(produce_desc)
            idaapi.register_action(load_desc)

            idaapi.attach_action_to_menu('File/Produce file/',
                self.produce_action_name,
                idaapi.SETMENU_APP)
            idaapi.attach_action_to_menu('File/Load file/',
                self.load_action_name,
                idaapi.SETMENU_APP)
        else:
            self.menu_context.append(
                idaapi.add_menu_item(
                    os.path.join(self.menu_tab, 'Load file/'),
                    "Rizzo signature file...", "", 0, rizzo_load, (None,)))

            self.menu_context.append(
                idaapi.add_menu_item(
                    os.path.join(self.menu_tab, 'Produce file/'),
                    "Rizzo signature file...", "", 0, rizzo_produce, (None,)))

        return idaapi.PLUGIN_KEEP

    def term(self):
        if idaapi.IDA_SDK_VERSION >= 700:
            idaapi.detach_action_from_menu(
                self.menu_tab, self.produce_action_name)
            idaapi.detach_action_from_menu(
                self.menu_tab, self.load_action_name)
        else:
            if self.menu_context is not None:
                idaapi.del_menu_item(self.menu_context)
        return None

    def run(self, arg):
        return None

    def rizzo_script(self):
        idaapi.IDAPython_ExecScript(self.script, globals())


def PLUGIN_ENTRY():
    return RizzoPlugin()

rizzo_plugin = PLUGIN_ENTRY()
rizzo_plugin.init()


repository: zeropointdynamics/zemu-client
file: ida-plugin/zemu_ida.py
description: Client for generating dynamic binary traces using the Zemu API.
wanted_name: ZemuOverlay
comment: Load an overlay file generated by zemu.
created_at: 2019-10-01 20:08:10+00:00
pushed_at: 2020-03-11 23:00:02+00:00
url:
https://github.com/zeropointdynamics/zemu-client/blob/8c2c5a8f52bc0c9d95a6b8deb490f70bfb3cc934/ida-plugin/zemu_ida.py

content:
'''

Copyright (c) 2019 Zeropoint Dynamics, LLC

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

'''

import idc
import idaapi

class ApplyZemuOverlay(idaapi.action_handler_t):
    '''
    Applies the overlay from the user-selected overlay file.
    '''

    def __init__(self):
        idaapi.action_handler_t.__init__(self)

    def activate(self, ctx):
        import json

        filepath = idc.AskFile(False, '*.zmu;*.overlay;*',
            'Load Zemu Overlay...')
        if filepath is None:
            return
        f = open(filepath, 'r')
        zemu_data = f.read()
        f.close()

        zemu_data = zemu_data[len('DISAS\n'):]
        zemu_dump = json.loads(zemu_data)

        # Apply the overlay data
        for comment in zemu_dump['comments']:
            ea = comment['address']
            comment_text = str(comment['text'])
            color = comment.get('color', 0x73f0df)

            # Set color of instruction line
            idaapi.set_item_color(ea, color)
            idaapi.set_cmt(ea, comment_text, False)

            # Set function name if not already changed
            idc.GetFunctionAttr(ea, idc.FUNCATTR_START)
            name = idc.GetFunctionName(ea)
            if len(name) > 0 and name.startswith('zmu_') == False:
                idc.MakeName(ea, 'zmu_' + name)

        return 1

    def update(self, ctx):
        return idaapi.AST_ENABLE_ALWAYS


class zemuoverlay_t(idaapi.plugin_t):
    '''
    Adds a Zemu {View} menu option for loading an overlay.
    '''
    flags = 0
    comment = 'Load an overlay file generated by zemu.'
    help = comment
    wanted_name = 'ZemuOverlay'
    wanted_hotkey = ''
    menu_name = 'View/'
    menu_context = []

    def init(self):
        zemu_overlay_action = idaapi.action_desc_t(
            'zemuoverlay:action', 'Load Zemu Overlay...',
            ApplyZemuOverlay(), '', 'Load Zemu Overlay...', 199)

        idaapi.register_action(zemu_overlay_action)
        idaapi.attach_action_to_menu(
            'View/', 'zemuoverlay:action', idaapi.SETMENU_APP)

        return idaapi.PLUGIN_KEEP

    def term(self):
        pass

    def run(self, arg):
        pass

def PLUGIN_ENTRY():
    return zemuoverlay_t()


repository: Introspelliam/ida-python
file: examples/ex_prefix_plugin.py
wanted_name: user defined prefix
comment: This is a user defined prefix sample plugin
created_at: 2019-03-17 16:34:44+00:00
pushed_at: 2019-03-17 17:41:26+00:00
url:
https://github.com/Introspelliam/ida-python/blob/e3f6e6a763f442add7c2422cd410b2b03fd3e1fb/examples/ex_prefix_plugin.py

content:
from __future__ import print_function
import idaapi

PREFIX = idaapi.SCOLOR_INV + ' ' + idaapi.SCOLOR_INV

class prefix_plugin_t(idaapi.plugin_t):
    flags = 0
    comment = "This is a user defined prefix sample plugin"
    help = "This is help"
    wanted_name = "user defined prefix"
    wanted_hotkey = ""


    def user_prefix(self, ea, lnnum, indent, line, bufsize):
        #print("ea=%x lnnum=%d indent=%d line=%s bufsize=%d" % (ea, lnnum, indent, line, bufsize))

        if (ea % 2 == 0) and indent == -1:
            return PREFIX
        else:
            return ""


    def init(self):
        self.prefix_installed = idaapi.set_user_defined_prefix(8, self.user_prefix)
        if self.prefix_installed:
            print("prefix installed")

        return idaapi.PLUGIN_KEEP


    def run(self, arg):
        pass


    def term(self):
        if self.prefix_installed:
            idaapi.set_user_defined_prefix(0, None)
            print("prefix uninstalled!")


def PLUGIN_ENTRY():
    return prefix_plugin_t()


repository: SapphireServer/FFXIVOpcodes
file: xiv-opcode-parser.py
description: A shared effort to maintain FFXIV opcode mappings
wanted_name: Find FFXIV Opcodes
comment: Does magic and shit
created_at: 2019-11-14 10:44:37+00:00
pushed_at: 2023-05-22 22:36:35+00:00
url:
https://github.com/SapphireServer/FFXIVOpcodes/blob/0a1433d0e97ab8e1685bc155694a9e6bc91f54d7/xiv-opcode-parser.py

content:
"""
IDA Script to parse the ZoneDownHandler and generate a schema of it to identify packets cross-patch.
"""

import ida_search
import ida_ida
import idautils
import ida_idaapi
import idc
import ida_funcs
import ida_bytes
import idaapi
import ida_name

import json

from PyQt5.Qt import QApplication

main_jumptable = ida_idaapi.BADADDR

def log(str, indent=0):
    print('[FFXIV Opcode Parser] %s%s' % ('  ' * indent, str))

def set_clipboard(data):
    QApplication.clipboard().setText(data)

def set_clipboard_json(data):
    set_clipboard(json.dumps(data, indent=2, separators=(',', ':')))
    log('copied parsed data to clipboard')

def find_pattern(pattern):
    return ida_search.find_binary(ida_ida.cvar.inf.omin_ea, ida_ida.cvar.inf.omax_ea, pattern, 16, ida_search.SEARCH_DOWN)

def get_executable_watermark():
    rev_ea = find_pattern('48 8D 0D ? ? ? ? 45 33 F6 41 8B DE')

    ea = idc.GetOperandValue(rev_ea, 1)

    return idc.GetString(ea, -1, idc.ASCSTR_C)

def get_clean_rev(rev_str):
    return rev_str.replace('ff14', '').replace('/', '').replace('*', '')


def find_switch(func_ea):
    # get all chunks that belong to a function, because apparently they're not contiguous or some shit
    for (start_ea, end_ea) in idautils.Chunks(func_ea):
        for head in idautils.Heads(start_ea, end_ea):
            switch = idaapi.get_switch_info_ex(head)

            if switch != None:
                log('found switch @ %x, cases: %d' % (head, switch.get_jtable_size()))
                return (head, switch)

    return (None, None)

def find_block(ea, blocks):
    for block in blocks:
        if block.startEA == ea:
            return block

def ea_to_rva(ea):
    return ea - idaapi.get_imagebase()

def get_bytes_str(start_ea, end_ea):
    size = end_ea - start_ea

    bytes = []
    for ea in range(start_ea, end_ea):
        b = '{:02x}'.format(ida_bytes.get_byte(ea))
        bytes.append(b)

    return ' '.join(bytes)

def get_func_name(ea):
    name = ida_funcs.get_func_name(ea)
    demangled = ida_name.demangle_name(name, idc.get_inf_attr(idc.INF_LONG_DN))

    return demangled or name

def postprocess_func(fn, depth = 0):
    func = {
        'ea': fn.startEA,
        'rva': ea_to_rva(fn.startEA),
        'body': get_bytes_str(fn.startEA, fn.endEA)
    }

    # total aids
    switch_ea, switch = find_switch(fn.startEA)

    if switch and switch_ea != main_jumptable:
        sw = func['switch'] = {}

        res = idaapi.calc_switch_cases(switch_ea, switch)
        
        case_ids = []
        for case in res.cases:
            for i in case:
                case_ids.append(int(i))

        sw['cases'] = [i for i in set(case_ids)]

    else:
        func['switch'] = None

    return func

def process_func(func, start_ea, end_ea):
    for head in idautils.Heads(start_ea, end_ea):
        flags = idaapi.getFlags(head)
        if idaapi.isCode(flags):

            mnem = idc.GetMnem(head)

            if mnem == 'call' or mnem == 'jmp':
                op_ea = idc.GetOperandValue(head, 0)
                fn = ida_funcs.get_func(op_ea)

                if fn:
                    fn_info = postprocess_func(fn)

                    if fn_info:
                        func['calls'][get_func_name(op_ea)] = fn_info

def process_case(case, id):
    func = case['func'] = {}
    body = func['body'] = get_bytes_str(case['start_ea'], case['end_ea'])
    func['calls'] = {}

    process_func(func, case['start_ea'], case['end_ea'])



def run():
    func_ea = find_pattern('48 89 74 24 ? 57 48 83 EC 50 8B F2 49 8B F8')
    output = {}

    output['rev'] = get_executable_watermark()
    output['clean_rev'] = get_clean_rev(output['rev'])

    if func_ea == ida_idaapi.BADADDR:
        log('couldn''t find server opcode handler')
        return

    func_end_ea = idc.get_func_attr(func_ea, idc.FUNCATTR_END)
    output['size'] = func_size = func_end_ea - func_ea

    log('found opcode handler @ %x, end: %x, size: %x' % (func_ea, func_end_ea, func_size))
    output['rva'] = ea_to_rva(func_ea)
    output['start_ea'] = func_ea
    output['end_ea'] = func_end_ea

    # find switch
    head, switch = find_switch(func_ea)

    global main_jumptable
    main_jumptable = head

    if switch == None:
        log('failed to find switch in opcode handler')
        return
    
    # get switch cases
    res = idaapi.calc_switch_cases(head, switch)

    # get basic blocks
    blocks = idaapi.FlowChart(idaapi.get_func(func_ea))

    output['cases'] = case_infos = []

    for idx, case in enumerate(res.cases):
        case_ea = res.targets[idx];
        rel_ea = case_ea - func_ea

        case_info = {
            'rva': ea_to_rva(case_ea),
            'start_ea': case_ea,
            'rel_ea': rel_ea
        }

        block = find_block(case_ea, blocks)
        
        if block != None:
            case_info['end_ea'] = block.endEA;
            case_info['size'] = block.endEA - case_ea

            # -1 to make it actually clickable in the output window and it goes to the right place
            #log('end: %x, size: %x' % ((block.endEA - 1), case_info['size']), 2)

        else:
            log('failed to get block for %x' % case_ea)
            continue

        case_info['opcodes'] = [int(oc) for oc in case]

        case_infos.append(case_info)

    log('got %d case info objs, switch blocks: %d' % (len(case_infos), len(res.cases)))

    for k, v in enumerate(case_infos):
        process_case(v, k)

    set_clipboard_json(output)

class xiv_opcode_parser_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL

    wanted_name = 'Find FFXIV Opcodes'
    wanted_hotkey = ''

    comment = 'Does magic and shit'
    help = 'no'
 
    def init(self):
        return idaapi.PLUGIN_OK
 
    def run(self, arg):
        run()
 
    def term(self):
        pass
 
def PLUGIN_ENTRY():
    return xiv_opcode_parser_t()


repository: manesec/tools4mane
file: Config/gdb/idaHelper.py
description: Some tools/script make by mane to help for pentest.
wanted_name: auto-load-xmlrpc-server
comment: This plugin auto load simple xmlrpc server
created_at: 2022-02-01 11:01:45+00:00
pushed_at: 2025-05-03 18:41:37+00:00
url:
https://github.com/manesec/tools4mane/blob/ec69e48018d203ec01d3dc52745bdaf8270f7c5d/Config/gdb/idaHelper.py

content:
from __future__ import annotations

import datetime
import threading
import xmlrpc.client as xmlclient
from xml.sax.saxutils import escape
from xmlrpc.server import SimpleXMLRPCServer

import idaapi
import idc

listen = "0.0.0.0"
listen_port = 31337

class autoLoadXMLRPC(idaapi.plugin_t):
    print("[+] Power By mane ...")
    flags = idaapi.PLUGIN_DRAW
    comment = "This plugin auto load simple xmlrpc server"
    help = "No help is needed"
    wanted_name = "auto-load-xmlrpc-server"

    def init(self):
        
        # Wait for any processing to get done
        idaapi.auto_wait()

        # On Windows with NTFS filesystem a filepath with ':'
        # is treated as NTFS ADS (Alternative Data Stream)
        # and so saving file with such name fails
        dt = datetime.datetime.now().isoformat().replace(":", "-")

        # Save the database so nothing gets lost.
        idc.save_database(idc.get_idb_path() + "." + dt)


        DEBUG_MARSHALLING = False


        def create_marshaller(use_format=None, just_to_str=False):
            assert (
                use_format or just_to_str
            ), "Either pass format to use or make it converting the value to str."

            def wrapper(_marshaller, value, appender):
                if use_format:
                    marshalled = use_format % value
                elif just_to_str:
                    marshalled = f"<value><string>{escape(str(value))}</string></value>"

                if DEBUG_MARSHALLING:
                    print(f"Marshalled: '{marshalled}'")

                appender(marshalled)

            return wrapper


        xmlclient.Marshaller.dispatch[type(1 << 63)] = create_marshaller("<value><i8>%d</i8></value>")
        xmlclient.Marshaller.dispatch[int] = create_marshaller("<value><i8>%d</i8></value>")
        xmlclient.Marshaller.dispatch[idaapi.cfuncptr_t] = create_marshaller(just_to_str=True)

        host = "0.0.0.0"
        port = 31337

        mutex = threading.Condition()


        def wrap(f):
            def wrapper(*a, **kw):
                rv = []
                error = []

                def work():
                    try:
                        result = f(*a, **kw)
                        rv.append(result)
                    except Exception as e:
                        error.append(e)

                with mutex:
                    flags = idaapi.MFF_WRITE
                    if f == idc.set_color:
                        flags |= idaapi.MFF_NOWAIT
                        rv.append(None)
                    idaapi.execute_sync(work, flags)

                if error:
                    msg = f"Failed on calling {f.__module__}.{f.__name__} with args: {a}, kwargs: {kw}\nException: {str(error[0])}"
                    print("[!!!] ERROR:", msg)
                    raise error[0]

                return rv[0]

            return wrapper


        def register_module(module):
            for name, function in module.__dict__.items():
                if hasattr(function, "__call__"):
                    server.register_function(wrap(function), name)


        def decompile(addr):
            """
            Function that overwrites `idaapi.decompile` for xmlrpc so that instead
            of throwing an exception on `idaapi.DecompilationFailure` it just returns `None`.
            (so that we don't have to parse xmlrpc Fault's exception string on pwndbg side
            as it differs between IDA versions).
            """
            try:
                return idaapi.decompile(addr)
            except idaapi.DecompilationFailure:
                return None


        def get_decompile_coord_by_ea(cfunc, addr):
            if idaapi.IDA_SDK_VERSION >= 720:
                item = cfunc.body.find_closest_addr(addr)
                y_holder = idaapi.int_pointer()
                if not cfunc.find_item_coords(item, None, y_holder):
                    return None
                y = y_holder.value()
            else:
                lnmap = {}
                for i, line in enumerate(cfunc.pseudocode):
                    phead = idaapi.ctree_item_t()
                    pitem = idaapi.ctree_item_t()
                    ptail = idaapi.ctree_item_t()
                    ret = cfunc.get_line_item(line.line, 0, True, phead, pitem, ptail)
                    if ret and pitem.it:
                        lnmap[pitem.it.ea] = i
                y = None
                closest_ea = idaapi.BADADDR
                for ea, line in lnmap.items():
                    if closest_ea == idaapi.BADADDR or abs(closest_ea - addr) > abs(ea - addr):
                        closest_ea = ea
                        y = lnmap[ea]

            return y


        def decompile_context(addr, context_lines):
            cfunc = decompile(addr)
            if cfunc is None:
                return None
            y = get_decompile_coord_by_ea(cfunc, addr)
            if y is None:
                return cfunc
            lines = cfunc.get_pseudocode()
            retlines = []
            for lnnum in range(max(0, y - context_lines), min(len(lines), y + context_lines)):
                retlines.append(idaapi.tag_remove(lines[lnnum].line))
                if lnnum == y:
                    retlines[-1] = ">" + retlines[-1][1:]
            return "\n".join(retlines)


        def versions():
            """Returns IDA & Python versions"""
            import sys

            return {
                "python": sys.version,
                "ida": idaapi.get_kernel_version(),
                "hexrays": idaapi.get_hexrays_version() if idaapi.init_hexrays_plugin() else None,
            }


        server = SimpleXMLRPCServer((host, port), logRequests=False, allow_none=True)
        register_module(idaapi)
        register_module(
            idc
        )  # prioritize idc functions over above (e.g. idc.get_next_seg/ida_segment.get_next_seg)

        server.register_function(lambda a: eval(a, globals(), locals()), "eval")
        server.register_function(wrap(decompile))  # overwrites idaapi/ida_hexrays.decompile
        server.register_function(wrap(decompile_context), "decompile_context")  # support context decompile
        server.register_function(wrap(versions))
        server.register_introspection_functions()

        print(f"IDA Pro xmlrpc hosted on http://{host}:{port}")
        print("Call `shutdown()` to shutdown the IDA Pro xmlrpc server.")

        thread = threading.Thread(target=server.serve_forever)
        thread.daemon = True
        thread.start()


        def shutdown():
            global server
            global thread
            server.shutdown()
            server.server_close()
            del server
            del thread

    
    
        return idaapi.PLUGIN_KEEP # On garde le plugin chargé en mémoire

    def run(self, arg):
        return

    def term(self):
        pass

def PLUGIN_ENTRY():
    return autoLoadXMLRPC()


repository: ZjW1nd/Keypatch_9.0
file: keypatch9.0.py
description: Keypatch plugin for ida9.0 from @y0ush4's Issue
url: https://github.com/ZjW1nd/Keypatch_9.0/blob/2b71e4fed4e06c9129c79686a1aa8c56d8afb06a/keypatch9.0.py
wanted_name: Keypatch Patcher
comment: Keypatch plugin for IDA Pro (using Keystone framework)
created_at: 2025-01-19 11:22:07+00:00
pushed_at: 2025-02-04 04:22:55+00:00
content:
# -*- coding: utf-8 -*-

# Keypatch IDA Plugin, powered by Keystone Engine (http://www.keystone-engine.org).
# By Nguyen Anh Quynh & Thanh Nguyen, 2018.

# Keypatch is released under the GPL v2. See COPYING for more information.
# Find docs & latest version at http://keystone-engine.org/keypatch

# This IDA plugin includes 3 tools inside: Patcher, Fill Range & Search.
# Access to these tools via menu "Edit | Keypatch", or via right-click popup menu "Keypatch".

# Hotkey Ctrl-Alt-K opens either Patcher or "Fill Range" window, depending on context.
#  - If there is no code selection, hotkey opens Patcher dialog
#  - If a range of code is selected, hotkey opens "Fill Range" dialog

# To revert (undo) the last patching, choose menu "Edit | Keypatch | Undo last patching".
# To check for update version, choose menu "Edit | Keypatch | Check for update".

# Thanks: y0ush4@gmail.com
# Fixed By zj_w1nd@qq.com
'''
    Do not use shortcut(Ctrk + Alt + K) for debugging, otherwise debugging features will not be available.
'''
#########################################################################################################
is_debug = False
if is_debug:
    ''' 
        Install pydevd:

        1. sudo pip install pydevd

        or

        2. Install pycharm-debug.egg, Ensure to use pycharm pro
        https://www.jetbrains.com/help/pycharm/remote-debugging-with-product.html

        # import site
        # site.addsitedir("/usr/local/lib/python2.7/site-packages")
    '''
    try:
        import pydevd

        pydevd.settrace(host='localhost',
                        port=51234,
                        stdoutToServer=True,
                        stderrToServer=True
                        )
    except Exception as e:
        print(e)
#########################################################################################################

import os
import re
import json
from keystone import *
import idc
import idaapi
import ida_ida
import six

################################ IDA 6/7 Compatibility import ###########################################
if idaapi.IDA_SDK_VERSION >= 700:
    import ida_search
    import ida_bytes
    from idc import (
        get_operand_type,
        print_operand,
        get_item_end,
        get_item_head,
        get_item_size,
        get_sreg,
        has_value,
        get_full_flags,
        get_wide_byte,
        get_func_attr,
        set_func_end,
        warning,
        get_screen_ea
    )
    from idaapi import (
        get_bytes,
        plan_and_wait,
        is_mapped,
        Choose,
        jumpto
    )
else:
    from idc import (
        GetOpType as get_operand_type,
        GetOpnd as print_operand,
        ItemEnd as get_item_end,
        ItemHead as get_item_head,
        ItemSize as get_item_size,
        GetReg as get_sreg,
        hasValue as has_value,
        GetFlags as get_full_flags,
        Byte as get_wide_byte,
        GetFunctionAttr as get_func_attr,
        Jump as jumpto,
        Warning as warning,
        ScreenEA as get_screen_ea
    )
    from idaapi import (
        get_many_bytes as get_bytes,
        analyze_area as plan_and_wait,
        isEnabled as is_mapped,
        func_setend as set_func_end,
        Choose2 as Choose
    )
#########################################################################################################

# bleeding-edge version
# on a new release, this should be sync with VERSION_STABLE file
VERSION = "2.2"


MAX_INSTRUCTION_STRLEN = 256
MAX_ENCODING_LEN = 40
MAX_ADDRESS_LEN = 40
ENCODING_ERR_OUTPUT = "..."
KP_GITHUB_VERSION = "https://raw.githubusercontent.com/keystone-engine/keypatch/master/VERSION_STABLE"
KP_HOMEPAGE = "http://keystone-engine.org/keypatch"

X86_NOP = "\x90"

# Configuration file
KP_CFGFILE = os.path.join(idaapi.get_user_idadir(), "keypatch.cfg")

# save all the info on patching
patch_info = []

################################ Python2/3 Compatibility function #######################################
# Convert unicode/string/bytes to string
def to_string(s):
  # python3 bytes
  if six.PY3 and isinstance(s, bytes):
      return s.decode('latin-1')
  # python2 unicode
  elif six.PY2 and isinstance(s, six.text_type):
      return s.encode('utf-8')
  return str(s)

def to_hexstr(buf, sep=' '):
    # for python3 bytes
    if six.PY3 and isinstance(buf, bytes):
        return sep.join("{0:02x}".format(c) for c in buf).upper()
    return sep.join("{0:02x}".format(ord(c)) for c in buf).upper()
#########################################################################################################

################################ IDA 6/7 Compatibility function #########################################
def get_dtype(ea, op_idx):
    if idaapi.IDA_SDK_VERSION >= 700:
        insn = idaapi.insn_t()
        idaapi.decode_insn(insn, ea)
        dtype = insn.ops[op_idx].dtype
        dtyp_size = idaapi.get_dtype_size(dtype)
    else:
        dtype = idaapi.cmd.Operands[op_idx].dtyp
        dtyp_size = idaapi.get_dtyp_size(dtype)
    return dtype, dtyp_size

def set_comment(ea, comment):
    if idaapi.IDA_SDK_VERSION >= 700:
        idc.set_cmt(ea, comment, 0)
    else:
        idc.MakeComm(ea, comment)

def get_comment(ea):
    if idaapi.IDA_SDK_VERSION >= 700:
        return idc.get_cmt(ea, 0)
    return idc.Comment(ea)

def read_range_selection():
    if idaapi.IDA_SDK_VERSION >= 700:
        return idaapi.read_range_selection(None)
    return idaapi.read_selection()
#########################################################################################################

# return a normalized code, or None if input is invalid
def convert_hexstr(code):
    # normalize code
    code = code.lower()
    code = code.replace(' ', '')    # remove space
    code = code.replace('h', '')    # remove trailing 'h' in 90h
    code = code.replace('0x', '')   # remove 0x
    code = code.replace('\\x', '')  # remove \x
    code = code.replace(',', '')    # remove ,
    code = code.replace(';', '')    # remove ;
    code = code.replace('"', '')    # remove "
    code = code.replace("'", '')    # remove '
    code = code.replace("+", '')    # remove +

    # single-digit hexcode?
    if len(code) == 1 and ((code >= '0' and code <= '9') or (code >= 'a' and code <= 'f')):
        # stick 0 in front (so 'a' --> '0a')
        code = '0' + code

    # odd-length is invalid
    if len(code) % 2 != 0:
        return None

    try:
        hex_data = code.decode('hex')
        # we want a list of int
        return [ord(i) for i in hex_data]
    except:
        # invalid hex
        return None


# download a file from @url, then return (result, file-content)
# return (0, content) on success, or ({1|2}, None) on download failure
def url_download(url):
    # Import for python2
    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)


def get_name_value(_from, name):
    """
    Fixed: the return value truncated(32 bit) of get_name_value function that analyzed 64 bit binary file about ida64 for win.

    eg:
    type == idaapi.NT_BYTE
    (type, value) = idaapi.get_name_value(idc.BADADDR, "wcschr") # ida64 for win

    value = 0x14003d3f0L is correct  ida64 > 7.x for macOS
    value = 0x4003d3f0L is truncated ida64 >= 6.x for win, ida64 == 6.x for macOS

    :param _from: ea
    :param name: name string
    :return: tuple
    """
    name = to_string(name)
    (typ, value) = idaapi.get_name_value(_from, name)
    if typ == idaapi.NT_BYTE:  # typ is byte name (regular name)
        value = idaapi.get_name_ea(_from, name)
    return (typ, value)

## Main Keypatch class
class Keypatch_Asm:
    # supported architectures
    arch_lists = {
        "X86 16-bit": (KS_ARCH_X86, KS_MODE_16),                # X86 16-bit
        "X86 32-bit": (KS_ARCH_X86, KS_MODE_32),                # X86 32-bit
        "X86 64-bit": (KS_ARCH_X86, KS_MODE_64),                # X86 64-bit
        "ARM": (KS_ARCH_ARM, KS_MODE_ARM),                      # ARM
        "ARM Thumb": (KS_ARCH_ARM, KS_MODE_THUMB),              # ARM Thumb
        "ARM64 (ARMV8)": (KS_ARCH_ARM64, KS_MODE_LITTLE_ENDIAN),# ARM64
        "Hexagon": (KS_ARCH_HEXAGON, KS_MODE_BIG_ENDIAN),       # Hexagon
        "Mips32": (KS_ARCH_MIPS, KS_MODE_MIPS32),               # Mips32
        "Mips64": (KS_ARCH_MIPS, KS_MODE_MIPS64),               # Mips64
        "PowerPC 32": (KS_ARCH_PPC, KS_MODE_PPC32),             # PPC32
        "PowerPC 64": (KS_ARCH_PPC, KS_MODE_PPC64),             # PPC64
        "Sparc 32": (KS_ARCH_SPARC, KS_MODE_SPARC32),           # Sparc32
        "Sparc 64": (KS_ARCH_SPARC, KS_MODE_SPARC64),           # Sparc64
        "SystemZ": (KS_ARCH_SYSTEMZ, KS_MODE_BIG_ENDIAN),       # SystemZ
    }

    endian_lists = {
        "Little Endian": KS_MODE_LITTLE_ENDIAN,                 # little endian
        "Big Endian": KS_MODE_BIG_ENDIAN,                       # big endian
    }

    syntax_lists = {
        "Intel": KS_OPT_SYNTAX_INTEL,
        "Nasm": KS_OPT_SYNTAX_NASM,
        "AT&T": KS_OPT_SYNTAX_ATT
    }

    def __init__(self, arch=None, mode=None):
        # update current arch and mode
        self.update_hardware_mode()

        # override arch & mode if provided
        if arch is not None:
            self.arch = arch
        if mode is not None:
            self.mode = mode

        # IDA uses Intel syntax by default
        self.syntax = KS_OPT_SYNTAX_INTEL

    # return Keystone arch & mode (with endianess included)
    @staticmethod
    def get_hardware_mode():
        (arch, mode) = (None, None)

        if idaapi.IDA_SDK_VERSION >= 900:
            is_64bit = ida_ida.idainfo_is_64bit()
            is_32bit = ida_ida.idainfo_is_32bit()
            cpuname = ida_ida.inf_get_procname().lower()
            is_be = ida_ida.inf_is_be()
        else:
            # heuristically detect hardware setup
            info = idaapi.get_inf_structure()
            is_64bit = info.is_64bit()
            is_32bit = info.is_32bit()
        
            try:
                cpuname = info.procname.lower()
            except:
                cpuname = info.procName.lower()

            try:
                # since IDA7 beta 3 (170724) renamed inf.mf -> is_be()/set_be()
                is_be = idaapi.cvar.inf.is_be()
            except:
                # older IDA versions
                is_be = idaapi.cvar.inf.mf
                
        # print("Keypatch BIG_ENDIAN = %s" %is_be)
        
        if cpuname == "metapc":
            arch = KS_ARCH_X86
            if is_64bit:
                mode = KS_MODE_64
            elif is_32bit:
                mode = KS_MODE_32
            else:
                mode = KS_MODE_16
        elif cpuname.startswith("arm"):
            # ARM or ARM64
            if is_64bit:
                arch = KS_ARCH_ARM64
                if is_be:
                    mode = KS_MODE_BIG_ENDIAN
                else:
                    mode = KS_MODE_LITTLE_ENDIAN
            else:
                arch = KS_ARCH_ARM
                # either big-endian or little-endian
                if is_be:
                    mode = KS_MODE_ARM | KS_MODE_BIG_ENDIAN
                else:
                    mode = KS_MODE_ARM | KS_MODE_LITTLE_ENDIAN
        elif cpuname.startswith("sparc"):
            arch = KS_ARCH_SPARC
            if is_64bit:
                mode = KS_MODE_SPARC64
            else:
                mode = KS_MODE_SPARC32
            if is_be:
                mode |= KS_MODE_BIG_ENDIAN
            else:
                mode |= KS_MODE_LITTLE_ENDIAN
        elif cpuname.startswith("ppc"):
            arch = KS_ARCH_PPC
            if is_64bit:
                mode = KS_MODE_PPC64
            else:
                mode = KS_MODE_PPC32
            if cpuname == "ppc":
                # do not support Little Endian mode for PPC
                mode += KS_MODE_BIG_ENDIAN
        elif cpuname.startswith("mips"):
            arch = KS_ARCH_MIPS
            if is_64bit:
                mode = KS_MODE_MIPS64
            else:
                mode = KS_MODE_MIPS32
            if is_be:
                mode |= KS_MODE_BIG_ENDIAN
            else:
                mode |= KS_MODE_LITTLE_ENDIAN
        elif cpuname.startswith("systemz") or cpuname.startswith("s390x"):
            arch = KS_ARCH_SYSTEMZ
            mode = KS_MODE_BIG_ENDIAN

        return (arch, mode)

    def update_hardware_mode(self):
        (self.arch, self.mode) = self.get_hardware_mode()

    # normalize assembly code
    # remove comment at the end of assembly code
    @staticmethod
    def asm_normalize(text):
        text = ' '.join(text.split())
        if text.rfind(';') != -1:
            return text[:text.rfind(';')].strip()

        return text.strip()

    @staticmethod
    # check if input address is valid
    # return
    #       -1  invalid address at target binary
    #        0  type mismatch of input address
    #        1  valid address at target binary
    def check_address(address):
        try:
            if is_mapped(address):
                return 1
            else:
                return -1
        except:
            # invalid type
            return 0

    ### resolve IDA names from input asm code
    # todo: a better syntax parser for all archs
    def ida_resolve(self, assembly, address=idc.BADADDR):
        def _resolve(_op, ignore_kw=True):
            names = re.findall(r"[\$a-z0-9_:\.]+", _op, re.I)

            # try to resolve all names
            for name in names:
                # ignore known keywords
                if ignore_kw and name in ('byte', 'near', 'short', 'word', 'dword', 'ptr', 'offset'):
                    continue

                sym = name

                # split segment reg
                parts = name.partition(':')
                if parts[2] != '':
                    sym = parts[2]

                (typ, value) = get_name_value(address, sym)

                # skip if name doesn't exist or segment / segment registers
                if typ in (idaapi.NT_SEG, idaapi.NT_NONE):
                    continue

                _op = _op.replace(sym, '0x{0:X}'.format(value))

            return _op

        if self.check_address(address) == 0:
            print("Keypatch: WARNING: invalid input address {0}".format(address))
            return assembly

        # for now, we only support IDA name resolve for X86, ARM, ARM64, MIPS, PPC, SPARC
        if not (self.arch in (KS_ARCH_X86, KS_ARCH_ARM, KS_ARCH_ARM64, KS_ARCH_MIPS, KS_ARCH_PPC, KS_ARCH_SPARC)):
            return assembly

        _asm = assembly.partition(' ')
        mnem = _asm[0]
        opers = _asm[2].split(',')

        for idx, op in enumerate(opers):
            _op = list(op.partition('['))
            ignore_kw = True
            if _op[1] == '':
                _op[2] = _op[0]
                _op[0] = ''
            else:
                _op[0] = _resolve(_op[0], ignore_kw=True)
                ignore_kw = False

            _op[2] = _resolve(_op[2], ignore_kw=ignore_kw)

            opers[idx] = ''.join(_op)

        asm = "{0} {1}".format(mnem, ','.join(opers))
        return asm

    # return bytes of instruction or data
    # return None on failure
    def ida_get_item(self, address, hex_output=False):
        if self.check_address(address) != 1:
            # not a valid address
            return (None, 0)

        # return None if address is in the middle of instruction / data
        if address != get_item_head(address):
            return (None, 0)

        size = get_item_size(address)
        item = get_bytes(address, size)

        if item is None:
            return (None, 0)

        if hex_output:
            item = to_hexstr(item)

        return (item, size)

    @staticmethod
    def get_op_dtype_name(ea, op_idx):
        dtyp_lists = {
            idaapi.dt_byte: 'byte',     #  8 bit
            idaapi.dt_word: 'word',     #  16 bit
            idaapi.dt_dword: 'dword',   #  32 bit
            idaapi.dt_float: 'dword',   #  4 byte
            idaapi.dt_double: 'dword',  #  8 byte
            #idaapi.dt_tbyte = 5        #  variable size (ph.tbyte_size)
            #idaapi.dt_packreal = 6         #  packed real format for mc68040
            idaapi.dt_qword: 'qword',   #  64 bit
            idaapi.dt_byte16: 'xmmword',#  128 bit
            #idaapi.dt_code = 9         #  ptr to code (not used?)
            #idaapi.dt_void = 10        #  none
            #idaapi.dt_fword = 11       #  48 bit
            #idaapi.dt_bitfild = 12     #  bit field (mc680x0)
            #idaapi.dt_string = 13      #  pointer to asciiz string
            #idaapi.dt_unicode = 14     #  pointer to unicode string
            #idaapi.dt_3byte = 15       #  3-byte data
            #idaapi.dt_ldbl = 16        #  long double (which may be different from tbyte)
            idaapi.dt_byte32: 'ymmword',# 256 bit
        }

        dtype, dtyp_size = get_dtype(ea, op_idx)
        if dtype == idaapi.dt_tbyte and dtyp_size == 10:
          return 'xword'

        dtyp_name = dtyp_lists.get(dtype, None)

        return dtyp_name

    # return asm instructions from start to end
    def ida_get_disasm_range(self, start, end):
        codes = []
        while start < end:
            asm = self.asm_normalize(idc.GetDisasm(start))
            if asm == None:
                asm = ''
            codes.append(asm)
            start = start + get_item_size(start)

        return codes

    # get disasm from IDA
    # return '' on invalid address
    def ida_get_disasm(self, address, fixup=False):

        def GetMnem(asm):
            sp = asm.find(' ')
            if sp == -1:
                return asm
            return asm[:sp]

        if self.check_address(address) != 1:
            # not a valid address
            return ''

        # return if address is in the middle of instruction / data
        if address != get_item_head(address):
            return ''

        asm = self.asm_normalize(idc.GetDisasm(address))
        # for now, only support IDA syntax fixup for Intel CPU
        if not fixup or self.arch != KS_ARCH_X86:
            return asm

        # KS_ARCH_X86 mode
        # rebuild disasm code from IDA
        i = 0
        mnem = GetMnem(asm)
        if mnem == '' or mnem in ('rep', 'repne', 'repe'):
            return asm

        opers = []
        while get_operand_type(address, i) > 0 and i < 6:
            t = get_operand_type(address, i)
            o = print_operand(address, i)

            if t in (idc.o_mem, idc.o_displ):
                parts = list(o.partition(':'))
                if parts[2] == '':
                    parts[2] = parts[0]
                    parts[0] = ''

                if '[' not in parts[2]:
                    parts[2] = '[{0}]'.format(parts[2])

                o = ''.join(parts)

                if 'ptr ' not in o:
                    dtyp_name = self.get_op_dtype_name(address, i)
                    if dtyp_name != None:
                        o = "{0} ptr {1}".format(dtyp_name, o)

            opers.append(o)
            i += 1

        asm = mnem
        for o in opers:
            if o != '':
                asm = "{0} {1},".format(asm, o)

        asm = asm.strip(',')
        return asm

    # assemble code with Keystone
    # return (encoding, count), or (None, 0) on failure
    def assemble(self, assembly, address, arch=None, mode=None, syntax=None):

        # return assembly with arithmetic equation evaluated
        def eval_operand(assembly, start, stop, prefix=''):
            imm = assembly[start+1:stop]
            try:
                eval_imm = eval(imm)
                if eval_imm > 0x80000000:
                    eval_imm = 0xffffffff - eval_imm
                    eval_imm += 1
                    eval_imm = -eval_imm
                return assembly.replace(prefix + imm, prefix + hex(eval_imm))
            except:
                return assembly

        # IDA uses different syntax from Keystone
        # sometimes, we can convert code to be consumable by Keystone
        def fix_ida_syntax(assembly):

            # return True if this insn needs to be fixed
            def check_arm_arm64_insn(arch, mnem):
                if arch == KS_ARCH_ARM:
                    if mnem.startswith("ldr") or mnem.startswith("str"):
                        return True
                    return False
                elif arch == KS_ARCH_ARM64:
                    if mnem.startswith("ldr") or mnem.startswith("str"):
                        return True
                    return mnem in ("stp")
                return False

            # return True if this insn needs to be fixed
            def check_ppc_insn(mnem):
                return mnem in ("stw")

            # replace the right most string occurred
            def rreplace(s, old, new):
                li = s.rsplit(old, 1)
                return new.join(li)

            # convert some ARM pre-UAL assembly to UAL, so Keystone can handle it
            # example: streqb --> strbeq
            def fix_arm_ual(mnem, assembly):
                # TODO: this is not an exhaustive list yet
                if len(mnem) != 6:
                    return assembly

                if (mnem[-1] in ('s', 'b', 'h', 'd')):
                    #print(">> 222", mnem[3:5])
                    if mnem[3:5] in ("cc", "eq", "ne", "hs", "lo", "mi", "pl", "vs", "vc", "hi", "ls", "ge", "lt", "gt", "le", "al"):
                        return assembly.replace(mnem, mnem[:3] + mnem[-1] + mnem[3:5], 1)

                return assembly

            if self.arch != KS_ARCH_X86:
                assembly = assembly.lower()
            else:
                # Keystone does not support immediate 0bh, but only 0Bh
                assembly = assembly.upper()

            # however, 0X must be converted to 0x
            # Keystone should fix this limitation in the future
            assembly = assembly.replace("0X", " 0x")

            _asm = assembly.partition(' ')
            mnem = _asm[0]
            if mnem == '':
                return assembly

            # for PPC, Keystone does not accept registers with 'r' prefix,
            # but only the number behind. lets try to fix that here by
            # removing the prefix 'r'.
            if self.arch == KS_ARCH_PPC:
                for n in range(32):
                    r = " r%u," %n
                    if r in assembly:
                        assembly = assembly.replace(r, " %u," %n)
                for n in range(32):
                    r = "(r%u)" %n
                    if r in assembly:
                        assembly = assembly.replace(r, "(%u)" %n)
                for n in range(32):
                    r = ", r%u" %n
                    if assembly.endswith(r):
                        assembly = rreplace(assembly, r, ", %u" %n)

            if self.arch == KS_ARCH_X86:
                if mnem == "RETN":
                    # replace retn with ret
                    return assembly.replace('RETN', 'RET', 1)
                if 'OFFSET ' in assembly:
                    return assembly.replace('OFFSET ', ' ')
                if mnem in ('CALL', 'JMP') or mnem.startswith('LOOP'):
                    # remove 'NEAR PTR'
                    if ' NEAR PTR ' in assembly:
                        return assembly.replace(' NEAR PTR ', ' ')
                elif mnem[0] == 'J':
                    # JMP instruction
                    if ' SHORT ' in assembly:
                        # remove ' short '
                        return assembly.replace(' SHORT ', ' ')
            elif self.arch in (KS_ARCH_ARM, KS_ARCH_ARM64, KS_ARCH_PPC):
                # *** ARM
                # LDR     R1, [SP+rtld_fini],#4
                # STR     R2, [SP,#-4+rtld_fini]!
                # STR     R0, [SP,#fini]!
                # STR     R12, [SP,#4+var_8]!

                # *** ARM64
                # STP     X29, X30, [SP,#-0x10+var_150]!
                # STR     W0, [X29,#0x150+var_8]
                # LDR     X0, [X0,#(qword_4D6678 - 0x4D6660)]
                # TODO:
                # ADRP    X19, #interactive@PAGE

                # *** PPC
                # stw     r5, 0x120+var_108(r1)
                
                if self.arch == KS_ARCH_ARM and mode == KS_MODE_THUMB:
                    assembly = assembly.replace('movt.w', 'movt')

                if self.arch == KS_ARCH_ARM:
                    #print(">> before UAL fix: ", assembly)
                    assembly = fix_arm_ual(mnem, assembly)
                    #print(">> after UAL fix: ", assembly)

                if check_arm_arm64_insn(self.arch, mnem) or (("[" in assembly) and ("]" in assembly)):
                    bang = assembly.find('#')
                    bracket = assembly.find(']')
                    if bang != -1 and bracket != -1 and bang < bracket:
                        return eval_operand(assembly, bang, bracket, '#')
                    elif '+0x0]' in assembly:
                        return assembly.replace('+0x0]', ']')
                elif check_ppc_insn(mnem):
                    start = assembly.find(', ')
                    stop = assembly.find('(')
                    if start != -1 and stop != -1 and start < stop:
                        return eval_operand(assembly, start, stop)
            return assembly

        def is_thumb(address):
            return get_sreg(address, 'T') == 1

        if self.check_address(address) == 0:
            return (None, 0)

        # use default syntax, arch and mode if not provided
        if syntax is None:
            syntax = self.syntax
        if arch is None:
            arch = self.arch
        if mode is None:
            mode = self.mode

        if arch == KS_ARCH_ARM and is_thumb(address):
            mode = KS_MODE_THUMB

        try:
            ks = Ks(arch, mode)
            if arch == KS_ARCH_X86:
                ks.syntax = syntax
            encoding, count = ks.asm(fix_ida_syntax(assembly), address)
        except KsError as e:
            # keep the below code for debugging
            #print("Keypatch Error: {0}".format(e))
            #print("Original asm: {0}".format(assembly))
            #print("Fixed up asm: {0}".format(fix_ida_syntax(assembly)))
            encoding, count = None, 0

        return (encoding, count)


    # patch at address, return the number of written bytes & original data
    # this process can fail in some cases
    @staticmethod
    def patch_raw(address, patch_data, size):
        ea = address
        orig_data = ''

        while ea < (address + size):

            if not has_value(get_full_flags(ea)):
                print("Keypatch: FAILED to read data at 0x{0:X}".format(ea))
                break

            orig_byte = get_wide_byte(ea)
            orig_data += chr(orig_byte)
            patch_byte = ord(patch_data[ea - address])

            if patch_byte != orig_byte:
                # patch one byte
                if idaapi.patch_byte(ea, patch_byte) != 1:
                    print("Keypatch: FAILED to patch byte at 0x{0:X} [0x{1:X}]".format(ea, patch_byte))
                    break
            ea += 1
        return (ea - address, orig_data)

    # patch at address, return the number of written bytes & original data
    # on patch failure, we revert to the original code, then return (None, None)
    def patch(self, address, patch_data, size):
        # save original function end to fix IDA re-analyze issue after patching
        orig_func_end = get_func_attr(address, idc.FUNCATTR_END)

        (patched_len, orig_data) = self.patch_raw(address, patch_data, size)

        if size != patched_len:
            # patch failure
            if patched_len > 0:
                # revert the changes
                (rlen, _) = self.patch_raw(address, orig_data, patched_len)
                if rlen == patched_len:
                    print("Keypatch: successfully reverted changes of {0:d} byte(s) at 0x{1:X} [{2}]".format(
                                        patched_len, address, to_hexstr(orig_data)))
                else:
                    print("Keypatch: FAILED to revert changes of {0:d} byte(s) at 0x{1:X} [{2}]".format(
                                        patched_len, address, to_hexstr(orig_data)))

            return (None, None)

        # ask IDA to re-analyze the patched area
        if orig_func_end == idc.BADADDR:
            # only analyze patched bytes, otherwise it would take a lot of time to re-analyze the whole binary
            plan_and_wait(address, address + patched_len + 1)
        else:
            plan_and_wait(address, orig_func_end)

            # try to fix IDA function re-analyze issue after patching
            set_func_end(address, orig_func_end)

        return (patched_len, orig_data)

    # return number of bytes patched
    # return
    #    0  Invalid assembly
    #   -1  PatchByte failure
    #   -2  Can't read original data
    #   -3  Invalid address
    def patch_code(self, address, assembly, syntax, padding, save_origcode, orig_asm=None, patch_data=None, patch_comment=None, undo=False):
        global patch_info

        if self.check_address(address) != 1:
            # not a valid address
            return -3

        orig_comment = get_comment(address)
        if orig_comment is None:
            orig_comment = ''

        nop_comment = ""
        padding_len = 0
        if not undo:
            # we are patching via Patcher
            (orig_encoding, orig_len) = self.ida_get_item(address)
            if (orig_encoding, orig_len) == (None, 0):
                return -2

            (encoding, count) = self.assemble(assembly, address, syntax=syntax)
            if encoding is None:
                return 0

            patch_len = len(encoding)
            patch_data = ''.join(chr(c) for c in encoding)

            if patch_data == orig_encoding:
                #print("Keypatch: no need to patch, same encoding data [{0}] at 0x{1:X}".format(to_hexstr(orig_encoding), address))
                return orig_len

            # for now, only support NOP padding on Intel CPU
            if padding and self.arch == KS_ARCH_X86:
                if patch_len < orig_len:
                    padding_len = orig_len - patch_len
                    patch_len = orig_len
                    patch_data = patch_data.ljust(patch_len, X86_NOP)
                elif patch_len > orig_len:
                    patch_end = address + patch_len - 1
                    ins_end = get_item_end(patch_end)
                    padding_len = ins_end - patch_end - 1

                    if padding_len > 0:
                        patch_len = ins_end - address
                        patch_data = patch_data.ljust(patch_len, X86_NOP)

                if padding_len > 0:
                    nop_comment = "\nKeypatch padded NOP to next boundary: {0} bytes".format(padding_len)

            orig_asm = self.ida_get_disasm_range(address, address + patch_len)
        else:
            # we are reverting the change via "Undo" menu
            patch_len = len(patch_data)

        (plen, p_orig_data) = self.patch(address, patch_data, patch_len)
        if plen is None:
            # failed to patch
            return -1

        if not undo: # we are patching
            new_patch_comment = None
            if save_origcode is True:
                # append original instruction to comments
                if orig_comment == '':
                    new_patch_comment = "Keypatch modified this from:\n  {0}{1}".format('\n  '.join(orig_asm), nop_comment)
                else:
                    new_patch_comment = "\nKeypatch modified this from:\n  {0}{1}".format('\n  '.join(orig_asm), nop_comment)

                new_comment = "{0}{1}".format(orig_comment, new_patch_comment)
                set_comment(address, new_comment)

            if padding_len == 0:
                print("Keypatch: successfully patched {0:d} byte(s) at 0x{1:X} from [{2}] to [{3}]".format(plen,
                                        address, to_hexstr(p_orig_data), to_hexstr(patch_data)))
            else:
                print("Keypatch: successfully patched {0:d} byte(s) at 0x{1:X} from [{2}] to [{3}], with {4} byte(s) NOP padded".format(plen,
                                        address, to_hexstr(p_orig_data), to_hexstr(patch_data), padding_len))
            # save this patching for future "undo"
            patch_info.append((address, assembly, p_orig_data, new_patch_comment))
        else:   # we are reverting
            if patch_comment:
                # clean previous IDA comment by replacing it with ''
                new_comment = orig_comment.replace(patch_comment, '')
                set_comment(address, new_comment)

            print("Keypatch: successfully reverted {0:d} byte(s) at 0x{1:X} from [{2}] to [{3}]".format(plen,
                                        address, to_hexstr(p_orig_data), to_hexstr(patch_data)))

        return plen

    # fill a range of code [addr_begin, addr_end].
    # return the length of patched area
    # on failure, return 0 = wrong input, -1 = failed to patch
    def fill_code(self, addr_begin, addr_end, assembly, syntax, padding, save_origcode, orig_asm=None):
        # treat input as assembly code first
        (encoding, _) =  self.assemble(assembly, addr_begin, syntax=syntax)

        if encoding is None:
            # input might be a hexcode string. try to convert it to raw bytes
            encoding = convert_hexstr(assembly)

        if encoding is None:
            # invalid input: this is neither assembly nor hexcode string
            return 0

        # save original assembly code before overwritting them
        orig_asm = self.ida_get_disasm_range(addr_begin, addr_end)

        # save original comment at addr_begin
        # TODO: save comments in this range, but how to interleave them?
        orig_comment = get_comment(addr_begin)
        if orig_comment is None:
            orig_comment = ''

        patch_data = ""
        assembly_new = []
        size = addr_end - addr_begin
        # calculate filling data
        encode_chr = ''.join(chr(c) for c in encoding)
        while True:
            if len(patch_data) + len(encode_chr) <= size:
                patch_data = patch_data + encode_chr
                assembly_new += [assembly.strip()]
            else:
                break

        # for now, only support NOP padding on Intel CPU
        if padding and self.arch == KS_ARCH_X86:
            for i in range(size -len(patch_data)):
                assembly_new += ["nop"]
            patch_data = patch_data.ljust(size, X86_NOP)

        (plen, p_orig_data) = self.patch(addr_begin, patch_data, len(patch_data))
        if plen is None:
            # failed to patch
            return -1

        new_patch_comment = ''
        # append original instruction to comments
        if save_origcode is True:
            if orig_comment == '':
                new_patch_comment = "Keypatch filled range [0x{0:X}:0x{1:X}] ({2} bytes), replaced:\n  {3}".format(addr_begin, addr_end - 1, addr_end - addr_begin, '\n  '.join(orig_asm))
            else:
                new_patch_comment = "\nKeypatch filled range [0x{0:X}:0x{1:X}] ({2} bytes), replaced:\n  {3}".format(addr_begin, addr_end - 1, addr_end - addr_begin, '\n  '.join(orig_asm))

            new_comment = "{0}{1}".format(orig_comment, new_patch_comment)
            set_comment(addr_begin, new_comment)

        print("Keypatch: successfully filled range [0x{0:X}:0x{1:X}] ({2} bytes) with \"{3}\", replaced \"{4}\"".format(
                    addr_begin, addr_end - 1, addr_end - addr_begin, assembly, '; '.join(orig_asm)))

        # save this modification for future "undo"
        patch_info.append((addr_begin, '\n  '.join(assembly_new), p_orig_data, new_patch_comment))

        return plen


    ### Form helper functions
    @staticmethod
    def dict_to_ordered_list(dictionary):
        l = sorted(list(dictionary.items()), key=lambda t: t[0], reverse=False)
        keys = [i[0] for i in l]
        values = [i[1] for i in l]

        return (keys, values)

    def get_value_by_idx(self, dictionary, idx, default=None):
        (keys, values) = self.dict_to_ordered_list(dictionary)

        try:
            val = values[idx]
        except IndexError:
            val = default

        return val

    def find_idx_by_value(self, dictionary, value, default=None):
        (keys, values) = self.dict_to_ordered_list(dictionary)

        try:
            idx = values.index(value)
        except:
            idx = default

        return idx

    def get_arch_by_idx(self, idx):
        return self.get_value_by_idx(self.arch_lists, idx)

    def find_arch_idx(self, arch, mode):
        return self.find_idx_by_value(self.arch_lists, (arch, mode))

    def get_syntax_by_idx(self, idx):
        return self.get_value_by_idx(self.syntax_lists, idx, self.syntax)

    def find_syntax_idx(self, syntax):
        return self.find_idx_by_value(self.syntax_lists, syntax)
    ### /Form helper functions


# Common ancestor form to be derived by Patcher, FillRange & Search
class Keypatch_Form(idaapi.Form):
    # prepare for form initializing
    def setup(self, kp_asm, address, assembly=None):
        self.kp_asm = kp_asm
        self.address = address

        # update ordered list of arch and syntax
        self.syntax_keys = self.kp_asm.dict_to_ordered_list(self.kp_asm.syntax_lists)[0]
        self.arch_keys = self.kp_asm.dict_to_ordered_list(self.kp_asm.arch_lists)[0]

        # update current arch & mode
        self.kp_asm.update_hardware_mode()

        # find right value for c_arch & c_endian controls
        mode = self.kp_asm.mode
        self.endian_id = 0   # little endian
        if self.kp_asm.mode & KS_MODE_BIG_ENDIAN:
            self.endian_id = 1   # big endian
            mode = self.kp_asm.mode - KS_MODE_BIG_ENDIAN

        self.arch_id = self.kp_asm.find_arch_idx(self.kp_asm.arch, mode)

        self.syntax_id = 0  # to make non-X86 arch happy
        if self.kp_asm.arch == KS_ARCH_X86:
            self.syntax_id = self.kp_asm.find_syntax_idx(self.kp_asm.syntax)

        # get original instruction and bytes
        self.orig_asm = kp_asm.ida_get_disasm(address)
        (self.orig_encoding, self.orig_len) = kp_asm.ida_get_item(address, hex_output=True)
        if self.orig_encoding is None:
            self.orig_encoding = ''

        if assembly is None:
            self.asm = self.kp_asm.ida_get_disasm(self.address, fixup=True)
        else:
            self.asm = assembly

    # update Encoding control
    # return True on success, False on failure
    def _update_encoding(self, arch, mode):
        try:
            syntax = None
            if arch == KS_ARCH_X86:
                syntax_id = self.GetControlValue(self.c_syntax)
                syntax = self.kp_asm.get_syntax_by_idx(syntax_id)

            address = self.GetControlValue(self.c_addr)
            try:
                is_mapped(address)
            except:
                # invalid address value
                address = 0

            assembly = self.GetControlValue(self.c_assembly)
            raw_assembly = self.kp_asm.ida_resolve(assembly, address)
            self.SetControlValue(self.c_raw_assembly, raw_assembly)

            (encoding, count) =  self.kp_asm.assemble(raw_assembly, address, arch=arch,
                                                    mode=mode, syntax=syntax)

            if encoding is None:
                self.SetControlValue(self.c_encoding, ENCODING_ERR_OUTPUT)
                self.SetControlValue(self.c_encoding_len, 0)
                return False
            else:
                text = ""
                for byte in encoding:
                    text += "%02X " % byte
                text.strip()
                if text == "":
                    # error?
                    self.SetControlValue(self.c_encoding, ENCODING_ERR_OUTPUT)
                    return False
                else:
                    self.SetControlValue(self.c_encoding, text.strip())
                    self.SetControlValue(self.c_encoding_len, len(encoding))
                    return True
        except Exception as e:
            print(str(e))
            import traceback
            traceback.print_exc()
            self.SetControlValue(self.c_encoding, ENCODING_ERR_OUTPUT)
            return False

    # callback to be executed when any form control changed
    def OnFormChange(self, fid):
        return 1

    # update Patcher & Fillrange controls
    def update_patchform(self, fid):
        self.EnableField(self.c_endian, False)
        self.EnableField(self.c_addr, False)

        (arch, mode) = (self.kp_asm.arch, self.kp_asm.mode)
        # assembly is focused
        self.SetFocusedField(self.c_assembly)

        if arch == KS_ARCH_X86:
            # do not show Endian control
            self.ShowField(self.c_endian, False)
            # allow to choose Syntax
            self.ShowField(self.c_syntax, True)
            self.ShowField(self.c_opt_padding, True)
        else:   # do not show Syntax control for non-X86 mode
            self.ShowField(self.c_syntax, False)
            # for now, we do not support padding for non-X86 archs
            self.ShowField(self.c_opt_padding, False)
            #self.EnableField(self.c_opt_padding, False)

        # update other controls & Encoding with live assembling
        self.update_controls(arch, mode)

        return 1

    # update some controls - including Encoding control
    def update_controls(self, arch, mode):
        # Fixup & Encoding-len are read-only controls
        self.EnableField(self.c_raw_assembly, False)
        self.EnableField(self.c_encoding_len, False)

        # Encoding is enable to allow user to select & copy
        self.EnableField(self.c_encoding, True)

        if self.GetControlValue(self.c_endian) == 1:
            endian = KS_MODE_BIG_ENDIAN
        else:
            endian = KS_MODE_LITTLE_ENDIAN

        # update encoding with live assembling
        self._update_encoding(arch, mode | endian)

        return 1

    # get Patcher/FillRange options
    def get_opts(self, name=None):
        names = self.c_opt_chk.children_names
        val = self.c_opt_chk.value
        opts = {}
        for i in range(len(names)):
            opts[names[i]] = val & (2**i)

        if name != None:
            opts[name] = val

        return opts


# Fill Range form
class Keypatch_FillRange(Keypatch_Form):
    def __init__(self, kp_asm, addr_begin, addr_end, assembly=None, opts=None):
        self.setup(kp_asm, addr_begin, assembly)
        self.addr_end = addr_end

        # create FillRange form
        super(Keypatch_FillRange, self).__init__(
            r"""STARTITEM {id:c_assembly}
BUTTON YES* Patch
KEYPATCH:: Fill Range

            {FormChangeCb}
            <Endian     :{c_endian}>
            <~S~yntax     :{c_syntax}>
            <Start      :{c_addr}>
            <End        :{c_addr_end}>
            <Size       :{c_size}>
            <~A~ssembly   :{c_assembly}>
             <##-   Fixup :{c_raw_assembly}>
             <##-   Encode:{c_encoding}>
             <##-   Size  :{c_encoding_len}>
            <~N~OPs padding until next instruction boundary:{c_opt_padding}>
            <Save ~o~riginal instructions in IDA comment:{c_opt_comment}>{c_opt_chk}>
            """, {
            'c_endian': self.DropdownListControl(
                          items = list(self.kp_asm.endian_lists.keys()),
                          readonly = True,
                          selval = self.endian_id),
            'c_addr': self.NumericInput(value=addr_begin, swidth=MAX_ADDRESS_LEN, tp=self.FT_ADDR),
            'c_addr_end': self.NumericInput(value=addr_end - 1, swidth=MAX_ADDRESS_LEN, tp=self.FT_ADDR),
            'c_assembly': self.StringInput(value=self.asm[:MAX_INSTRUCTION_STRLEN], width=MAX_INSTRUCTION_STRLEN),
            'c_size': self.NumericInput(value=addr_end - addr_begin, swidth=8, tp=self.FT_DEC),
            'c_raw_assembly': self.StringInput(value='', width=MAX_INSTRUCTION_STRLEN),
            'c_encoding': self.StringInput(value='', width=MAX_ENCODING_LEN),
            'c_encoding_len': self.NumericInput(value=0, swidth=8, tp=self.FT_DEC),
            'c_syntax': self.DropdownListControl(
                          items = self.syntax_keys,
                          readonly = True,
                          selval = self.syntax_id),
            'c_opt_chk':idaapi.Form.ChkGroupControl(('c_opt_padding', 'c_opt_comment', ''), value=opts['c_opt_chk']),
            'FormChangeCb': self.FormChangeCb(self.OnFormChange),
            })

        self.Compile()

    # callback to be executed when any form control changed
    def OnFormChange(self, fid):
        # make some controls read-only in FillRange mode
        self.EnableField(self.c_size, False)
        self.EnableField(self.c_addr_end, False)

        return self.update_patchform(fid)


# Patcher form
class Keypatch_Patcher(Keypatch_Form):
    def __init__(self, kp_asm, address, assembly=None, opts=None):
        self.setup(kp_asm, address, assembly)

        # create Patcher form
        super(Keypatch_Patcher, self).__init__(
            r"""STARTITEM {id:c_assembly}
BUTTON YES* Patch
KEYPATCH:: Patcher

            {FormChangeCb}
            <Endian     :{c_endian}>
            <~S~yntax     :{c_syntax}>
            <Address    :{c_addr}>
            <Original   :{c_orig_assembly}>
             <##-   Encode:{c_orig_encoding}>
             <##-   Size  :{c_orig_len}>
            <~A~ssembly   :{c_assembly}>
             <##-   Fixup :{c_raw_assembly}>
             <##-   Encode:{c_encoding}>
             <##-   Size  :{c_encoding_len}>
            <~N~OPs padding until next instruction boundary:{c_opt_padding}>
            <Save ~o~riginal instructions in IDA comment:{c_opt_comment}>{c_opt_chk}>
            """, {
            'c_endian': self.DropdownListControl(
                          items = list(self.kp_asm.endian_lists.keys()),
                          readonly = True,
                          selval = self.endian_id),
            'c_addr': self.NumericInput(value=address, swidth=MAX_ADDRESS_LEN, tp=self.FT_ADDR),
            'c_assembly': self.StringInput(value=self.asm[:MAX_INSTRUCTION_STRLEN], width=MAX_INSTRUCTION_STRLEN),
            'c_orig_assembly': self.StringInput(value=self.orig_asm[:MAX_INSTRUCTION_STRLEN], width=MAX_INSTRUCTION_STRLEN),
            'c_orig_encoding': self.StringInput(value=self.orig_encoding[:MAX_ENCODING_LEN], width=MAX_ENCODING_LEN),
            'c_orig_len': self.NumericInput(value=self.orig_len, swidth=8, tp=self.FT_DEC),
            'c_raw_assembly': self.StringInput(value='', width=MAX_INSTRUCTION_STRLEN),
            'c_encoding': self.StringInput(value='', width=MAX_ENCODING_LEN),
            'c_encoding_len': self.NumericInput(value=0, swidth=8, tp=self.FT_DEC),
            'c_syntax': self.DropdownListControl(
                          items = self.syntax_keys,
                          readonly = True,
                          selval = self.syntax_id),
            'c_opt_chk':self.ChkGroupControl(('c_opt_padding', 'c_opt_comment', ''), value=opts['c_opt_chk']),
            'FormChangeCb': self.FormChangeCb(self.OnFormChange),
            })

        self.Compile()

    # callback to be executed when any form control changed
    def OnFormChange(self, fid):
        # make some fields read-only in Patch mode
        self.EnableField(self.c_orig_assembly, False)
        self.EnableField(self.c_orig_encoding, False)
        self.EnableField(self.c_orig_len, False)

        return self.update_patchform(fid)


# Search position chooser
class SearchResultChooser(Choose):
    def __init__(self, title, items, flags=0, width=None, height=None, embedded=False, modal=False):
        super(SearchResultChooser, self).__init__(
            title,
            [["Address", Choose.CHCOL_HEX|40]],
            flags = flags,
            width = width,
            height = height,
            embedded = embedded)
        self.n = 0
        self.items = items
        self.selcount = 0
        self.modal = modal

    def OnClose(self):
        return

    def OnSelectLine(self, n):
        self.selcount += 1
        jumpto(self.items[n][0])

    def OnGetLine(self, n):
        res = self.items[n]
        res = [idc.atoa(res[0])]
        return res

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

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


# Search form
class Keypatch_Search(Keypatch_Form):
    def __init__(self, kp_asm, address, assembly=None):
        self.setup(kp_asm, address, assembly)

        # create Search form
        super(Keypatch_Search, self).__init__(
            r"""STARTITEM {id:c_assembly}
BUTTON YES* Search
KEYPATCH:: Search

            {FormChangeCb}
            <A~r~ch       :{c_arch}>
            <E~n~dian     :{c_endian}>
            <~S~yntax     :{c_syntax}>
            <A~d~dress    :{c_addr}>
            <~A~ssembly   :{c_assembly}>
             <##-   Fixup :{c_raw_assembly}>
             <##-   Encode:{c_encoding}>
             <##-   Size  :{c_encoding_len}>
            """, {
            'c_addr': self.NumericInput(value=address, swidth=MAX_ADDRESS_LEN, tp=self.FT_ADDR),
            'c_assembly': self.StringInput(value=self.asm[:MAX_INSTRUCTION_STRLEN], width=MAX_INSTRUCTION_STRLEN),
            'c_raw_assembly': self.StringInput(value='', width=MAX_INSTRUCTION_STRLEN),
            'c_encoding': self.StringInput(value='', width=MAX_ENCODING_LEN),
            'c_encoding_len': self.NumericInput(value=0, swidth=8, tp=self.FT_DEC),
            'c_arch': self.DropdownListControl(
                          items = self.arch_keys,
                          readonly = True,
                          selval = self.arch_id,
                          width = 32),
            'c_endian': self.DropdownListControl(
                          items = list(self.kp_asm.endian_lists.keys()),
                          readonly = True,
                          selval = self.endian_id),
            'c_syntax': self.DropdownListControl(
                          items = self.syntax_keys,
                          readonly = True,
                          selval = self.syntax_id),
            'FormChangeCb': self.FormChangeCb(self.OnFormChange),
            })

        self.Compile()

    ################################ IDA 6/7 Compatibility method #######################################
    def find_binary(self, ea):
        if idaapi.IDA_SDK_VERSION >= 900:
        # for ida 9.0sp1 we use this new api:
            return ida_bytes.find_bytes(self.GetControlValue(self.c_encoding), ea)
        if idaapi.IDA_SDK_VERSION >= 700:
            return ida_search.find_binary(ea, idaapi.cvar.inf.max_ea, self.GetControlValue(self.c_encoding), 16, idc.SEARCH_DOWN)
        return idc.FindBinary(ea, idc.SEARCH_DOWN, self.GetControlValue(self.c_encoding))
    #####################################################################################################

    # callback to be executed when any form control changed
    def OnFormChange(self, fid):
        # handle the search button
        if fid == -2:
            address = 0
            addresses = []
            while address != idc.BADADDR:
                address = self.find_binary(address)
                if address == idc.BADADDR:
                    break
                addresses.append([address])
                address = address + 1
            c = SearchResultChooser("Searching for [{0}]".format(self.GetControlValue(self.c_raw_assembly)), addresses)
            c.show()
            return 1

        # only Search mode allows to select arch+mode
        arch_id = self.GetControlValue(self.c_arch)
        (arch, mode) = self.kp_asm.get_arch_by_idx(arch_id)

        # assembly is focused
        self.SetFocusedField(self.c_assembly)

        if arch == KS_ARCH_X86:
            # enable Syntax and disable Endian for x86
            self.ShowField(self.c_syntax, True)
            self.EnableField(self.c_syntax, True)
            self.syntax_id = self.GetControlValue(self.c_syntax)
            self.EnableField(self.c_endian, False)
            # set Endian index properly
            self.SetControlValue(self.c_endian, 0)
        elif arch in (KS_ARCH_ARM64, KS_ARCH_HEXAGON, KS_ARCH_SYSTEMZ):
            # no Syntax & Endian option for these archs
            self.ShowField(self.c_syntax, False)
            self.EnableField(self.c_syntax, False)
            self.EnableField(self.c_endian, False)
            # set Endian index properly
            self.SetControlValue(self.c_endian, (mode & KS_MODE_BIG_ENDIAN != 0))
        elif (arch == KS_ARCH_PPC) and (mode & KS_MODE_PPC32 != 0):
            # no Syntax & Endian option for these archs
            self.ShowField(self.c_syntax, False)
            self.EnableField(self.c_syntax, False)
            self.EnableField(self.c_endian, False)
            # set Endian index properly
            self.SetControlValue(self.c_endian, (mode & KS_MODE_BIG_ENDIAN != 0))
        else:
            # no Syntax & Endian option
            self.ShowField(self.c_syntax, False)
            self.EnableField(self.c_syntax, False)
            self.EnableField(self.c_endian, True)

        if self.GetControlValue(self.c_endian) == 1:
            endian = KS_MODE_BIG_ENDIAN
        else:
            endian = KS_MODE_LITTLE_ENDIAN

        # update other controls & Encoding with live assembling
        self.update_controls(arch, mode)

        return 1


# About form
class About_Form(idaapi.Form):
    def __init__(self, version):
        # create About form
        super(About_Form, self).__init__(
            r"""STARTITEM 0
BUTTON YES* Open Keypatch Website
KEYPATCH:: About

            {FormChangeCb}
            Keypatch IDA plugin v%s, using Keystone Engine v%s.
            (c) Nguyen Anh Quynh + Thanh Nguyen, 2018.

            Keypatch is released under the GPL v2.
            Find more info at http://www.keystone-engine.org/keypatch
            """ %(version, keystone.__version__), {
            '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(KP_HOMEPAGE, new = 2)

        return 1


# Check-for-update form
class Update_Form(idaapi.Form):
    def __init__(self, version, message):
        # create Update form
        super(Update_Form, self).__init__(
            r"""STARTITEM 0
BUTTON YES* Open Keypatch Website
KEYPATCH:: Check for update

            {FormChangeCb}
            Your Keypatch 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(KP_HOMEPAGE, new = 2)

        return 1


try:
    # adapted from pull request #7 by @quangnh89
    class Kp_Menu_Context(idaapi.action_handler_t):

        @classmethod
        def get_name(self):
            return self.__name__

        @classmethod
        def get_label(self):
            return self.label

        @classmethod
        def register(self, plugin, label):
            self.plugin = plugin
            self.label = label
            instance = self()
            return idaapi.register_action(idaapi.action_desc_t(
                self.get_name(),  # Name. Acts as an ID. Must be unique.
                instance.get_label(),  # Label. That's what users see.
                instance  # Handler. Called when activated, and for updating
            ))

        @classmethod
        def unregister(self):
            """Unregister the action.
            After unregistering the class cannot be used.
            """
            idaapi.unregister_action(self.get_name())

        @classmethod
        def activate(self, ctx):
            # dummy method
            return 1

        @classmethod
        def update(self, ctx):
            try:
                if idaapi.IDA_SDK_VERSION >= 900:
                    # Since IDA 9.0, form_type is deprecated, should use widget_type
                    if ctx.widget_type == idaapi.BWN_DISASM:
                        return idaapi.AST_ENABLE_FOR_FORM
                    else:
                        return idaapi.AST_DISABLE_FOR_FORM
                else:
                    if ctx.form_type == idaapi.BWN_DISASM:
                        return idaapi.AST_ENABLE_FOR_FORM
                    else:
                        return idaapi.AST_DISABLE_FOR_FORM
            except Exception as e:
                # Add exception for main menu on >= IDA 7.0
                return idaapi.AST_ENABLE_ALWAYS
            
    # context menu for Patcher
    class Kp_MC_Patcher(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.patcher()
            return 1

    # context menu for Fill Range
    class Kp_MC_Fill_Range(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.fill_range()
            return 1

    # context menu for Undo
    class Kp_MC_Undo(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.undo()
            return 1

    # context menu for Search
    class Kp_MC_Search(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.search()
            return 1

    # context menu for Check Update
    class Kp_MC_Updater(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.updater()
            return 1

    # context menu for About
    class Kp_MC_About(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.about()
            return 1

except:
    pass

# hooks for popup menu
class Hooks(idaapi.UI_Hooks):
    if idaapi.IDA_SDK_VERSION >= 700:
        # IDA >= 700 right click widget popup
        def finish_populating_widget_popup(self, form, popup):
            if idaapi.get_widget_type(form) == idaapi.BWN_DISASM:
                try:
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Patcher.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Fill_Range.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Undo.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, "-", 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Search.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, "-", 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Updater.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_About.get_name(), 'Keypatch/')
                except:
                    pass
    else:
        # IDA < 700 right click popup
        def finish_populating_tform_popup(self, form, popup):
            # We'll add our action to all "IDA View-*"s.
            # If we wanted to add it only to "IDA View-A", we could
            # also discriminate on the widget's title:
            #
            #  if idaapi.get_tform_title(form) == "IDA View-A":
            #      ...
            #
            if idaapi.get_tform_type(form) == idaapi.BWN_DISASM:
                try:
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Patcher.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Fill_Range.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Undo.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, "-", 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Search.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, "-", 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Updater.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_About.get_name(), 'Keypatch/')
                except:
                    pass


# check if we already initialized Keypatch
kp_initialized = False

#--------------------------------------------------------------------------
# Plugin
#--------------------------------------------------------------------------
class Keypatch_Plugin_t(idaapi.plugin_t):
    comment = "Keypatch plugin for IDA Pro (using Keystone framework)"
    help = "Find more information on Keypatch at http://keystone-engine.org/keypatch"
    wanted_name = "Keypatch Patcher"
    wanted_hotkey = "Ctrl-Alt-K"
    flags = idaapi.PLUGIN_KEEP

    def load_configuration(self):
        # default
        self.opts = {}

        # load configuration from file
        try:
            f = open(KP_CFGFILE, "rt")
            self.opts = json.load(f)
            f.close()
        except IOError:
            print("Keypatch: FAILED to load config file. Use default setup now.")
        except Exception as e:
            print("Keypatch: FAILED to load config file, with exception: {0}".format(str(e)))

        # use default values if not defined in config file
        if 'c_opt_padding' not in self.opts:
            self.opts['c_opt_padding'] = 1

        if 'c_opt_comment' not in self.opts:
            self.opts['c_opt_comment'] = 2

        self.opts['c_opt_chk'] = self.opts['c_opt_padding'] | self.opts['c_opt_comment']

    def init(self):
        global kp_initialized

        # register popup menu handlers
        try:
            Kp_MC_Patcher.register(self, "Patcher    (Ctrl-Alt-K)")
            Kp_MC_Fill_Range.register(self, "Fill Range")
            Kp_MC_Undo.register(self, "Undo last patching")
            Kp_MC_Search.register(self, "Search")
            Kp_MC_Updater.register(self, "Check for update")
            Kp_MC_About.register(self, "About")
        except:
            pass

        # setup popup menu
        self.hooks = Hooks()
        self.hooks.hook()

        self.opts = None
        if kp_initialized == False:
            kp_initialized = True

            if idaapi.IDA_SDK_VERSION >= 700:
                # Add menu IDA >= 7.0
                idaapi.attach_action_to_menu("Edit/Keypatch/Patcher", Kp_MC_Patcher.get_name(), idaapi.SETMENU_APP)
                idaapi.attach_action_to_menu("Edit/Keypatch/About", Kp_MC_About.get_name(), idaapi.SETMENU_APP)
                idaapi.attach_action_to_menu("Edit/Keypatch/Check for update", Kp_MC_Updater.get_name(), idaapi.SETMENU_APP)
                idaapi.attach_action_to_menu("Edit/Keypatch/Search", Kp_MC_Search.get_name(), idaapi.SETMENU_APP)
                idaapi.attach_action_to_menu("Edit/Keypatch/Undo last patching", Kp_MC_Undo.get_name(), idaapi.SETMENU_APP)
                idaapi.attach_action_to_menu("Edit/Keypatch/Fill Range", Kp_MC_Fill_Range.get_name(), idaapi.SETMENU_APP)
            else:
                # add Keypatch menu
                menu = idaapi.add_menu_item("Edit/Keypatch/", "Patcher     (Ctrl-Alt-K)", "", 1, self.patcher, None)
                if menu is not None:
                    idaapi.add_menu_item("Edit/Keypatch/", "About", "", 1, self.about, None)
                    idaapi.add_menu_item("Edit/Keypatch/", "Check for update", "", 1, self.updater, None)
                    idaapi.add_menu_item("Edit/Keypatch/", "-", "", 1, self.menu_null, None)
                    idaapi.add_menu_item("Edit/Keypatch/", "Search", "", 1, self.search, None)
                    idaapi.add_menu_item("Edit/Keypatch/", "-", "", 1, self.menu_null, None)
                    idaapi.add_menu_item("Edit/Keypatch/", "Undo last patching", "", 1, self.undo, None)
                    idaapi.add_menu_item("Edit/Keypatch/", "Fill Range", "", 1, self.fill_range, None)
                elif idaapi.IDA_SDK_VERSION < 680:
                    # older IDAPython (such as in IDAPro 6.6) does add new submenu.
                    # in this case, put Keypatch menu in menu Edit \ Patch program
                    # not sure about v6.7, so to be safe we just check against v6.8
                    idaapi.add_menu_item("Edit/Patch program/", "-", "", 0, self.menu_null, None)
                    idaapi.add_menu_item("Edit/Patch program/", "Keypatch:: About", "", 0, self.about, None)
                    idaapi.add_menu_item("Edit/Patch program/", "Keypatch:: Check for update", "", 0, self.updater, None)
                    idaapi.add_menu_item("Edit/Patch program/", "Keypatch:: Search", "", 0, self.search, None)
                    idaapi.add_menu_item("Edit/Patch program/", "Keypatch:: Undo last patching", "", 0, self.undo, None)
                    idaapi.add_menu_item("Edit/Patch program/", "Keypatch:: Fill Range", "", 0, self.fill_range, None)
                    idaapi.add_menu_item("Edit/Patch program/", "Keypatch:: Patcher     (Ctrl-Alt-K)", "", 0, self.patcher, None)

            print("=" * 80)
            print("Keypatch v{0} (c) Nguyen Anh Quynh & Thanh Nguyen, 2018".format(VERSION))
            print("Keypatch is using Keystone v{0}".format(keystone.__version__))
            print("Keypatch Patcher's shortcut key is Ctrl-Alt-K")
            print("Use the same hotkey Ctrl-Alt-K to open 'Fill Range' window on a selected range of code")
            print("To revert (undo) the last patching, choose menu Edit | Keypatch | Undo last patching")
            print("Keypatch Search is available from menu Edit | Keypatch | Search")
            print("Find more information about Keypatch at http://keystone-engine.org/keypatch")

            self.load_configuration()

            print("=" * 80)
            self.kp_asm = Keypatch_Asm()

        return idaapi.PLUGIN_KEEP

    def term(self):
        if self.hooks is not None:
            self.hooks.unhook()
            self.hooks = None

        if self.opts is None:
            return
        #save configuration to file
        try:
            json.dump(self.opts, open(KP_CFGFILE, "wt"))
        except Exception as e:
            print("Keypatch: FAILED to save config file, with exception: {0}".format(str(e)))
        else:
            print("Keypatch: configuration is saved to {0}".format(KP_CFGFILE))

    # null handler
    def menu_null(self):
        pass

    # handler for About menu
    def about(self):
        f = About_Form(VERSION)
        f.Execute()
        f.Free()

    # handler for Check-for-Update menu
    def updater(self):
        (r, content) = url_download(KP_GITHUB_VERSION)
        content = to_string(content)
        if r == 0:
            # find stable version
            sig = 'VERSION_STABLE = "'
            tmp = content[content.find(sig)+len(sig):]
            version_stable = tmp[:tmp.find('"')]

            # compare with the current version
            if version_stable == VERSION:
                f = Update_Form(VERSION, "Good, you are already on the latest stable version!")
                f.Execute()
                # free this form
                f.Free()
            else:
                f = Update_Form(VERSION, "Download latest stable version {0} from http://keystone-engine.org/keypatch".format(version_stable))
                f.Execute()
                # free this form
                f.Free()
        else:
            # fail to download
            warning("ERROR: Keypatch failed to connect to internet (Github). Try again later.")
            print("Keypatch: FAILED to connect to Github to check for latest update. Try again later.")

    # handler for Undo menu
    def undo(self):
        global patch_info
        if len(patch_info) == 0:
            # TODO: disable Undo menu?
            warning("ERROR: Keypatch already got to the last undo patching!")
        else:
            (address, assembly, p_orig_data, patch_comment) = patch_info[-1]

            # undo the patch
            self.kp_asm.patch_code(address, None, None, None, None, orig_asm=[assembly], patch_data=p_orig_data, patch_comment=patch_comment, undo=True)
            del(patch_info[-1])

    # handler for Search menu
    def search(self):
        address = get_screen_ea()
        f = Keypatch_Search(self.kp_asm, address)
        f.Execute()
        f.Free()

    # handler for Patcher menu
    def patcher(self):
        # be sure that this arch is supported by Keystone
        if self.kp_asm.arch is None:
            warning("ERROR: Keypatch cannot handle this architecture (unsupported by Keystone), quit!")
            return

        selection, addr_begin, addr_end = read_range_selection()
        if selection:
            # call Fill Range function on this selected code
            return self.fill_range()

        address = get_screen_ea()

        if self.opts is None:
            self.load_configuration()

        init_assembly = None
        while True:
            f = Keypatch_Patcher(self.kp_asm, address, assembly=init_assembly, opts=self.opts)
            ok = f.Execute()
            if ok == 1:
                try:
                    syntax = None
                    if f.kp_asm.arch == KS_ARCH_X86:
                        syntax_id = f.c_syntax.value
                        syntax = self.kp_asm.get_syntax_by_idx(syntax_id)

                    assembly = f.c_assembly.value
                    self.opts = f.get_opts('c_opt_chk')
                    padding = (self.opts.get("c_opt_padding", 0) != 0)
                    comment = (self.opts.get("c_opt_comment", 0) != 0)

                    raw_assembly = self.kp_asm.ida_resolve(assembly, address)

                    print("Keypatch: attempting to modify \"{0}\" at 0x{1:X} to \"{2}\"".format(
                            self.kp_asm.ida_get_disasm(address), address, assembly))

                    length = self.kp_asm.patch_code(address, raw_assembly, syntax, padding, comment, None)
                    if length > 0:
                        # update start address pointing to the next instruction
                        init_assembly = None
                        address += length
                    else:
                        init_assembly = f.c_assembly.value
                        if length == 0:
                            warning("ERROR: Keypatch found invalid assembly [{0}]".format(assembly))
                        elif length == -1:
                            warning("ERROR: Keypatch failed to patch binary at 0x{0:X}!".format(address))
                        elif length == -2:
                            warning("ERROR: Keypatch can't read original data at 0x{0:X}, try again".format(address))

                except KsError as e:
                    print("Keypatch Error: {0}".format(e))
            else:   # Cancel
                break
            f.Free()

    # handler for Fill Range menu
    def fill_range(self):
        # be sure that this arch is supported by Keystone
        if self.kp_asm.arch is None:
            warning("ERROR: Keypatch cannot handle this architecture (unsupported by Keystone), quit!")
            return
               
        selection, addr_begin, addr_end = read_range_selection()
        if not selection:
            warning("ERROR: Keypatch requires a range to be selected for fill in, try again")
            return

        if self.opts is None:
            self.load_configuration()

        init_assembly = None
        f = Keypatch_FillRange(self.kp_asm, addr_begin, addr_end, assembly=init_assembly, opts=self.opts)
        ok = f.Execute()
        if ok == 1:
            try:
                syntax = None
                if f.kp_asm.arch == KS_ARCH_X86:
                    syntax_id = f.c_syntax.value
                    syntax = self.kp_asm.get_syntax_by_idx(syntax_id)

                assembly = f.c_assembly.value
                self.opts = f.get_opts('c_opt_chk')
                padding = (self.opts.get("c_opt_padding", 0) != 0)
                comment = (self.opts.get("c_opt_comment", 0) != 0)

                raw_assembly = self.kp_asm.ida_resolve(assembly, addr_begin)

                print("Keypatch: attempting to fill range [0x{0:X}:0x{1:X}] with \"{2}\"".format(
                    addr_begin, addr_end - 1, assembly))

                length = self.kp_asm.fill_code(addr_begin, addr_end, raw_assembly, syntax, padding, comment, None)
                if length == 0:
                    warning("ERROR: Keypatch failed to process this input.")
                    print("Keypatch: FAILED to process this input '{0}'".format(assembly))
                elif length == -1:
                    warning("ERROR: Keypatch failed to patch binary at 0x{0:X}!".format(addr_begin))

            except KsError as e:
                print("Keypatch Error: {0}".format(e))

        # free this form
        f.Free()

    def run(self, arg):
        self.patcher()


# register IDA plugin
def PLUGIN_ENTRY():
    return Keypatch_Plugin_t()


repository: awgil/ffxiv_reverse
file: idaplugins/ffutil.py
description: Various utilities for reversing FF14
wanted_name: ffutil
comment: A bunch of utilities to simplify reversing FF14
created_at: 2023-02-06 18:28:16+00:00
pushed_at: 2023-10-03 10:25:46+00:00
url:
https://github.com/awgil/ffxiv_reverse/blob/f35b6226c1478234ca2b7149f82d251cffca2f56/idaplugins/ffutil.py

content:
import idaapi
import ida_kernwin
import ida_struct
import ida_funcs
import ida_name
from PyQt5 import QtCore, QtGui, QtWidgets

def functions():
	start_ea = 0
	while True:
		nextfn = ida_funcs.get_next_func(start_ea)
		if not nextfn:
			break
		start_ea = nextfn.start_ea
		yield start_ea

def names():
	for i in range(ida_name.get_nlist_size()):
		yield (ida_name.get_nlist_ea(i), ida_name.get_nlist_name(i))

class dialog(QtWidgets.QDialog):
	_layout = None

	def __init__(self, name):
		QtWidgets.QDialog.__init__(self)
		self.setWindowTitle(name)
		self._layout = QtWidgets.QVBoxLayout()
		self.setLayout(self._layout)

	def add_widget(self, widget):
		self._layout.addWidget(widget)
		return widget

	def add_layout(self, layout):
		self._layout.addLayout(layout)
		return layout

	def add_buttons(self, ok_name = 'OK'):
		buttons = self.add_layout(QtWidgets.QHBoxLayout())
		ok = QtWidgets.QPushButton(ok_name)
		ok.setDefault(True)
		ok.clicked.connect(self.accept)
		buttons.addWidget(ok)
		cancel = QtWidgets.QPushButton('Cancel')
		cancel.clicked.connect(self.reject)
		buttons.addWidget(cancel)

	def run(self, ok_name = 'OK'):
		self.add_buttons(ok_name)
		return self.exec()

class smart_rename(idaapi.action_handler_t):
	def __init__(self):
		idaapi.action_handler_t.__init__(self)

	def activate(self, ctx):
		highlight = ida_kernwin.get_highlight(ida_kernwin.get_current_viewer())
		struct = ida_struct.get_struc_id(highlight[0]) if highlight else idaapi.BADADDR
		if struct == idaapi.BADADDR:
			print('Select a structure to rename')
			return 0

		original_name = ida_struct.get_struc_name(struct)
		related_names = self._related_names(original_name)

		dlg = dialog('Rename class and members')
		dlg_rename = dlg.add_widget(QtWidgets.QLineEdit(original_name))

		dlg_items = dlg.add_widget(QtWidgets.QListWidget())
		dlg_items_list = []
		for ea, name in related_names:
			item = QtWidgets.QListWidgetItem(name, dlg_items)
			item.setData(QtCore.Qt.UserRole, ea)
			item.setCheckState(QtCore.Qt.Checked)
			dlg_items_list.append(item)

		if dlg.run('Rename') == 0:
			return 0

		new_name = dlg_rename.text()
		if new_name == original_name:
			print('Name is not changed, nothing to do...')
			return 0

		print(f'Renaming struct {original_name} to {new_name}')
		ida_struct.set_struc_name(struct, new_name)
		for item in dlg_items_list:
			if item.checkState() == QtCore.Qt.Checked:
				ea = item.data(QtCore.Qt.UserRole)
				rel_ori_name = item.text()
				rel_new_name = rel_ori_name.replace(original_name, new_name)
				print(f'Renaming related {rel_ori_name} to {rel_new_name}')
				ida_name.set_name(ea, rel_new_name)
		return 1

	def update(self, ctx):
		return idaapi.AST_ENABLE_ALWAYS

	def _related_names(self, struct_name):
		return [(ea, name) for ea, name in names() if self._is_related_name(name, struct_name)]

	def _is_related_name(self, name, struct_name):
		adj_name = name[5:] if name.startswith('vtbl_') else name
		pos = adj_name.find(struct_name)
		if pos < 0:
			return False
		if pos > 0 and adj_name[pos-1] != ':' and adj_name[pos-1] != '.':
			return False
		pos += len(struct_name)
		if pos > len(adj_name):
			return False
		return pos == len(adj_name) or adj_name[pos] == ':' or adj_name[pos] == '.'

class vtable_sync(idaapi.action_handler_t):
	def __init__(self):
		idaapi.action_handler_t.__init__(self)

	def activate(self, ctx):
		ea = ida_kernwin.get_screen_ea()
		name = ida_name.get_name(ea)
		if not name.startswith('vtbl_'):
			print('Place cursor at the beginning of a virtual table and make sure it is called vtbl_ClassName')
			return 0

		name = name[5:]
		vt_struct = ida_struct.get_struc_id(name + '_vtbl')
		if vt_struct != idaapi.BADADDR:
			return self._sync_vtable(ea, name, vt_struct)
		else:
			return self._create_vtable(ea, name)

	def update(self, ctx):
		return idaapi.AST_ENABLE_ALWAYS

	def _sync_vtable(self, ea, classname, vtstruct):
		print('todo')
		return 0

	def _create_vtable(self, ea, classname):
		next_ea = min([addr for addr, name in names() if addr > ea])
		size = int((next_ea - ea) / 8)

		dlg = dialog(f'Create vtable for {classname}')
		
		if dlg.run('Create') == 0:
			return 0

		print(f'creating: max {size} entries')
		return 0

class ffutil(idaapi.plugin_t):
	flags = idaapi.PLUGIN_UNL
	comment = 'A bunch of utilities to simplify reversing FF14'
	help = ''
	wanted_name = 'ffutil'
	wanted_hotkey = 'Ctrl+Alt+M'

	def init(self):
		print('ffutil init')
		self._register_action('smart_rename', 'Rename class and all methods', smart_rename(), 'Ctrl+Shift+N')
		self._register_action('vtable_sync', 'Sync vtable definitions and signatures', vtable_sync(), 'Ctrl+Shift+V')
		return idaapi.PLUGIN_KEEP

	def run(self, arg=None):
		print('ffutil run')

	def term(self):
		print('ffutil term')

	def _register_action(self, name, label, handler, shortcut):
		deco_name = 'ffutil:' + name
		idaapi.unregister_action(deco_name)
		idaapi.register_action(idaapi.action_desc_t(deco_name, label, handler, shortcut))

def PLUGIN_ENTRY():
	return ffutil()


repository: BeaCox/ida-plugins
file: keypatch/keypatch.py
description: Some ida plugins i use (for ida9.0)
wanted_name: Keypatch Patcher
comment: Keypatch plugin for IDA Pro (using Keystone framework)
created_at: 2024-09-20 03:25:27+00:00
pushed_at: 2024-09-20 06:07:05+00:00
url:
https://github.com/BeaCox/ida-plugins/blob/83c25f5ab8bd3df32e0445431ce552c756580f89/keypatch/keypatch.py

content:
# -*- coding: utf-8 -*-

# Keypatch IDA Plugin, powered by Keystone Engine (http://www.keystone-engine.org).
# By Nguyen Anh Quynh & Thanh Nguyen, 2018.

# Keypatch is released under the GPL v2. See COPYING for more information.
# Find docs & latest version at http://keystone-engine.org/keypatch

# This IDA plugin includes 3 tools inside: Patcher, Fill Range & Search.
# Access to these tools via menu "Edit | Keypatch", or via right-click popup menu "Keypatch".

# Hotkey Ctrl-Alt-K opens either Patcher or "Fill Range" window, depending on context.
#  - If there is no code selection, hotkey opens Patcher dialog
#  - If a range of code is selected, hotkey opens "Fill Range" dialog

# To revert (undo) the last patching, choose menu "Edit | Keypatch | Undo last patching".
# To check for update version, choose menu "Edit | Keypatch | Check for update".

'''
    Do not use shortcut(Ctrk + Alt + K) for debugging, otherwise debugging features will not be available.
'''
#########################################################################################################
is_debug = False
if is_debug:
    ''' 
        Install pydevd:

        1. sudo pip install pydevd

        or

        2. Install pycharm-debug.egg, Ensure to use pycharm pro
        https://www.jetbrains.com/help/pycharm/remote-debugging-with-product.html

        # import site
        # site.addsitedir("/usr/local/lib/python2.7/site-packages")
    '''
    try:
        import pydevd

        pydevd.settrace(host='localhost',
                        port=51234,
                        stdoutToServer=True,
                        stderrToServer=True
                        )
    except Exception as e:
        print(e)
#########################################################################################################

import os
import re
import json
from keystone import *
import idc
import idaapi
import six
import ida_ida


################################ IDA 6/7 Compatibility import ###########################################
if idaapi.IDA_SDK_VERSION >= 700:
    import ida_search
    from idc import (
        get_operand_type,
        print_operand,
        get_item_end,
        get_item_head,
        get_item_size,
        get_sreg,
        has_value,
        get_full_flags,
        get_wide_byte,
        get_func_attr,
        set_func_end,
        warning,
        get_screen_ea
    )
    from idaapi import (
        get_bytes,
        plan_and_wait,
        is_mapped,
        Choose,
        jumpto
    )
else:
    from idc import (
        GetOpType as get_operand_type,
        GetOpnd as print_operand,
        ItemEnd as get_item_end,
        ItemHead as get_item_head,
        ItemSize as get_item_size,
        GetReg as get_sreg,
        hasValue as has_value,
        GetFlags as get_full_flags,
        Byte as get_wide_byte,
        GetFunctionAttr as get_func_attr,
        Jump as jumpto,
        Warning as warning,
        ScreenEA as get_screen_ea
    )
    from idaapi import (
        get_many_bytes as get_bytes,
        analyze_area as plan_and_wait,
        isEnabled as is_mapped,
        func_setend as set_func_end,
        Choose2 as Choose
    )
#########################################################################################################

# bleeding-edge version
# on a new release, this should be sync with VERSION_STABLE file
VERSION = "2.2"


MAX_INSTRUCTION_STRLEN = 256
MAX_ENCODING_LEN = 40
MAX_ADDRESS_LEN = 40
ENCODING_ERR_OUTPUT = "..."
KP_GITHUB_VERSION = "https://raw.githubusercontent.com/keystone-engine/keypatch/master/VERSION_STABLE"
KP_HOMEPAGE = "http://keystone-engine.org/keypatch"

X86_NOP = "\x90"

# Configuration file
KP_CFGFILE = os.path.join(idaapi.get_user_idadir(), "keypatch.cfg")

# save all the info on patching
patch_info = []

################################ Python2/3 Compatibility function #######################################
# Convert unicode/string/bytes to string
def to_string(s):
  # python3 bytes
  if six.PY3 and isinstance(s, bytes):
      return s.decode('latin-1')
  # python2 unicode
  elif six.PY2 and isinstance(s, six.text_type):
      return s.encode('utf-8')
  return str(s)

def to_hexstr(buf, sep=' '):
    # for python3 bytes
    if six.PY3 and isinstance(buf, bytes):
        return sep.join("{0:02x}".format(c) for c in buf).upper()
    return sep.join("{0:02x}".format(ord(c)) for c in buf).upper()
#########################################################################################################

################################ IDA 6/7 Compatibility function #########################################
def get_dtype(ea, op_idx):
    if idaapi.IDA_SDK_VERSION >= 700:
        insn = idaapi.insn_t()
        idaapi.decode_insn(insn, ea)
        dtype = insn.ops[op_idx].dtype
        dtyp_size = idaapi.get_dtype_size(dtype)
    else:
        dtype = idaapi.cmd.Operands[op_idx].dtyp
        dtyp_size = idaapi.get_dtyp_size(dtype)
    return dtype, dtyp_size

def set_comment(ea, comment):
    if idaapi.IDA_SDK_VERSION >= 700:
        idc.set_cmt(ea, comment, 0)
    else:
        idc.MakeComm(ea, comment)

def get_comment(ea):
    if idaapi.IDA_SDK_VERSION >= 700:
        return idc.get_cmt(ea, 0)
    return idc.Comment(ea)

def read_range_selection():
    if idaapi.IDA_SDK_VERSION >= 700:
        return idaapi.read_range_selection(None)
    return idaapi.read_selection()
#########################################################################################################

# return a normalized code, or None if input is invalid
def convert_hexstr(code):
    # normalize code
    code = code.lower()
    code = code.replace(' ', '')    # remove space
    code = code.replace('h', '')    # remove trailing 'h' in 90h
    code = code.replace('0x', '')   # remove 0x
    code = code.replace('\\x', '')  # remove \x
    code = code.replace(',', '')    # remove ,
    code = code.replace(';', '')    # remove ;
    code = code.replace('"', '')    # remove "
    code = code.replace("'", '')    # remove '
    code = code.replace("+", '')    # remove +

    # single-digit hexcode?
    if len(code) == 1 and ((code >= '0' and code <= '9') or (code >= 'a' and code <= 'f')):
        # stick 0 in front (so 'a' --> '0a')
        code = '0' + code

    # odd-length is invalid
    if len(code) % 2 != 0:
        return None

    try:
        hex_data = code.decode('hex')
        # we want a list of int
        return [ord(i) for i in hex_data]
    except:
        # invalid hex
        return None


# download a file from @url, then return (result, file-content)
# return (0, content) on success, or ({1|2}, None) on download failure
def url_download(url):
    # Import for python2
    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)


def get_name_value(_from, name):
    """
    Fixed: the return value truncated(32 bit) of get_name_value function that analyzed 64 bit binary file about ida64 for win.

    eg:
    type == idaapi.NT_BYTE
    (type, value) = idaapi.get_name_value(idc.BADADDR, "wcschr") # ida64 for win

    value = 0x14003d3f0L is correct  ida64 > 7.x for macOS
    value = 0x4003d3f0L is truncated ida64 >= 6.x for win, ida64 == 6.x for macOS

    :param _from: ea
    :param name: name string
    :return: tuple
    """
    name = to_string(name)
    (typ, value) = idaapi.get_name_value(_from, name)
    if typ == idaapi.NT_BYTE:  # typ is byte name (regular name)
        value = idaapi.get_name_ea(_from, name)
    return (typ, value)


## Main Keypatch class
class Keypatch_Asm:
    # supported architectures
    arch_lists = {
        "X86 16-bit": (KS_ARCH_X86, KS_MODE_16),                # X86 16-bit
        "X86 32-bit": (KS_ARCH_X86, KS_MODE_32),                # X86 32-bit
        "X86 64-bit": (KS_ARCH_X86, KS_MODE_64),                # X86 64-bit
        "ARM": (KS_ARCH_ARM, KS_MODE_ARM),                      # ARM
        "ARM Thumb": (KS_ARCH_ARM, KS_MODE_THUMB),              # ARM Thumb
        "ARM64 (ARMV8)": (KS_ARCH_ARM64, KS_MODE_LITTLE_ENDIAN),# ARM64
        "Hexagon": (KS_ARCH_HEXAGON, KS_MODE_BIG_ENDIAN),       # Hexagon
        "Mips32": (KS_ARCH_MIPS, KS_MODE_MIPS32),               # Mips32
        "Mips64": (KS_ARCH_MIPS, KS_MODE_MIPS64),               # Mips64
        "PowerPC 32": (KS_ARCH_PPC, KS_MODE_PPC32),             # PPC32
        "PowerPC 64": (KS_ARCH_PPC, KS_MODE_PPC64),             # PPC64
        "Sparc 32": (KS_ARCH_SPARC, KS_MODE_SPARC32),           # Sparc32
        "Sparc 64": (KS_ARCH_SPARC, KS_MODE_SPARC64),           # Sparc64
        "SystemZ": (KS_ARCH_SYSTEMZ, KS_MODE_BIG_ENDIAN),       # SystemZ
    }

    endian_lists = {
        "Little Endian": KS_MODE_LITTLE_ENDIAN,                 # little endian
        "Big Endian": KS_MODE_BIG_ENDIAN,                       # big endian
    }

    syntax_lists = {
        "Intel": KS_OPT_SYNTAX_INTEL,
        "Nasm": KS_OPT_SYNTAX_NASM,
        "AT&T": KS_OPT_SYNTAX_ATT
    }

    def __init__(self, arch=None, mode=None):
        # update current arch and mode
        self.update_hardware_mode()

        # override arch & mode if provided
        if arch is not None:
            self.arch = arch
        if mode is not None:
            self.mode = mode

        # IDA uses Intel syntax by default
        self.syntax = KS_OPT_SYNTAX_INTEL

    # return Keystone arch & mode (with endianess included)
    @staticmethod
    def get_hardware_mode():
        (arch, mode) = (None, None)

        # heuristically detect hardware setup
        if idaapi.IDA_SDK_VERSION < 900:
            info = idaapi.get_inf_structure()
        
        try:
            if idaapi.IDA_SDK_VERSION < 900:
                cpuname = info.procname.lower()
            else:
                cpuname = ida_ida.inf_get_procname().lower()
        except:
            if idaapi.IDA_SDK_VERSION < 900:
                cpuname = info.procname.lower()
            else:
                cpuname = ida_ida.inf_get_procname().lower()

        try:
            # since IDA7 beta 3 (170724) renamed inf.mf -> is_be()/set_be()
            if idaapi.IDA_SDK_VERSION < 900:
                is_be = idaapi.cvar.inf.is_be()
            else:
                is_be = ida_ida.inf_is_be()
        except:
            # older IDA versions
            if idaapi.IDA_SDK_VERSION < 900:
                is_be = idaapi.cvar.inf.is_be()
            else:
                is_be = ida_ida.inf_is_be()
        # print("Keypatch BIG_ENDIAN = %s" %is_be)
        
        if idaapi.IDA_SDK_VERSION < 900:
            if cpuname == "metapc":
                arch = KS_ARCH_X86
                if info.is_64bit():
                    mode = KS_MODE_64
                elif info.is_32bit():
                    mode = KS_MODE_32
                else:
                    mode = KS_MODE_16

            elif cpuname.startswith("arm"):
                # ARM or ARM64
                if info.is_64bit():
                    arch = KS_ARCH_ARM64
                    if is_be:
                        mode = KS_MODE_BIG_ENDIAN
                    else:
                        mode = KS_MODE_LITTLE_ENDIAN
                else:
                    arch = KS_ARCH_ARM
                    # either big-endian or little-endian
                    if is_be:
                        mode = KS_MODE_ARM | KS_MODE_BIG_ENDIAN
                    else:
                        mode = KS_MODE_ARM | KS_MODE_LITTLE_ENDIAN
            elif cpuname.startswith("sparc"):
                arch = KS_ARCH_SPARC
                if info.is_64bit():
                    mode = KS_MODE_SPARC64
                else:
                    mode = KS_MODE_SPARC32
                if is_be:
                    mode |= KS_MODE_BIG_ENDIAN
                else:
                    mode |= KS_MODE_LITTLE_ENDIAN
            elif cpuname.startswith("ppc"):
                arch = KS_ARCH_PPC
                if info.is_64bit():
                    mode = KS_MODE_PPC64
                else:
                    mode = KS_MODE_PPC32
                if cpuname == "ppc":
                    # do not support Little Endian mode for PPC
                    mode += KS_MODE_BIG_ENDIAN
            elif cpuname.startswith("mips"):
                arch = KS_ARCH_MIPS
                if info.is_64bit():
                    mode = KS_MODE_MIPS64
                else:
                    mode = KS_MODE_MIPS32
                if is_be:
                    mode |= KS_MODE_BIG_ENDIAN
                else:
                    mode |= KS_MODE_LITTLE_ENDIAN
            elif cpuname.startswith("systemz") or cpuname.startswith("s390x"):
                arch = KS_ARCH_SYSTEMZ
                mode = KS_MODE_BIG_ENDIAN
        else:
            if cpuname == "metapc":
                arch = KS_ARCH_X86
                if ida_ida.inf_is_64bit():
                    mode = KS_MODE_64
                elif ida_ida.inf_is_32bit_exactly():
                    mode = KS_MODE_32
                else:
                    mode = KS_MODE_16

            elif cpuname.startswith("arm"):
                # ARM or ARM64
                if ida_ida.inf_is_64bit():
                    arch = KS_ARCH_ARM64
                    if is_be:
                        mode = KS_MODE_BIG_ENDIAN
                    else:
                        mode = KS_MODE_LITTLE_ENDIAN
                else:
                    arch = KS_ARCH_ARM
                    # either big-endian or little-endian
                    if is_be:
                        mode = KS_MODE_ARM | KS_MODE_BIG_ENDIAN
                    else:
                        mode = KS_MODE_ARM | KS_MODE_LITTLE_ENDIAN
            elif cpuname.startswith("sparc"):
                arch = KS_ARCH_SPARC
                if ida_ida.inf_is_64bit():
                    mode = KS_MODE_SPARC64
                else:
                    mode = KS_MODE_SPARC32
                if is_be:
                    mode |= KS_MODE_BIG_ENDIAN
                else:
                    mode |= KS_MODE_LITTLE_ENDIAN
            elif cpuname.startswith("ppc"):
                arch = KS_ARCH_PPC
                if ida_ida.inf_is_64bit():
                    mode = KS_MODE_PPC64
                else:
                    mode = KS_MODE_PPC32
                if cpuname == "ppc":
                    # do not support Little Endian mode for PPC
                    mode += KS_MODE_BIG_ENDIAN
            elif cpuname.startswith("mips"):
                arch = KS_ARCH_MIPS
                if ida_ida.inf_is_64bit():
                    mode = KS_MODE_MIPS64
                else:
                    mode = KS_MODE_MIPS32
                if is_be:
                    mode |= KS_MODE_BIG_ENDIAN
                else:
                    mode |= KS_MODE_LITTLE_ENDIAN
            elif cpuname.startswith("systemz") or cpuname.startswith("s390x"):
                arch = KS_ARCH_SYSTEMZ
                mode = KS_MODE_BIG_ENDIAN            
                        
        return (arch, mode)

    def update_hardware_mode(self):
        (self.arch, self.mode) = self.get_hardware_mode()

    # normalize assembly code
    # remove comment at the end of assembly code
    @staticmethod
    def asm_normalize(text):
        text = ' '.join(text.split())
        if text.rfind(';') != -1:
            return text[:text.rfind(';')].strip()

        return text.strip()

    @staticmethod
    # check if input address is valid
    # return
    #       -1  invalid address at target binary
    #        0  type mismatch of input address
    #        1  valid address at target binary
    def check_address(address):
        try:
            if is_mapped(address):
                return 1
            else:
                return -1
        except:
            # invalid type
            return 0

    ### resolve IDA names from input asm code
    # todo: a better syntax parser for all archs
    def ida_resolve(self, assembly, address=idc.BADADDR):
        def _resolve(_op, ignore_kw=True):
            names = re.findall(r"[\$a-z0-9_:\.]+", _op, re.I)

            # try to resolve all names
            for name in names:
                # ignore known keywords
                if ignore_kw and name in ('byte', 'near', 'short', 'word', 'dword', 'ptr', 'offset'):
                    continue

                sym = name

                # split segment reg
                parts = name.partition(':')
                if parts[2] != '':
                    sym = parts[2]

                (typ, value) = get_name_value(address, sym)

                # skip if name doesn't exist or segment / segment registers
                if typ in (idaapi.NT_SEG, idaapi.NT_NONE):
                    continue

                _op = _op.replace(sym, '0x{0:X}'.format(value))

            return _op

        if self.check_address(address) == 0:
            print("Keypatch: WARNING: invalid input address {0}".format(address))
            return assembly

        # for now, we only support IDA name resolve for X86, ARM, ARM64, MIPS, PPC, SPARC
        if not (self.arch in (KS_ARCH_X86, KS_ARCH_ARM, KS_ARCH_ARM64, KS_ARCH_MIPS, KS_ARCH_PPC, KS_ARCH_SPARC)):
            return assembly

        _asm = assembly.partition(' ')
        mnem = _asm[0]
        opers = _asm[2].split(',')

        for idx, op in enumerate(opers):
            _op = list(op.partition('['))
            ignore_kw = True
            if _op[1] == '':
                _op[2] = _op[0]
                _op[0] = ''
            else:
                _op[0] = _resolve(_op[0], ignore_kw=True)
                ignore_kw = False

            _op[2] = _resolve(_op[2], ignore_kw=ignore_kw)

            opers[idx] = ''.join(_op)

        asm = "{0} {1}".format(mnem, ','.join(opers))
        return asm

    # return bytes of instruction or data
    # return None on failure
    def ida_get_item(self, address, hex_output=False):
        if self.check_address(address) != 1:
            # not a valid address
            return (None, 0)

        # return None if address is in the middle of instruction / data
        if address != get_item_head(address):
            return (None, 0)

        size = get_item_size(address)
        item = get_bytes(address, size)

        if item is None:
            return (None, 0)

        if hex_output:
            item = to_hexstr(item)

        return (item, size)

    @staticmethod
    def get_op_dtype_name(ea, op_idx):
        dtyp_lists = {
            idaapi.dt_byte: 'byte',     #  8 bit
            idaapi.dt_word: 'word',     #  16 bit
            idaapi.dt_dword: 'dword',   #  32 bit
            idaapi.dt_float: 'dword',   #  4 byte
            idaapi.dt_double: 'dword',  #  8 byte
            #idaapi.dt_tbyte = 5        #  variable size (ph.tbyte_size)
            #idaapi.dt_packreal = 6         #  packed real format for mc68040
            idaapi.dt_qword: 'qword',   #  64 bit
            idaapi.dt_byte16: 'xmmword',#  128 bit
            #idaapi.dt_code = 9         #  ptr to code (not used?)
            #idaapi.dt_void = 10        #  none
            #idaapi.dt_fword = 11       #  48 bit
            #idaapi.dt_bitfild = 12     #  bit field (mc680x0)
            #idaapi.dt_string = 13      #  pointer to asciiz string
            #idaapi.dt_unicode = 14     #  pointer to unicode string
            #idaapi.dt_3byte = 15       #  3-byte data
            #idaapi.dt_ldbl = 16        #  long double (which may be different from tbyte)
            idaapi.dt_byte32: 'ymmword',# 256 bit
        }

        dtype, dtyp_size = get_dtype(ea, op_idx)
        if dtype == idaapi.dt_tbyte and dtyp_size == 10:
          return 'xword'

        dtyp_name = dtyp_lists.get(dtype, None)

        return dtyp_name

    # return asm instructions from start to end
    def ida_get_disasm_range(self, start, end):
        codes = []
        while start < end:
            asm = self.asm_normalize(idc.GetDisasm(start))
            if asm == None:
                asm = ''
            codes.append(asm)
            start = start + get_item_size(start)

        return codes

    # get disasm from IDA
    # return '' on invalid address
    def ida_get_disasm(self, address, fixup=False):

        def GetMnem(asm):
            sp = asm.find(' ')
            if sp == -1:
                return asm
            return asm[:sp]

        if self.check_address(address) != 1:
            # not a valid address
            return ''

        # return if address is in the middle of instruction / data
        if address != get_item_head(address):
            return ''

        asm = self.asm_normalize(idc.GetDisasm(address))
        # for now, only support IDA syntax fixup for Intel CPU
        if not fixup or self.arch != KS_ARCH_X86:
            return asm

        # KS_ARCH_X86 mode
        # rebuild disasm code from IDA
        i = 0
        mnem = GetMnem(asm)
        if mnem == '' or mnem in ('rep', 'repne', 'repe'):
            return asm

        opers = []
        while get_operand_type(address, i) > 0 and i < 6:
            t = get_operand_type(address, i)
            o = print_operand(address, i)

            if t in (idc.o_mem, idc.o_displ):
                parts = list(o.partition(':'))
                if parts[2] == '':
                    parts[2] = parts[0]
                    parts[0] = ''

                if '[' not in parts[2]:
                    parts[2] = '[{0}]'.format(parts[2])

                o = ''.join(parts)

                if 'ptr ' not in o:
                    dtyp_name = self.get_op_dtype_name(address, i)
                    if dtyp_name != None:
                        o = "{0} ptr {1}".format(dtyp_name, o)

            opers.append(o)
            i += 1

        asm = mnem
        for o in opers:
            if o != '':
                asm = "{0} {1},".format(asm, o)

        asm = asm.strip(',')
        return asm

    # assemble code with Keystone
    # return (encoding, count), or (None, 0) on failure
    def assemble(self, assembly, address, arch=None, mode=None, syntax=None):

        # return assembly with arithmetic equation evaluated
        def eval_operand(assembly, start, stop, prefix=''):
            imm = assembly[start+1:stop]
            try:
                eval_imm = eval(imm)
                if eval_imm > 0x80000000:
                    eval_imm = 0xffffffff - eval_imm
                    eval_imm += 1
                    eval_imm = -eval_imm
                return assembly.replace(prefix + imm, prefix + hex(eval_imm))
            except:
                return assembly

        # IDA uses different syntax from Keystone
        # sometimes, we can convert code to be consumable by Keystone
        def fix_ida_syntax(assembly):

            # return True if this insn needs to be fixed
            def check_arm_arm64_insn(arch, mnem):
                if arch == KS_ARCH_ARM:
                    if mnem.startswith("ldr") or mnem.startswith("str"):
                        return True
                    return False
                elif arch == KS_ARCH_ARM64:
                    if mnem.startswith("ldr") or mnem.startswith("str"):
                        return True
                    return mnem in ("stp")
                return False

            # return True if this insn needs to be fixed
            def check_ppc_insn(mnem):
                return mnem in ("stw")

            # replace the right most string occurred
            def rreplace(s, old, new):
                li = s.rsplit(old, 1)
                return new.join(li)

            # convert some ARM pre-UAL assembly to UAL, so Keystone can handle it
            # example: streqb --> strbeq
            def fix_arm_ual(mnem, assembly):
                # TODO: this is not an exhaustive list yet
                if len(mnem) != 6:
                    return assembly

                if (mnem[-1] in ('s', 'b', 'h', 'd')):
                    #print(">> 222", mnem[3:5])
                    if mnem[3:5] in ("cc", "eq", "ne", "hs", "lo", "mi", "pl", "vs", "vc", "hi", "ls", "ge", "lt", "gt", "le", "al"):
                        return assembly.replace(mnem, mnem[:3] + mnem[-1] + mnem[3:5], 1)

                return assembly

            if self.arch != KS_ARCH_X86:
                assembly = assembly.lower()
            else:
                # Keystone does not support immediate 0bh, but only 0Bh
                assembly = assembly.upper()

            # however, 0X must be converted to 0x
            # Keystone should fix this limitation in the future
            assembly = assembly.replace("0X", " 0x")

            _asm = assembly.partition(' ')
            mnem = _asm[0]
            if mnem == '':
                return assembly

            # for PPC, Keystone does not accept registers with 'r' prefix,
            # but only the number behind. lets try to fix that here by
            # removing the prefix 'r'.
            if self.arch == KS_ARCH_PPC:
                for n in range(32):
                    r = " r%u," %n
                    if r in assembly:
                        assembly = assembly.replace(r, " %u," %n)
                for n in range(32):
                    r = "(r%u)" %n
                    if r in assembly:
                        assembly = assembly.replace(r, "(%u)" %n)
                for n in range(32):
                    r = ", r%u" %n
                    if assembly.endswith(r):
                        assembly = rreplace(assembly, r, ", %u" %n)

            if self.arch == KS_ARCH_X86:
                if mnem == "RETN":
                    # replace retn with ret
                    return assembly.replace('RETN', 'RET', 1)
                if 'OFFSET ' in assembly:
                    return assembly.replace('OFFSET ', ' ')
                if mnem in ('CALL', 'JMP') or mnem.startswith('LOOP'):
                    # remove 'NEAR PTR'
                    if ' NEAR PTR ' in assembly:
                        return assembly.replace(' NEAR PTR ', ' ')
                elif mnem[0] == 'J':
                    # JMP instruction
                    if ' SHORT ' in assembly:
                        # remove ' short '
                        return assembly.replace(' SHORT ', ' ')
            elif self.arch in (KS_ARCH_ARM, KS_ARCH_ARM64, KS_ARCH_PPC):
                # *** ARM
                # LDR     R1, [SP+rtld_fini],#4
                # STR     R2, [SP,#-4+rtld_fini]!
                # STR     R0, [SP,#fini]!
                # STR     R12, [SP,#4+var_8]!

                # *** ARM64
                # STP     X29, X30, [SP,#-0x10+var_150]!
                # STR     W0, [X29,#0x150+var_8]
                # LDR     X0, [X0,#(qword_4D6678 - 0x4D6660)]
                # TODO:
                # ADRP    X19, #interactive@PAGE

                # *** PPC
                # stw     r5, 0x120+var_108(r1)
                
                if self.arch == KS_ARCH_ARM and mode == KS_MODE_THUMB:
                    assembly = assembly.replace('movt.w', 'movt')

                if self.arch == KS_ARCH_ARM:
                    #print(">> before UAL fix: ", assembly)
                    assembly = fix_arm_ual(mnem, assembly)
                    #print(">> after UAL fix: ", assembly)

                if check_arm_arm64_insn(self.arch, mnem) or (("[" in assembly) and ("]" in assembly)):
                    bang = assembly.find('#')
                    bracket = assembly.find(']')
                    if bang != -1 and bracket != -1 and bang < bracket:
                        return eval_operand(assembly, bang, bracket, '#')
                    elif '+0x0]' in assembly:
                        return assembly.replace('+0x0]', ']')
                elif check_ppc_insn(mnem):
                    start = assembly.find(', ')
                    stop = assembly.find('(')
                    if start != -1 and stop != -1 and start < stop:
                        return eval_operand(assembly, start, stop)
            return assembly

        def is_thumb(address):
            return get_sreg(address, 'T') == 1

        if self.check_address(address) == 0:
            return (None, 0)

        # use default syntax, arch and mode if not provided
        if syntax is None:
            syntax = self.syntax
        if arch is None:
            arch = self.arch
        if mode is None:
            mode = self.mode

        if arch == KS_ARCH_ARM and is_thumb(address):
            mode = KS_MODE_THUMB

        try:
            ks = Ks(arch, mode)
            if arch == KS_ARCH_X86:
                ks.syntax = syntax
            encoding, count = ks.asm(fix_ida_syntax(assembly), address)
        except KsError as e:
            # keep the below code for debugging
            #print("Keypatch Error: {0}".format(e))
            #print("Original asm: {0}".format(assembly))
            #print("Fixed up asm: {0}".format(fix_ida_syntax(assembly)))
            encoding, count = None, 0

        return (encoding, count)


    # patch at address, return the number of written bytes & original data
    # this process can fail in some cases
    @staticmethod
    def patch_raw(address, patch_data, size):
        ea = address
        orig_data = ''

        while ea < (address + size):

            if not has_value(get_full_flags(ea)):
                print("Keypatch: FAILED to read data at 0x{0:X}".format(ea))
                break

            orig_byte = get_wide_byte(ea)
            orig_data += chr(orig_byte)
            patch_byte = ord(patch_data[ea - address])

            if patch_byte != orig_byte:
                # patch one byte
                if idaapi.patch_byte(ea, patch_byte) != 1:
                    print("Keypatch: FAILED to patch byte at 0x{0:X} [0x{1:X}]".format(ea, patch_byte))
                    break
            ea += 1
        return (ea - address, orig_data)

    # patch at address, return the number of written bytes & original data
    # on patch failure, we revert to the original code, then return (None, None)
    def patch(self, address, patch_data, size):
        # save original function end to fix IDA re-analyze issue after patching
        orig_func_end = get_func_attr(address, idc.FUNCATTR_END)

        (patched_len, orig_data) = self.patch_raw(address, patch_data, size)

        if size != patched_len:
            # patch failure
            if patched_len > 0:
                # revert the changes
                (rlen, _) = self.patch_raw(address, orig_data, patched_len)
                if rlen == patched_len:
                    print("Keypatch: successfully reverted changes of {0:d} byte(s) at 0x{1:X} [{2}]".format(
                                        patched_len, address, to_hexstr(orig_data)))
                else:
                    print("Keypatch: FAILED to revert changes of {0:d} byte(s) at 0x{1:X} [{2}]".format(
                                        patched_len, address, to_hexstr(orig_data)))

            return (None, None)

        # ask IDA to re-analyze the patched area
        if orig_func_end == idc.BADADDR:
            # only analyze patched bytes, otherwise it would take a lot of time to re-analyze the whole binary
            plan_and_wait(address, address + patched_len + 1)
        else:
            plan_and_wait(address, orig_func_end)

            # try to fix IDA function re-analyze issue after patching
            set_func_end(address, orig_func_end)

        return (patched_len, orig_data)

    # return number of bytes patched
    # return
    #    0  Invalid assembly
    #   -1  PatchByte failure
    #   -2  Can't read original data
    #   -3  Invalid address
    def patch_code(self, address, assembly, syntax, padding, save_origcode, orig_asm=None, patch_data=None, patch_comment=None, undo=False):
        global patch_info

        if self.check_address(address) != 1:
            # not a valid address
            return -3

        orig_comment = get_comment(address)
        if orig_comment is None:
            orig_comment = ''

        nop_comment = ""
        padding_len = 0
        if not undo:
            # we are patching via Patcher
            (orig_encoding, orig_len) = self.ida_get_item(address)
            if (orig_encoding, orig_len) == (None, 0):
                return -2

            (encoding, count) = self.assemble(assembly, address, syntax=syntax)
            if encoding is None:
                return 0

            patch_len = len(encoding)
            patch_data = ''.join(chr(c) for c in encoding)

            if patch_data == orig_encoding:
                #print("Keypatch: no need to patch, same encoding data [{0}] at 0x{1:X}".format(to_hexstr(orig_encoding), address))
                return orig_len

            # for now, only support NOP padding on Intel CPU
            if padding and self.arch == KS_ARCH_X86:
                if patch_len < orig_len:
                    padding_len = orig_len - patch_len
                    patch_len = orig_len
                    patch_data = patch_data.ljust(patch_len, X86_NOP)
                elif patch_len > orig_len:
                    patch_end = address + patch_len - 1
                    ins_end = get_item_end(patch_end)
                    padding_len = ins_end - patch_end - 1

                    if padding_len > 0:
                        patch_len = ins_end - address
                        patch_data = patch_data.ljust(patch_len, X86_NOP)

                if padding_len > 0:
                    nop_comment = "\nKeypatch padded NOP to next boundary: {0} bytes".format(padding_len)

            orig_asm = self.ida_get_disasm_range(address, address + patch_len)
        else:
            # we are reverting the change via "Undo" menu
            patch_len = len(patch_data)

        (plen, p_orig_data) = self.patch(address, patch_data, patch_len)
        if plen is None:
            # failed to patch
            return -1

        if not undo: # we are patching
            new_patch_comment = None
            if save_origcode is True:
                # append original instruction to comments
                if orig_comment == '':
                    new_patch_comment = "Keypatch modified this from:\n  {0}{1}".format('\n  '.join(orig_asm), nop_comment)
                else:
                    new_patch_comment = "\nKeypatch modified this from:\n  {0}{1}".format('\n  '.join(orig_asm), nop_comment)

                new_comment = "{0}{1}".format(orig_comment, new_patch_comment)
                set_comment(address, new_comment)

            if padding_len == 0:
                print("Keypatch: successfully patched {0:d} byte(s) at 0x{1:X} from [{2}] to [{3}]".format(plen,
                                        address, to_hexstr(p_orig_data), to_hexstr(patch_data)))
            else:
                print("Keypatch: successfully patched {0:d} byte(s) at 0x{1:X} from [{2}] to [{3}], with {4} byte(s) NOP padded".format(plen,
                                        address, to_hexstr(p_orig_data), to_hexstr(patch_data), padding_len))
            # save this patching for future "undo"
            patch_info.append((address, assembly, p_orig_data, new_patch_comment))
        else:   # we are reverting
            if patch_comment:
                # clean previous IDA comment by replacing it with ''
                new_comment = orig_comment.replace(patch_comment, '')
                set_comment(address, new_comment)

            print("Keypatch: successfully reverted {0:d} byte(s) at 0x{1:X} from [{2}] to [{3}]".format(plen,
                                        address, to_hexstr(p_orig_data), to_hexstr(patch_data)))

        return plen

    # fill a range of code [addr_begin, addr_end].
    # return the length of patched area
    # on failure, return 0 = wrong input, -1 = failed to patch
    def fill_code(self, addr_begin, addr_end, assembly, syntax, padding, save_origcode, orig_asm=None):
        # treat input as assembly code first
        (encoding, _) =  self.assemble(assembly, addr_begin, syntax=syntax)

        if encoding is None:
            # input might be a hexcode string. try to convert it to raw bytes
            encoding = convert_hexstr(assembly)

        if encoding is None:
            # invalid input: this is neither assembly nor hexcode string
            return 0

        # save original assembly code before overwritting them
        orig_asm = self.ida_get_disasm_range(addr_begin, addr_end)

        # save original comment at addr_begin
        # TODO: save comments in this range, but how to interleave them?
        orig_comment = get_comment(addr_begin)
        if orig_comment is None:
            orig_comment = ''

        patch_data = ""
        assembly_new = []
        size = addr_end - addr_begin
        # calculate filling data
        encode_chr = ''.join(chr(c) for c in encoding)
        while True:
            if len(patch_data) + len(encode_chr) <= size:
                patch_data = patch_data + encode_chr
                assembly_new += [assembly.strip()]
            else:
                break

        # for now, only support NOP padding on Intel CPU
        if padding and self.arch == KS_ARCH_X86:
            for i in range(size -len(patch_data)):
                assembly_new += ["nop"]
            patch_data = patch_data.ljust(size, X86_NOP)

        (plen, p_orig_data) = self.patch(addr_begin, patch_data, len(patch_data))
        if plen is None:
            # failed to patch
            return -1

        new_patch_comment = ''
        # append original instruction to comments
        if save_origcode is True:
            if orig_comment == '':
                new_patch_comment = "Keypatch filled range [0x{0:X}:0x{1:X}] ({2} bytes), replaced:\n  {3}".format(addr_begin, addr_end - 1, addr_end - addr_begin, '\n  '.join(orig_asm))
            else:
                new_patch_comment = "\nKeypatch filled range [0x{0:X}:0x{1:X}] ({2} bytes), replaced:\n  {3}".format(addr_begin, addr_end - 1, addr_end - addr_begin, '\n  '.join(orig_asm))

            new_comment = "{0}{1}".format(orig_comment, new_patch_comment)
            set_comment(addr_begin, new_comment)

        print("Keypatch: successfully filled range [0x{0:X}:0x{1:X}] ({2} bytes) with \"{3}\", replaced \"{4}\"".format(
                    addr_begin, addr_end - 1, addr_end - addr_begin, assembly, '; '.join(orig_asm)))

        # save this modification for future "undo"
        patch_info.append((addr_begin, '\n  '.join(assembly_new), p_orig_data, new_patch_comment))

        return plen


    ### Form helper functions
    @staticmethod
    def dict_to_ordered_list(dictionary):
        l = sorted(list(dictionary.items()), key=lambda t: t[0], reverse=False)
        keys = [i[0] for i in l]
        values = [i[1] for i in l]

        return (keys, values)

    def get_value_by_idx(self, dictionary, idx, default=None):
        (keys, values) = self.dict_to_ordered_list(dictionary)

        try:
            val = values[idx]
        except IndexError:
            val = default

        return val

    def find_idx_by_value(self, dictionary, value, default=None):
        (keys, values) = self.dict_to_ordered_list(dictionary)

        try:
            idx = values.index(value)
        except:
            idx = default

        return idx

    def get_arch_by_idx(self, idx):
        return self.get_value_by_idx(self.arch_lists, idx)

    def find_arch_idx(self, arch, mode):
        return self.find_idx_by_value(self.arch_lists, (arch, mode))

    def get_syntax_by_idx(self, idx):
        return self.get_value_by_idx(self.syntax_lists, idx, self.syntax)

    def find_syntax_idx(self, syntax):
        return self.find_idx_by_value(self.syntax_lists, syntax)
    ### /Form helper functions


# Common ancestor form to be derived by Patcher, FillRange & Search
class Keypatch_Form(idaapi.Form):
    # prepare for form initializing
    def setup(self, kp_asm, address, assembly=None):
        self.kp_asm = kp_asm
        self.address = address

        # update ordered list of arch and syntax
        self.syntax_keys = self.kp_asm.dict_to_ordered_list(self.kp_asm.syntax_lists)[0]
        self.arch_keys = self.kp_asm.dict_to_ordered_list(self.kp_asm.arch_lists)[0]

        # update current arch & mode
        self.kp_asm.update_hardware_mode()

        # find right value for c_arch & c_endian controls
        mode = self.kp_asm.mode
        self.endian_id = 0   # little endian
        if self.kp_asm.mode & KS_MODE_BIG_ENDIAN:
            self.endian_id = 1   # big endian
            mode = self.kp_asm.mode - KS_MODE_BIG_ENDIAN

        self.arch_id = self.kp_asm.find_arch_idx(self.kp_asm.arch, mode)

        self.syntax_id = 0  # to make non-X86 arch happy
        if self.kp_asm.arch == KS_ARCH_X86:
            self.syntax_id = self.kp_asm.find_syntax_idx(self.kp_asm.syntax)

        # get original instruction and bytes
        self.orig_asm = kp_asm.ida_get_disasm(address)
        (self.orig_encoding, self.orig_len) = kp_asm.ida_get_item(address, hex_output=True)
        if self.orig_encoding is None:
            self.orig_encoding = ''

        if assembly is None:
            self.asm = self.kp_asm.ida_get_disasm(self.address, fixup=True)
        else:
            self.asm = assembly

    # update Encoding control
    # return True on success, False on failure
    def _update_encoding(self, arch, mode):
        try:
            syntax = None
            if arch == KS_ARCH_X86:
                syntax_id = self.GetControlValue(self.c_syntax)
                syntax = self.kp_asm.get_syntax_by_idx(syntax_id)

            address = self.GetControlValue(self.c_addr)
            try:
                is_mapped(address)
            except:
                # invalid address value
                address = 0

            assembly = self.GetControlValue(self.c_assembly)
            raw_assembly = self.kp_asm.ida_resolve(assembly, address)
            self.SetControlValue(self.c_raw_assembly, raw_assembly)

            (encoding, count) =  self.kp_asm.assemble(raw_assembly, address, arch=arch,
                                                    mode=mode, syntax=syntax)

            if encoding is None:
                self.SetControlValue(self.c_encoding, ENCODING_ERR_OUTPUT)
                self.SetControlValue(self.c_encoding_len, 0)
                return False
            else:
                text = ""
                for byte in encoding:
                    text += "%02X " % byte
                text.strip()
                if text == "":
                    # error?
                    self.SetControlValue(self.c_encoding, ENCODING_ERR_OUTPUT)
                    return False
                else:
                    self.SetControlValue(self.c_encoding, text.strip())
                    self.SetControlValue(self.c_encoding_len, len(encoding))
                    return True
        except Exception as e:
            print(str(e))
            import traceback
            traceback.print_exc()
            self.SetControlValue(self.c_encoding, ENCODING_ERR_OUTPUT)
            return False

    # callback to be executed when any form control changed
    def OnFormChange(self, fid):
        return 1

    # update Patcher & Fillrange controls
    def update_patchform(self, fid):
        self.EnableField(self.c_endian, False)
        self.EnableField(self.c_addr, False)

        (arch, mode) = (self.kp_asm.arch, self.kp_asm.mode)
        # assembly is focused
        self.SetFocusedField(self.c_assembly)

        if arch == KS_ARCH_X86:
            # do not show Endian control
            self.ShowField(self.c_endian, False)
            # allow to choose Syntax
            self.ShowField(self.c_syntax, True)
            self.ShowField(self.c_opt_padding, True)
        else:   # do not show Syntax control for non-X86 mode
            self.ShowField(self.c_syntax, False)
            # for now, we do not support padding for non-X86 archs
            self.ShowField(self.c_opt_padding, False)
            #self.EnableField(self.c_opt_padding, False)

        # update other controls & Encoding with live assembling
        self.update_controls(arch, mode)

        return 1

    # update some controls - including Encoding control
    def update_controls(self, arch, mode):
        # Fixup & Encoding-len are read-only controls
        self.EnableField(self.c_raw_assembly, False)
        self.EnableField(self.c_encoding_len, False)

        # Encoding is enable to allow user to select & copy
        self.EnableField(self.c_encoding, True)

        if self.GetControlValue(self.c_endian) == 1:
            endian = KS_MODE_BIG_ENDIAN
        else:
            endian = KS_MODE_LITTLE_ENDIAN

        # update encoding with live assembling
        self._update_encoding(arch, mode | endian)

        return 1

    # get Patcher/FillRange options
    def get_opts(self, name=None):
        names = self.c_opt_chk.children_names
        val = self.c_opt_chk.value
        opts = {}
        for i in range(len(names)):
            opts[names[i]] = val & (2**i)

        if name != None:
            opts[name] = val

        return opts


# Fill Range form
class Keypatch_FillRange(Keypatch_Form):
    def __init__(self, kp_asm, addr_begin, addr_end, assembly=None, opts=None):
        self.setup(kp_asm, addr_begin, assembly)
        self.addr_end = addr_end

        # create FillRange form
        super(Keypatch_FillRange, self).__init__(
            r"""STARTITEM {id:c_assembly}
BUTTON YES* Patch
KEYPATCH:: Fill Range

            {FormChangeCb}
            <Endian     :{c_endian}>
            <~S~yntax     :{c_syntax}>
            <Start      :{c_addr}>
            <End        :{c_addr_end}>
            <Size       :{c_size}>
            <~A~ssembly   :{c_assembly}>
             <##-   Fixup :{c_raw_assembly}>
             <##-   Encode:{c_encoding}>
             <##-   Size  :{c_encoding_len}>
            <~N~OPs padding until next instruction boundary:{c_opt_padding}>
            <Save ~o~riginal instructions in IDA comment:{c_opt_comment}>{c_opt_chk}>
            """, {
            'c_endian': self.DropdownListControl(
                          items = list(self.kp_asm.endian_lists.keys()),
                          readonly = True,
                          selval = self.endian_id),
            'c_addr': self.NumericInput(value=addr_begin, swidth=MAX_ADDRESS_LEN, tp=self.FT_ADDR),
            'c_addr_end': self.NumericInput(value=addr_end - 1, swidth=MAX_ADDRESS_LEN, tp=self.FT_ADDR),
            'c_assembly': self.StringInput(value=self.asm[:MAX_INSTRUCTION_STRLEN], width=MAX_INSTRUCTION_STRLEN),
            'c_size': self.NumericInput(value=addr_end - addr_begin, swidth=8, tp=self.FT_DEC),
            'c_raw_assembly': self.StringInput(value='', width=MAX_INSTRUCTION_STRLEN),
            'c_encoding': self.StringInput(value='', width=MAX_ENCODING_LEN),
            'c_encoding_len': self.NumericInput(value=0, swidth=8, tp=self.FT_DEC),
            'c_syntax': self.DropdownListControl(
                          items = self.syntax_keys,
                          readonly = True,
                          selval = self.syntax_id),
            'c_opt_chk':idaapi.Form.ChkGroupControl(('c_opt_padding', 'c_opt_comment', ''), value=opts['c_opt_chk']),
            'FormChangeCb': self.FormChangeCb(self.OnFormChange),
            })

        self.Compile()

    # callback to be executed when any form control changed
    def OnFormChange(self, fid):
        # make some controls read-only in FillRange mode
        self.EnableField(self.c_size, False)
        self.EnableField(self.c_addr_end, False)

        return self.update_patchform(fid)


# Patcher form
class Keypatch_Patcher(Keypatch_Form):
    def __init__(self, kp_asm, address, assembly=None, opts=None):
        self.setup(kp_asm, address, assembly)

        # create Patcher form
        super(Keypatch_Patcher, self).__init__(
            r"""STARTITEM {id:c_assembly}
BUTTON YES* Patch
KEYPATCH:: Patcher

            {FormChangeCb}
            <Endian     :{c_endian}>
            <~S~yntax     :{c_syntax}>
            <Address    :{c_addr}>
            <Original   :{c_orig_assembly}>
             <##-   Encode:{c_orig_encoding}>
             <##-   Size  :{c_orig_len}>
            <~A~ssembly   :{c_assembly}>
             <##-   Fixup :{c_raw_assembly}>
             <##-   Encode:{c_encoding}>
             <##-   Size  :{c_encoding_len}>
            <~N~OPs padding until next instruction boundary:{c_opt_padding}>
            <Save ~o~riginal instructions in IDA comment:{c_opt_comment}>{c_opt_chk}>
            """, {
            'c_endian': self.DropdownListControl(
                          items = list(self.kp_asm.endian_lists.keys()),
                          readonly = True,
                          selval = self.endian_id),
            'c_addr': self.NumericInput(value=address, swidth=MAX_ADDRESS_LEN, tp=self.FT_ADDR),
            'c_assembly': self.StringInput(value=self.asm[:MAX_INSTRUCTION_STRLEN], width=MAX_INSTRUCTION_STRLEN),
            'c_orig_assembly': self.StringInput(value=self.orig_asm[:MAX_INSTRUCTION_STRLEN], width=MAX_INSTRUCTION_STRLEN),
            'c_orig_encoding': self.StringInput(value=self.orig_encoding[:MAX_ENCODING_LEN], width=MAX_ENCODING_LEN),
            'c_orig_len': self.NumericInput(value=self.orig_len, swidth=8, tp=self.FT_DEC),
            'c_raw_assembly': self.StringInput(value='', width=MAX_INSTRUCTION_STRLEN),
            'c_encoding': self.StringInput(value='', width=MAX_ENCODING_LEN),
            'c_encoding_len': self.NumericInput(value=0, swidth=8, tp=self.FT_DEC),
            'c_syntax': self.DropdownListControl(
                          items = self.syntax_keys,
                          readonly = True,
                          selval = self.syntax_id),
            'c_opt_chk':self.ChkGroupControl(('c_opt_padding', 'c_opt_comment', ''), value=opts['c_opt_chk']),
            'FormChangeCb': self.FormChangeCb(self.OnFormChange),
            })

        self.Compile()

    # callback to be executed when any form control changed
    def OnFormChange(self, fid):
        # make some fields read-only in Patch mode
        self.EnableField(self.c_orig_assembly, False)
        self.EnableField(self.c_orig_encoding, False)
        self.EnableField(self.c_orig_len, False)

        return self.update_patchform(fid)


# Search position chooser
class SearchResultChooser(Choose):
    def __init__(self, title, items, flags=0, width=None, height=None, embedded=False, modal=False):
        super(SearchResultChooser, self).__init__(
            title,
            [["Address", Choose.CHCOL_HEX|40]],
            flags = flags,
            width = width,
            height = height,
            embedded = embedded)
        self.n = 0
        self.items = items
        self.selcount = 0
        self.modal = modal

    def OnClose(self):
        return

    def OnSelectLine(self, n):
        self.selcount += 1
        jumpto(self.items[n][0])

    def OnGetLine(self, n):
        res = self.items[n]
        res = [idc.atoa(res[0])]
        return res

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

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


# Search form
class Keypatch_Search(Keypatch_Form):
    def __init__(self, kp_asm, address, assembly=None):
        self.setup(kp_asm, address, assembly)

        # create Search form
        super(Keypatch_Search, self).__init__(
            r"""STARTITEM {id:c_assembly}
BUTTON YES* Search
KEYPATCH:: Search

            {FormChangeCb}
            <A~r~ch       :{c_arch}>
            <E~n~dian     :{c_endian}>
            <~S~yntax     :{c_syntax}>
            <A~d~dress    :{c_addr}>
            <~A~ssembly   :{c_assembly}>
             <##-   Fixup :{c_raw_assembly}>
             <##-   Encode:{c_encoding}>
             <##-   Size  :{c_encoding_len}>
            """, {
            'c_addr': self.NumericInput(value=address, swidth=MAX_ADDRESS_LEN, tp=self.FT_ADDR),
            'c_assembly': self.StringInput(value=self.asm[:MAX_INSTRUCTION_STRLEN], width=MAX_INSTRUCTION_STRLEN),
            'c_raw_assembly': self.StringInput(value='', width=MAX_INSTRUCTION_STRLEN),
            'c_encoding': self.StringInput(value='', width=MAX_ENCODING_LEN),
            'c_encoding_len': self.NumericInput(value=0, swidth=8, tp=self.FT_DEC),
            'c_arch': self.DropdownListControl(
                          items = self.arch_keys,
                          readonly = True,
                          selval = self.arch_id,
                          width = 32),
            'c_endian': self.DropdownListControl(
                          items = list(self.kp_asm.endian_lists.keys()),
                          readonly = True,
                          selval = self.endian_id),
            'c_syntax': self.DropdownListControl(
                          items = self.syntax_keys,
                          readonly = True,
                          selval = self.syntax_id),
            'FormChangeCb': self.FormChangeCb(self.OnFormChange),
            })

        self.Compile()

    ################################ IDA 6/7 Compatibility method #######################################
    def find_binary(self, ea):
        if idaapi.IDA_SDK_VERSION >= 700:
            return ida_search.find_binary(ea, idaapi.cvar.inf.max_ea, self.GetControlValue(self.c_encoding), 16, idc.SEARCH_DOWN)
        return idc.FindBinary(ea, idc.SEARCH_DOWN, self.GetControlValue(self.c_encoding))
    #####################################################################################################

    # callback to be executed when any form control changed
    def OnFormChange(self, fid):
        # handle the search button
        if fid == -2:
            address = 0
            addresses = []
            while address != idc.BADADDR:
                address = self.find_binary(address)
                if address == idc.BADADDR:
                    break
                addresses.append([address])
                address = address + 1
            c = SearchResultChooser("Searching for [{0}]".format(self.GetControlValue(self.c_raw_assembly)), addresses)
            c.show()
            return 1

        # only Search mode allows to select arch+mode
        arch_id = self.GetControlValue(self.c_arch)
        (arch, mode) = self.kp_asm.get_arch_by_idx(arch_id)

        # assembly is focused
        self.SetFocusedField(self.c_assembly)

        if arch == KS_ARCH_X86:
            # enable Syntax and disable Endian for x86
            self.ShowField(self.c_syntax, True)
            self.EnableField(self.c_syntax, True)
            self.syntax_id = self.GetControlValue(self.c_syntax)
            self.EnableField(self.c_endian, False)
            # set Endian index properly
            self.SetControlValue(self.c_endian, 0)
        elif arch in (KS_ARCH_ARM64, KS_ARCH_HEXAGON, KS_ARCH_SYSTEMZ):
            # no Syntax & Endian option for these archs
            self.ShowField(self.c_syntax, False)
            self.EnableField(self.c_syntax, False)
            self.EnableField(self.c_endian, False)
            # set Endian index properly
            self.SetControlValue(self.c_endian, (mode & KS_MODE_BIG_ENDIAN != 0))
        elif (arch == KS_ARCH_PPC) and (mode & KS_MODE_PPC32 != 0):
            # no Syntax & Endian option for these archs
            self.ShowField(self.c_syntax, False)
            self.EnableField(self.c_syntax, False)
            self.EnableField(self.c_endian, False)
            # set Endian index properly
            self.SetControlValue(self.c_endian, (mode & KS_MODE_BIG_ENDIAN != 0))
        else:
            # no Syntax & Endian option
            self.ShowField(self.c_syntax, False)
            self.EnableField(self.c_syntax, False)
            self.EnableField(self.c_endian, True)

        if self.GetControlValue(self.c_endian) == 1:
            endian = KS_MODE_BIG_ENDIAN
        else:
            endian = KS_MODE_LITTLE_ENDIAN

        # update other controls & Encoding with live assembling
        self.update_controls(arch, mode)

        return 1


# About form
class About_Form(idaapi.Form):
    def __init__(self, version):
        # create About form
        super(About_Form, self).__init__(
            r"""STARTITEM 0
BUTTON YES* Open Keypatch Website
KEYPATCH:: About

            {FormChangeCb}
            Keypatch IDA plugin v%s, using Keystone Engine v%s.
            (c) Nguyen Anh Quynh + Thanh Nguyen, 2018.

            Keypatch is released under the GPL v2.
            Find more info at http://www.keystone-engine.org/keypatch
            """ %(version, keystone.__version__), {
            '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(KP_HOMEPAGE, new = 2)

        return 1


# Check-for-update form
class Update_Form(idaapi.Form):
    def __init__(self, version, message):
        # create Update form
        super(Update_Form, self).__init__(
            r"""STARTITEM 0
BUTTON YES* Open Keypatch Website
KEYPATCH:: Check for update

            {FormChangeCb}
            Your Keypatch 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(KP_HOMEPAGE, new = 2)

        return 1


try:
    # adapted from pull request #7 by @quangnh89
    class Kp_Menu_Context(idaapi.action_handler_t):

        @classmethod
        def get_name(self):
            return self.__name__

        @classmethod
        def get_label(self):
            return self.label

        @classmethod
        def register(self, plugin, label):
            self.plugin = plugin
            self.label = label
            instance = self()
            return idaapi.register_action(idaapi.action_desc_t(
                self.get_name(),  # Name. Acts as an ID. Must be unique.
                instance.get_label(),  # Label. That's what users see.
                instance  # Handler. Called when activated, and for updating
            ))

        @classmethod
        def unregister(self):
            """Unregister the action.
            After unregistering the class cannot be used.
            """
            idaapi.unregister_action(self.get_name())

        @classmethod
        def activate(self, ctx):
            # dummy method
            return 1

        @classmethod
        def update(self, ctx):
            try:
                if ctx.widget_type == idaapi.BWN_DISASM:
                    return idaapi.AST_ENABLE_FOR_FORM
                else:
                    return idaapi.AST_DISABLE_FOR_FORM
            except Exception as e:
                # Add exception for main menu on >= IDA 7.0
                return idaapi.AST_ENABLE_ALWAYS
            
    # context menu for Patcher
    class Kp_MC_Patcher(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.patcher()
            return 1

    # context menu for Fill Range
    class Kp_MC_Fill_Range(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.fill_range()
            return 1

    # context menu for Undo
    class Kp_MC_Undo(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.undo()
            return 1

    # context menu for Search
    class Kp_MC_Search(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.search()
            return 1

    # context menu for Check Update
    class Kp_MC_Updater(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.updater()
            return 1

    # context menu for About
    class Kp_MC_About(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.about()
            return 1

except:
    pass

# hooks for popup menu
class Hooks(idaapi.UI_Hooks):
    if idaapi.IDA_SDK_VERSION >= 700:
        # IDA >= 700 right click widget popup
        def finish_populating_widget_popup(self, form, popup):
            if idaapi.get_widget_type(form) == idaapi.BWN_DISASM:
                try:
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Patcher.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Fill_Range.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Undo.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, "-", 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Search.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, "-", 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Updater.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_About.get_name(), 'Keypatch/')
                except:
                    pass
    else:
        # IDA < 700 right click popup
        def finish_populating_tform_popup(self, form, popup):
            # We'll add our action to all "IDA View-*"s.
            # If we wanted to add it only to "IDA View-A", we could
            # also discriminate on the widget's title:
            #
            #  if idaapi.get_tform_title(form) == "IDA View-A":
            #      ...
            #
            if idaapi.get_tform_type(form) == idaapi.BWN_DISASM:
                try:
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Patcher.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Fill_Range.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Undo.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, "-", 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Search.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, "-", 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_Updater.get_name(), 'Keypatch/')
                    idaapi.attach_action_to_popup(form, popup, Kp_MC_About.get_name(), 'Keypatch/')
                except:
                    pass


# check if we already initialized Keypatch
kp_initialized = False

#--------------------------------------------------------------------------
# Plugin
#--------------------------------------------------------------------------
class Keypatch_Plugin_t(idaapi.plugin_t):
    comment = "Keypatch plugin for IDA Pro (using Keystone framework)"
    help = "Find more information on Keypatch at http://keystone-engine.org/keypatch"
    wanted_name = "Keypatch Patcher"
    wanted_hotkey = "Ctrl-Alt-K"
    flags = idaapi.PLUGIN_KEEP

    def load_configuration(self):
        # default
        self.opts = {}

        # load configuration from file
        try:
            f = open(KP_CFGFILE, "rt")
            self.opts = json.load(f)
            f.close()
        except IOError:
            print("Keypatch: FAILED to load config file. Use default setup now.")
        except Exception as e:
            print("Keypatch: FAILED to load config file, with exception: {0}".format(str(e)))

        # use default values if not defined in config file
        if 'c_opt_padding' not in self.opts:
            self.opts['c_opt_padding'] = 1

        if 'c_opt_comment' not in self.opts:
            self.opts['c_opt_comment'] = 2

        self.opts['c_opt_chk'] = self.opts['c_opt_padding'] | self.opts['c_opt_comment']

    def init(self):
        global kp_initialized

        # register popup menu handlers
        try:
            Kp_MC_Patcher.register(self, "Patcher    (Ctrl-Alt-K)")
            Kp_MC_Fill_Range.register(self, "Fill Range")
            Kp_MC_Undo.register(self, "Undo last patching")
            Kp_MC_Search.register(self, "Search")
            Kp_MC_Updater.register(self, "Check for update")
            Kp_MC_About.register(self, "About")
        except:
            pass

        # setup popup menu
        self.hooks = Hooks()
        self.hooks.hook()

        self.opts = None
        if kp_initialized == False:
            kp_initialized = True

            if idaapi.IDA_SDK_VERSION >= 700:
                # Add menu IDA >= 7.0
                idaapi.attach_action_to_menu("Edit/Keypatch/Patcher", Kp_MC_Patcher.get_name(), idaapi.SETMENU_APP)
                idaapi.attach_action_to_menu("Edit/Keypatch/About", Kp_MC_About.get_name(), idaapi.SETMENU_APP)
                idaapi.attach_action_to_menu("Edit/Keypatch/Check for update", Kp_MC_Updater.get_name(), idaapi.SETMENU_APP)
                idaapi.attach_action_to_menu("Edit/Keypatch/Search", Kp_MC_Search.get_name(), idaapi.SETMENU_APP)
                idaapi.attach_action_to_menu("Edit/Keypatch/Undo last patching", Kp_MC_Undo.get_name(), idaapi.SETMENU_APP)
                idaapi.attach_action_to_menu("Edit/Keypatch/Fill Range", Kp_MC_Fill_Range.get_name(), idaapi.SETMENU_APP)
            else:
                # add Keypatch menu
                menu = idaapi.add_menu_item("Edit/Keypatch/", "Patcher     (Ctrl-Alt-K)", "", 1, self.patcher, None)
                if menu is not None:
                    idaapi.add_menu_item("Edit/Keypatch/", "About", "", 1, self.about, None)
                    idaapi.add_menu_item("Edit/Keypatch/", "Check for update", "", 1, self.updater, None)
                    idaapi.add_menu_item("Edit/Keypatch/", "-", "", 1, self.menu_null, None)
                    idaapi.add_menu_item("Edit/Keypatch/", "Search", "", 1, self.search, None)
                    idaapi.add_menu_item("Edit/Keypatch/", "-", "", 1, self.menu_null, None)
                    idaapi.add_menu_item("Edit/Keypatch/", "Undo last patching", "", 1, self.undo, None)
                    idaapi.add_menu_item("Edit/Keypatch/", "Fill Range", "", 1, self.fill_range, None)
                elif idaapi.IDA_SDK_VERSION < 680:
                    # older IDAPython (such as in IDAPro 6.6) does add new submenu.
                    # in this case, put Keypatch menu in menu Edit \ Patch program
                    # not sure about v6.7, so to be safe we just check against v6.8
                    idaapi.add_menu_item("Edit/Patch program/", "-", "", 0, self.menu_null, None)
                    idaapi.add_menu_item("Edit/Patch program/", "Keypatch:: About", "", 0, self.about, None)
                    idaapi.add_menu_item("Edit/Patch program/", "Keypatch:: Check for update", "", 0, self.updater, None)
                    idaapi.add_menu_item("Edit/Patch program/", "Keypatch:: Search", "", 0, self.search, None)
                    idaapi.add_menu_item("Edit/Patch program/", "Keypatch:: Undo last patching", "", 0, self.undo, None)
                    idaapi.add_menu_item("Edit/Patch program/", "Keypatch:: Fill Range", "", 0, self.fill_range, None)
                    idaapi.add_menu_item("Edit/Patch program/", "Keypatch:: Patcher     (Ctrl-Alt-K)", "", 0, self.patcher, None)

            print("=" * 80)
            print("Keypatch v{0} (c) Nguyen Anh Quynh & Thanh Nguyen, 2018".format(VERSION))
            print("Keypatch is using Keystone v{0}".format(keystone.__version__))
            print("Keypatch Patcher's shortcut key is Ctrl-Alt-K")
            print("Use the same hotkey Ctrl-Alt-K to open 'Fill Range' window on a selected range of code")
            print("To revert (undo) the last patching, choose menu Edit | Keypatch | Undo last patching")
            print("Keypatch Search is available from menu Edit | Keypatch | Search")
            print("Find more information about Keypatch at http://keystone-engine.org/keypatch")

            self.load_configuration()

            print("=" * 80)
            self.kp_asm = Keypatch_Asm()

        return idaapi.PLUGIN_KEEP

    def term(self):
        if self.hooks is not None:
            self.hooks.unhook()
            self.hooks = None

        if self.opts is None:
            return
        #save configuration to file
        try:
            json.dump(self.opts, open(KP_CFGFILE, "wt"))
        except Exception as e:
            print("Keypatch: FAILED to save config file, with exception: {0}".format(str(e)))
        else:
            print("Keypatch: configuration is saved to {0}".format(KP_CFGFILE))

    # null handler
    def menu_null(self):
        pass

    # handler for About menu
    def about(self):
        f = About_Form(VERSION)
        f.Execute()
        f.Free()

    # handler for Check-for-Update menu
    def updater(self):
        (r, content) = url_download(KP_GITHUB_VERSION)
        content = to_string(content)
        if r == 0:
            # find stable version
            sig = 'VERSION_STABLE = "'
            tmp = content[content.find(sig)+len(sig):]
            version_stable = tmp[:tmp.find('"')]

            # compare with the current version
            if version_stable == VERSION:
                f = Update_Form(VERSION, "Good, you are already on the latest stable version!")
                f.Execute()
                # free this form
                f.Free()
            else:
                f = Update_Form(VERSION, "Download latest stable version {0} from http://keystone-engine.org/keypatch".format(version_stable))
                f.Execute()
                # free this form
                f.Free()
        else:
            # fail to download
            warning("ERROR: Keypatch failed to connect to internet (Github). Try again later.")
            print("Keypatch: FAILED to connect to Github to check for latest update. Try again later.")

    # handler for Undo menu
    def undo(self):
        global patch_info
        if len(patch_info) == 0:
            # TODO: disable Undo menu?
            warning("ERROR: Keypatch already got to the last undo patching!")
        else:
            (address, assembly, p_orig_data, patch_comment) = patch_info[-1]

            # undo the patch
            self.kp_asm.patch_code(address, None, None, None, None, orig_asm=[assembly], patch_data=p_orig_data, patch_comment=patch_comment, undo=True)
            del(patch_info[-1])

    # handler for Search menu
    def search(self):
        address = get_screen_ea()
        f = Keypatch_Search(self.kp_asm, address)
        f.Execute()
        f.Free()

    # handler for Patcher menu
    def patcher(self):
        # be sure that this arch is supported by Keystone
        if self.kp_asm.arch is None:
            warning("ERROR: Keypatch cannot handle this architecture (unsupported by Keystone), quit!")
            return

        selection, addr_begin, addr_end = read_range_selection()
        if selection:
            # call Fill Range function on this selected code
            return self.fill_range()

        address = get_screen_ea()

        if self.opts is None:
            self.load_configuration()

        init_assembly = None
        while True:
            f = Keypatch_Patcher(self.kp_asm, address, assembly=init_assembly, opts=self.opts)
            ok = f.Execute()
            if ok == 1:
                try:
                    syntax = None
                    if f.kp_asm.arch == KS_ARCH_X86:
                        syntax_id = f.c_syntax.value
                        syntax = self.kp_asm.get_syntax_by_idx(syntax_id)

                    assembly = f.c_assembly.value
                    self.opts = f.get_opts('c_opt_chk')
                    padding = (self.opts.get("c_opt_padding", 0) != 0)
                    comment = (self.opts.get("c_opt_comment", 0) != 0)

                    raw_assembly = self.kp_asm.ida_resolve(assembly, address)

                    print("Keypatch: attempting to modify \"{0}\" at 0x{1:X} to \"{2}\"".format(
                            self.kp_asm.ida_get_disasm(address), address, assembly))

                    length = self.kp_asm.patch_code(address, raw_assembly, syntax, padding, comment, None)
                    if length > 0:
                        # update start address pointing to the next instruction
                        init_assembly = None
                        address += length
                    else:
                        init_assembly = f.c_assembly.value
                        if length == 0:
                            warning("ERROR: Keypatch found invalid assembly [{0}]".format(assembly))
                        elif length == -1:
                            warning("ERROR: Keypatch failed to patch binary at 0x{0:X}!".format(address))
                        elif length == -2:
                            warning("ERROR: Keypatch can't read original data at 0x{0:X}, try again".format(address))

                except KsError as e:
                    print("Keypatch Error: {0}".format(e))
            else:   # Cancel
                break
            f.Free()

    # handler for Fill Range menu
    def fill_range(self):
        # be sure that this arch is supported by Keystone
        if self.kp_asm.arch is None:
            warning("ERROR: Keypatch cannot handle this architecture (unsupported by Keystone), quit!")
            return
               
        selection, addr_begin, addr_end = read_range_selection()
        if not selection:
            warning("ERROR: Keypatch requires a range to be selected for fill in, try again")
            return

        if self.opts is None:
            self.load_configuration()

        init_assembly = None
        f = Keypatch_FillRange(self.kp_asm, addr_begin, addr_end, assembly=init_assembly, opts=self.opts)
        ok = f.Execute()
        if ok == 1:
            try:
                syntax = None
                if f.kp_asm.arch == KS_ARCH_X86:
                    syntax_id = f.c_syntax.value
                    syntax = self.kp_asm.get_syntax_by_idx(syntax_id)

                assembly = f.c_assembly.value
                self.opts = f.get_opts('c_opt_chk')
                padding = (self.opts.get("c_opt_padding", 0) != 0)
                comment = (self.opts.get("c_opt_comment", 0) != 0)

                raw_assembly = self.kp_asm.ida_resolve(assembly, addr_begin)

                print("Keypatch: attempting to fill range [0x{0:X}:0x{1:X}] with \"{2}\"".format(
                    addr_begin, addr_end - 1, assembly))

                length = self.kp_asm.fill_code(addr_begin, addr_end, raw_assembly, syntax, padding, comment, None)
                if length == 0:
                    warning("ERROR: Keypatch failed to process this input.")
                    print("Keypatch: FAILED to process this input '{0}'".format(assembly))
                elif length == -1:
                    warning("ERROR: Keypatch failed to patch binary at 0x{0:X}!".format(addr_begin))

            except KsError as e:
                print("Keypatch Error: {0}".format(e))

        # free this form
        f.Free()

    def run(self, arg):
        self.patcher()


# register IDA plugin
def PLUGIN_ENTRY():
    return Keypatch_Plugin_t()


repository: evolution536/ida-heuristics
file: heuristics.py
description: IDA plugin for architecture-agnostic heuristics
wanted_name: Heuristics
comment: Implements heuristics for obfuscated code by Tim Blazytko (from REcon)
created_at: 2024-07-09 10:46:20+00:00
pushed_at: 2024-07-26 21:40:56+00:00
url:
https://github.com/evolution536/ida-heuristics/blob/508269595b637bbf3212899a165cec51f30b201b/heuristics.py

content:
import os
import re
import idc
import idaapi
import ida_diskio
import idautils
import ida_bytes

# Global definitions.
VERSION = "1.0"

# Set initialized flag to false to begin with.
p_initialized = False

# Returns the index of the last element in the result list that falls within the 95th percentile or higher.
def get_95th_percentile(results):
	threshold = results[3] * 0.05
	slider = 0
	num = 0
	while slider < threshold:
		num += 1
		slider += results[0][num][1][1]
	return num

# Counts the number of instructions inside a function, taking into account that it might have multiple chunks.
def count_function_instructions(func_ea):
	# Keep track of total number of instructions in function.
	insn_count = 0

	# Get chunks in this function.
	for (start_ea, end_ea) in idautils.Chunks(func_ea):
		if start_ea == end_ea:
			continue

		# Increment instruction counter and move to next head.
		head = idc.next_head(start_ea)
		chunk_insn_count = 0
		while head < end_ea:
			chunk_insn_count += 1
			head = idc.next_head(head)

		# When done, add the instruction count inside this chunk to the total counter.
		insn_count += chunk_insn_count

	# Return total instruction count for the entire function.
	return insn_count

# Computes large basic blocks score for given function.
def func_large_basic_blocks(func_ea):
	# Get object reference to function.
	func_ptr = idaapi.get_func(func_ea)

	# Get flow chart of the function containing all basic blocks.
	# https://moritzraabe.de/2017/01/15/ida-pro-anti-disassembly-basic-blocks-and-idapython/
	flowchart = idaapi.FlowChart(func_ptr)

	# Get number of instructions inside the function.
	insn_count = count_function_instructions(func_ea)

	# Compute number of instructions divided by the number of basic blocks.
	return insn_count / max(1, flowchart.size)

# Scoring function that computes the cyclomatic complexity metric of given function.
def func_cyclomatic_complexity(func_ea):
	# Get object reference to function.
	func_ptr = idaapi.get_func(func_ea)

	# Get flow chart of the function containing all basic blocks.
	# https://moritzraabe.de/2017/01/15/ida-pro-anti-disassembly-basic-blocks-and-idapython/
	flowchart = idaapi.FlowChart(func_ptr, flags=idaapi.FC_PREDS | idaapi.FC_NOEXT)

	# Calculate number of edges inside the function by looking at predecessors and successors.
	edges = 0
	for block in flowchart:
		for succ in block.succs():
			edges += 1
		for pred in block.preds():
			edges += 1

	# Compute the cyclomatic complexity for this function.
	# We set the minimum value to 1, because in some heavily obfuscated binary files,
	# functions are not always parsed correctly and a couple of not-connected basic blocks
	# may reside in the function. In such a case, we would get a negative value that
	# screws up the statistics. Those are outliers and we will remove them this way.
	return max(1, edges - flowchart.size + 2)

# Calculates the numberof cross references to a given function.
def func_get_xrefs(func_ea):
	return len(list(idautils.XrefsTo(func_ea)))

# Gets tree of blocks dominated by a given basic block.
def get_blocks_dominated_by(block):
	result = set()
	worklist = [ block ]

	# Get successing blocks. We perform a depth-first search to get the entire tree.
	while worklist:
		next_block = worklist.pop(0)
		result.add(next_block.start_ea)

		# Get all successing blocks and add them to the worklist.
		for succ in next_block.succs():
			# Test whether we have already visited this block. If so, we will not add it to the results.
			if succ.start_ea in result:
				continue

			# Not visited yet, so add the successor to the worklist.
			worklist.append(succ)

	# Return the successors, or the dominated basic blocks.
	return result

# Computes control flow flattening score for given function.
def func_control_flow_flattening(func_ea):
	# Initialize score to zero.
	score = 0.0

	# Get object reference to function.
	func_ptr = idaapi.get_func(func_ea)

	# Get flow chart of the function containing all basic blocks.
	# https://moritzraabe.de/2017/01/15/ida-pro-anti-disassembly-basic-blocks-and-idapython/
	flowchart = idaapi.FlowChart(func_ptr, flags=idaapi.FC_PREDS | idaapi.FC_NOEXT)

	# Iterate over all blocks in this function and compute the score.
	for block in flowchart:
		# Get blocks dominated by the current one.
		dominated = get_blocks_dominated_by(block)

		# Test for a back edge in the tree of dominated blocks. That means, if there exists an edge
		# between any of the successing blocks back to the currently examined block, there is a cycle
		# and, assuming that the current block is the 'controller', the dominated block flows back there.
		if not any(b.start_ea in dominated and b.start_ea != block.start_ea for b in block.preds()):
			# None found, so this control flow path has a low probability of being flattened.
			continue

		# Compute the score with this block.
		score = max(score, len(dominated) / flowchart.size)

	# Return the accumulated score.
	return score

# Get dictionary of stack members of specified function.
def get_stack_variables(func_ea):
	members = {}
	base = None

	# Get stack frame.
	frame = idc.get_func_attr(func_ea, idc.FUNCATTR_FRAME)

	# Ignore functions that (don't have a stack frame at all I guess?).
	if frame == idc.BADADDR:
		return members, base

	# Iterate frame structure members and save them in the output.
	for frame_member in idautils.StructMembers(frame):
		member_offset, member_name, _ = frame_member
		members[member_name] = member_offset
		if member_name == ' r':
			base = member_offset
	return members, base

# Gets list of cross references to stack variables in function.
def get_stack_references(func_ea, members, stack_base):
	refs = {}
	for func_item in idautils.FuncItems(func_ea):
		flags = ida_bytes.get_full_flags(func_item)
		stkvar = 0 if ida_bytes.is_stkvar0(flags) else 1 if ida_bytes.is_stkvar1(flags) else None
		if stkvar is None:
			continue

		# Get operand that references the stack location as a string. We will split locations ourselves.
		opnd = idc.print_operand(func_item, stkvar).replace('[', '').replace(']', '')

		# Split the sub-operands from the memory locator.
		m = re.split(r'[+*-]', opnd)
		for x in m:
			# Test whether one of the fake offsets reside in this instruction.
			if x in members:
				# Add it to the dictionary if not already there.
				if x not in refs:
					refs[x] = [ func_item ]
				else:
					refs[x].append(func_item)

	return refs

# Computes wrong stack frame score for given function.
def func_wrong_stack_frame(func_ea):
	# Initialize score to zero.
	score = 0.0

	# Get stack members first.
	members, base = get_stack_variables(func_ea)

	# Get cross references to those stack members.
	total_refs = 0
	refs = get_stack_references(func_ea, members, base)
	for v, r in refs.items():
		if len(r) == 0:
			continue
		else: total_refs += len(r)

	# FIXME: see which type of metric would do best here.
	return len(members) / max(1, total_refs)

# Iterates over segments and functions, executing a scoring function and tracking the metrics.
def iterate_functions_and_track_statistics(proc):
	# Create dictionary that will keep track of functions and metrics.
	metric_dict = {}
	sum_of_metrics = 0
	min_metric = 0xFFFFFFFF
	max_metric = 0

	# Get all segments.
	seg_ea = idc.get_first_seg()
	while seg_ea != idc.BADADDR:
		# Get all functions within this segment.
		for func_ea in idautils.Functions(seg_ea, idc.get_segm_end(seg_ea)):
			# Get function name.
			func_name = idc.get_func_name(func_ea)

			# Compute metric and store it in the dictionary.
			metric = proc(func_ea)
			metric_dict[func_ea] = (func_name, metric)

			# Keep track of statistical values along the way.
			sum_of_metrics += metric
			min_metric = min(min_metric, metric)
			max_metric = max(max_metric, metric)

		# Get next segment ea.
		seg_ea = idc.get_next_seg(seg_ea)

	# Sort the metric dictionary by highest metric and return this list.
	return (sorted(metric_dict.items(), key=lambda item: item[1][1], reverse=True), min_metric, max_metric, sum_of_metrics)

# The callback executed to compute cyclomatic complexity.
def cyclomatic_complexity():
	# Compute the cyclomatic complexity over all functions in the database and return a sorted list of results.
	results = iterate_functions_and_track_statistics(func_cyclomatic_complexity)

	# Get 95th percentile of result data.
	num = get_95th_percentile(results)

	# Print the results.
	print("Total number of functions: %i. Lowest metric: %f, highest metric: %f." % (len(results[0]), results[1], results[2]))
	print("Top 5%% of functions with highest cyclomatic complexity")
	for i in range(num):
		entry = results[0][i]
		print("Function: 0x%X (%s) has cyclomatic complexity: %f" % (entry[0], entry[1][0], entry[1][1]))
	print("Highest 5%% covers %i out of %i functions." % (num, len(results[0])))

# The callback executed to compute large basic blocks.
def large_basic_blocks():
	# Compute the large basic blocks score over all functions in the database and return a sorted list of results.
	results = iterate_functions_and_track_statistics(func_large_basic_blocks)

	# Get 95th percentile of result data.
	num = get_95th_percentile(results)

	# Print the results.
	print("Total number of functions: %i. Lowest metric: %f, highest metric: %f." % (len(results[0]), results[1], results[2]))
	print("Top 5%% of functions with largest basic blocks.")
	for i in range(num):
		entry = results[0][i]
		print("Function: 0x%X (%s) has large basic block score: %f" % (entry[0], entry[1][0], entry[1][1]))
	print("Highest 5%% covers %i out of %i functions." % (num, len(results[0])))

# Gets overview of most frequently called functions in the database.
def most_frequently_called_functions():
	# Compute the number of cross references to all functions in the database and return a sorted list of results.
	results = iterate_functions_and_track_statistics(func_get_xrefs)

	# Get 95th percentile of result data.
	num = get_95th_percentile(results)

	# Print the results.
	print("Total number of functions: %i. Lowest metric: %f, highest metric: %f." % (len(results[0]), results[1], results[2]))
	print("Top 5%% of most frequently called functions.")
	for i in range(num):
		entry = results[0][i]
		print("Function: 0x%X (%s) has %f calling locations." % (entry[0], entry[1][0], entry[1][1]))
	print("Highest 5%% covers %i out of %i functions." % (num, len(results[0])))

# Computes control flow flattening score for functions.
def control_flow_flattening():
	# Compute the control flow flattening score over all functions in the database and return a sorted list of results.
	results = iterate_functions_and_track_statistics(func_control_flow_flattening)

	# Get 95th percentile of result data.
	num = get_95th_percentile(results)

	# Print the results.
	print("Total number of functions: %i. Lowest metric: %f, highest metric: %f." % (len(results[0]), results[1], results[2]))
	print("Top 5%% of highest control flow flattening scores.")
	for i in range(num):
		entry = results[0][i]
		print("Function: 0x%X (%s) has flattening score: %f." % (entry[0], entry[1][0], entry[1][1]))
	print("Highest 5%% covers %i out of %i functions." % (num, len(results[0])))

# Computes wrong stack frame score for functions.
def wrong_stack_frame():
	# Compute the wrong stack frame score over all functions in the database and return a sorted list of results.
	results = iterate_functions_and_track_statistics(func_wrong_stack_frame)

	# Get 95th percentile of result data.
	num = get_95th_percentile(results)

	# Print the results.
	print("Total number of functions: %i. Lowest metric: %f, highest metric: %f." % (len(results[0]), results[1], results[2]))
	print("Top 5%% of highest wrong stack frame scores.")
	for i in range(num):
		entry = results[0][i]
		print("Function: 0x%X (%s) has wrong stack frame score: %f." % (entry[0], entry[1][0], entry[1][1]))
	print("Highest 5%% covers %i out of %i functions." % (num, len(results[0])))

# Define the action_handler_t object that fires the callback function when each action is activated.
class ActionHandler(idaapi.action_handler_t):
	def __init__(self, callback):
		idaapi.action_handler_t.__init__(self)
		self.callback = callback

	def activate(self, ctx):
		self.callback()
		return 1

	def update(self, ctx):
		return idaapi.AST_ENABLE_ALWAYS

# Registers actions to IDA Pro.
def register_actions():
	actions = [
		{
			'id': 'heuristics:large_basic_blocks:',
			'name': 'Large Basic Blocks',
			'hotkey': 'Ctrl+Alt+B',
			'comment': 'Shows functions with large basic blocks',
			'callback': large_basic_blocks,
			'menu_location': 'Edit/Heuristics/Large Basic Blocks'
		},
		{
			'id': 'heuristics:cyclomatic_complexity',
			'name': 'Cyclomatic Complexity',
			'hotkey': 'Ctrl+Alt+C',
			'comment': 'Shows functions with top cyclomatic complexity',
			'callback': cyclomatic_complexity,
			'menu_location': 'Edit/Heuristics/Cyclomatic Complexity'
		},
		{
			'id' : 'heuristics:frequently_called_funcs',
			'name' : 'Frequenly Called Functions',
			'hotkey' : 'Ctrl+Alt+D',
			'comment' : 'Shows which functions are most frequently called',
			'callback' : most_frequently_called_functions,
			'menu_location' : 'Edit/Heuristics/Frequently Called Functions'
		},
		{
			'id' : 'heuristics:control_flow_flattening',
			'name' : 'Control Flow Flattening',
			'hotkey' : 'Ctrl+Alt+E',
			'comment' : 'Identifies functions with a high level of control flow flattening',
			'callback' : control_flow_flattening,
			'menu_location' : 'Edit/Heuristics/Control Flow Flattening'
		},
		{
			'id' : 'heuristics:wrong_stack_frame',
			'name' :' Wrong Stack Frame',
			'hotkey' : 'Ctrl+Alt+F',
			'comment' : 'Identifies functions with fake stack variables',
			'callback' : wrong_stack_frame,
			'menu_location' : 'Edit/Heuristics/Wrong Stack Frame'
		}
	]

	# Register the specified actions.
	for action in actions:
		if not idaapi.register_action(idaapi.action_desc_t(
			action['id'], action['name'],
			ActionHandler(action['callback']),
			action['hotkey'], action['comment'])):
				print('Failed to register ' + action['id'])

		# Attach registered action to menu bar.
		if not idaapi.attach_action_to_menu(action['menu_location'], action['id'], 0):
			print('Failed to attach to menu '+ action['id'])

class HeuristicsPlugin(idaapi.plugin_t):
	flags = idaapi.PLUGIN_HIDE
	comment = "Implements heuristics for obfuscated code by Tim Blazytko (from REcon)"
	help = "Do you even reverse bro?"
	wanted_name = "Heuristics"
	wanted_hotkey = ''

	def get_user_directory(self):
		user_dir = ida_diskio.get_user_idadir()
		plug_dir = os.path.join(user_dir, "plugins")
		res_dir = os.path.join(plug_dir, "heuristics")
		if not os.path.exists(res_dir):
			os.makedirs(res_dir, 0o755)
		return res_dir

	def init(self):
		global p_initialized

		# Register popup menu handlers.
		register_actions()

		# Set initialized flag to true.
		if p_initialized is False:
			p_initialized = True

		# Get user directory, so we can place files there if needed.
		self.user_directory = self.get_user_directory()

		# Print welcome message with some useful information.
		print("=" * 80)
		print("= Heuristics plugin v%s to detect obfuscation by Gijs Rijnders (evolution536)" % VERSION)
		print("= Implements heuristics as described by Tim Blazytko (REcon)")
		print("= !!! Please wait for auto-analysis to finish before running the script !!!")
		print("=" * 80)

		# Keep the plugin loaded, it might be called from the menu later.
		return idaapi.PLUGIN_KEEP

	def run(self, arg):
		return

	def term(self):
		pass

def PLUGIN_ENTRY():
	return HeuristicsPlugin()


repository: HexRabbit/diaphora_loader
file: diaphora_loader.py
description: A loader for IDA pro plugin diaphora
wanted_name: diaphora
comment: diaphora - powerful bindiff for IDA pro
created_at: 2019-08-16 06:19:24+00:00
pushed_at: 2019-08-16 06:28:10+00:00
url:
https://github.com/HexRabbit/diaphora_loader/blob/2a9cd812ec8669c40d6af99b3d929d3d34209538/diaphora_loader.py

content:
# -*- coding: utf-8 -*
from __future__ import print_function
import idaapi
import os

# IDA/
# └── plugins/
#     ├── diaphora_loader.py
#     └── diaphora/
#         └── diaphora.py

class diaphoraPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_FIX
    comment = 'diaphora - powerful bindiff for IDA pro'
    help = 'diaphora Plugin'
    wanted_name = 'diaphora'
    wanted_hotkey = ''
    dialog = None

    def init(self):
        action = idaapi.action_desc_t(
            'diaphora:proxy',             # The action name. This acts like an ID and must be unique
            'Run diaphora',               # The action text.
            runHandler(),                 # The action handler.
            None,                         # Optional: the action shortcut
            'Run diaphora'                # Optional: the action tooltip (available in menus/toolbar)
        )
        idaapi.register_action(action)
        idaapi.attach_action_to_menu(
            'File',                       # The relative path of where to add the action
            'diaphora:proxy',             # The action ID
            idaapi.SETMENU_APP)
        print("Diaphora :: Plugin loaded")
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        diaphora = runHandler()
        diaphora.activate(self)

    def term(self):
        pass

class runHandler(idaapi.action_handler_t):
    def __init__(self):
        idaapi.action_handler_t.__init__(self)

    def activate(self, ctx):
        g = globals()
        script = os.path.join(idaapi.idadir("plugins"), "diaphora", "diaphora.py")
        idaapi.IDAPython_ExecScript(script, g)

    def update(self, ctx):
        return idaapi.AST_ENABLE_ALWAYS

def PLUGIN_ENTRY():
    return diaphoraPlugin()


repository: mandiant/xrefer
file: plugins/xrefer.py
description: FLARE Team's Binary Navigator
url: https://github.com/mandiant/xrefer/blob/26f2400c8199350b31f89cc5e47d8b1cd2f1bb2e/plugins/xrefer.py
created_at: 2024-10-22 09:04:50+00:00
pushed_at: 2025-04-06 06:46:52+00:00
content:
# Copyright 2024 Google LLC
#
# 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.

import os
import sys

plugin_dir = os.path.dirname(os.path.realpath(__file__))
sys.path.append(os.path.join(plugin_dir, 'xrefer'))
sys.path.append(os.path.join(plugin_dir, 'xrefer', 'core'))
sys.path.append(os.path.join(plugin_dir, 'xrefer', 'lang'))
sys.path.append(os.path.join(plugin_dir, 'xrefer', 'legacy'))
sys.path.append(os.path.join(plugin_dir, 'xrefer', 'llm'))
sys.path.append(os.path.join(plugin_dir, 'xrefer', 'llm', 'templates'))
sys.path.append(os.path.join(plugin_dir, 'xrefer', 'loaders'))
sys.path.append(os.path.join(plugin_dir, 'xrefer', 'loaders', 'trace'))

from xrefer.plugin import XReferPlugin


def PLUGIN_ENTRY():
    """
    Entry point for the IDA plugin.

    Returns:
        XReferPlugin: A new instance of the XRefer plugin.
    """
    return XReferPlugin()


repository: OALabs/StrAnnotate-IDA
file: strannotate.py
description: Simple string annotation plugin for IDA
comment: StrAnnotate plugin for IDA
created_at: 2022-09-20 01:08:32+00:00
pushed_at: 2022-09-29 07:48:12+00:00
url:
https://github.com/OALabs/StrAnnotate-IDA/blob/5e61151c3e51802bd50fd761a5481b917ffbc072/strannotate.py

content:
########################################################################################
## 
## IDA plugin for annotating your IDB with imported strings
##
## To install:
##      Copy script into plugins directory, i.e: C:\Program Files\<ida version>\plugins
##
## Strings File:
##      The plugin requires a JSON file that contains a list of strings and their
##      respective offests using the following format.
##
##      "strings":[{"offset":<string_offset>,"value":<ascii_string>},...]
##
##
## To run:
##      Edit->Plugins->StrAnnotate
##      Use the dialogue box to select your strings JSON file and start annotating!
##
########################################################################################

import idaapi
import idautils
import ida_bytes
import idc
import ida_kernwin
import json
import string
import ida_loader

__AUTHOR__ = '@herrcore'

PLUGIN_NAME = "StrAnnotate"
PLUGIN_HOTKEY = ""
VERSION = '1.0.0'

p_initialized = False


def set_hexrays_comment(address, text):
    '''
    set comment in decompiled code
    '''
    try:
        cfunc = idaapi.decompile(address)
        tl = idaapi.treeloc_t()
        tl.ea = address
        tl.itp = idaapi.ITP_SEMI
        if cfunc is not None:
            cfunc.set_user_cmt(tl, text)
            cfunc.save_user_cmts() 
    except:
        print(f"Unable to comment pseudocode at {hex(address)}")


def set_comment(address, text):
    ## Set in dissassembly
    idc.set_cmt(address, text,0)
    ## Set in decompiled data
    set_hexrays_comment(address, text)


#--------------------------------------------------------------------------
# Plugin
#--------------------------------------------------------------------------
class StrAnnotate_Plugin_t(idaapi.plugin_t):
    comment = "StrAnnotate plugin for IDA"
    help = ""
    wanted_name = PLUGIN_NAME
    wanted_hotkey = PLUGIN_HOTKEY
    flags = idaapi.PLUGIN_KEEP


    def init(self):
        global p_initialized

        # register popup menu handlers
        try:
            Searcher.register(self, "StrAnnotate")
        except:
            pass

        if p_initialized is False:
            p_initialized = True
            idaapi.register_action(idaapi.action_desc_t(
                "StrAnnotate",
                "Import strings!",
                self.annotate,
                None,
                None,
                0))
            idaapi.attach_action_to_menu("Edit/StrAnnotate", "StrAnnotate", idaapi.SETMENU_APP)
            ## Print a nice header
            print("=" * 80)
            print("\nStrAnnotate v{0}".format(VERSION))
            print("\n")
            print("* Use this plugin to annotate your IDB with an externally generated strings table *")
            print("\n")
            print("=" * 80)

        return idaapi.PLUGIN_KEEP


    def term(self):
        pass


    def annotate(serlf, strings_file):
        # Opening JSON file
        file_data = None
        with open(strings_file,'r') as fp:
            file_data = fp.read()
        if file_data is None:
            print(f"Error reading file {strings_file}")
            return

        # Parse JSON
        try:
            json_data = json.loads(file_data)
        except:
            print(f"Error parsing file {strings_file}, invalid json format.")
            return

        # Get strings list
        strings_list = json_data.get("strings",None)
        if strings_list is None:
            print(f"Error parsing file {strings_file}, json has no key 'strings'.")
            return

        # Iterate through strings list 
        for string_entry in strings_list:
            # Validate entry
            string_offset = string_entry.get("offset",None)
            if string_offset is None:
                print(f"Invalid string entry {string_entry}, expected format \{'offset':<string_offset>,'value':<ascii_string>\}")
                return
            string_value = string_entry.get("value",None)
            if string_value is None:
                print(f"Invalid string entry {string_entry}, expected format \{'offset':<string_offset>,'value':<ascii_string>\}")
                return
            # Convert offset into ea
            string_address = ida_loader.get_fileregion_ea(string_offset)

            # Lol how to tell if x64 or x32
            if string_address == 0xffffffff or string_address == 0xffffffffffffffff:
                print(f"Invalid string offset {string_entry}, not in range of binary.")
                return

            # If the string is misaligned get the nearest head
            string_address_head = idaapi.get_item_head(string_address)

            # Print the string
            print(f"{hex(string_address_head)} : {string_value}")

            # Add string comment
            set_comment(string_address_head, string_value)


    def run(self, arg):
        strings_file = ida_kernwin.ask_file(0, "*.json", 'Choose string file (json format)...')
        if strings_file == None:
            print("ERROR: You must choose a strings file to annotate with")
        else:
            self.annotate(strings_file)


# register IDA plugin
def PLUGIN_ENTRY():
    return StrAnnotate_Plugin_t()


repository: g4hsean/BinAuthor
file: BinAuthor.py
url: https://github.com/g4hsean/BinAuthor/blob/bdcb86cf16303a5fa50fb2ffb4faf4f88a0050d7/BinAuthor.py
wanted_name: BinAuthor
comment: Author Identification using Novel Techniques
created_at: 2015-11-29 18:07:48+00:00
pushed_at: 2022-03-01 15:48:04+00:00
content:
import sys
from idautils import *
from idaapi import *
import os
from os import listdir
from os.path import isfile
import copy
from pymongo import MongoClient
from datetime import datetime
import hashlib
import idc
from sark.qt import QtGui, QtCore
import sark
import sark.ui
import sark.qt
from idaapi import plugin_t
from pprint import pprint
from idaapi import PluginForm

import pluginConfigurations
import BinAuthorPlugin.Algorithms.FunctionFliterAndColorizer as FunctionFilter
import BinAuthorPlugin.PluginMenuManager.BinAuthorManager as BinAuthorManager

class BinAuthor_plugin_t(plugin_t):
    flags = idaapi.PLUGIN_PROC
    comment = "Author Identification using Novel Techniques"
    help = "Help if a matter of trust."
    wanted_name = "BinAuthor"
    wanted_hotkey = ""
        
    def init(self):
        try:
            self.BinAuthorFunctionFilter = FunctionFilter.FunctionFilter()
        except:
            idaapi.msg("Failed to initialize Function Filter.\n")
            idaapi.msg("Failed to initialize BinAuthor.\n")
            del self.BinAuthorFeatureExtractor
            idaapi.msg("Errors and fun!\n")
            return idaapi.PLUGIN_SKIP
            
        try:
            self.BinAuthor_manager = BinAuthorManager.BinAuthorManager()
            self.BinAuthor_manager.buildMenu(self.BinAuthorFunctionFilter)
            return idaapi.PLUGIN_KEEP
        except:
            idaapi.msg("Failed to initialize BinAuthor.\n")
            del self.BinAuthor_manager
            idaapi.msg("Errors and fun!\n")
            return idaapi.PLUGIN_SKIP
        
        
    def term(self):
        self.BinAuthor_manager.del_menu_items()

    def run(self, arg):
        pass
def PLUGIN_ENTRY():
    return BinAuthor_plugin_t()


repository: LucienMP/idapro_m68k
file: plugins/m68kplugin.py
description: Extends existing support in IDA for the m68k by adding gdb step-over and type information support
wanted_name: m68kstepover
created_at: 2019-02-28 02:02:36+00:00
pushed_at: 2019-07-14 08:14:34+00:00
url:
https://github.com/LucienMP/idapro_m68k/blob/bed147fe59b0267de7a1502e592f3b597dba07fe/plugins/m68kplugin.py

content:
import idaapi
import ida_pro
import idautils

#from pprint import pprint

# Used by the loop for function purge
from idaapi import *

# Used for printing various data
import logging

"""
    This is a sample plugin for extending gdb support for step-over for the M68K,
    and to enable type support so you can press "y" on functions.

    When you install a hook in the HT_IDP category, it will be called
    before the processor module's notify() function, so if you return
    non-zero, your results will be used.

    1) Add GDB style XML files for the CPU register descriptions to $IDA_HOME/cfg/
    2) Add to the section in $IDA_HOME/cfg/dbg_gdb.cfg that integrates that XML for GDB
        CONFIGURATIONS = { "M68K" : ...
        IDA_FEATURES = { "m68k" : ... 
        NOT NEEDED> ARCH_MAP = { "m68k" : ... 
        NOT NEEDED> FEATURE_MAP = ...
    3) Add function to handle step-over understanding around branches
        The ev_calc_step_over interprets the current instruction.
        If it is a JSR/BSR routine call then it will advance the IP to the next instruction.
        In all other cases it will advance to the next insturction by passing back BADADDR
    4) Add function to handle type info (funciton name, and paramater propagation)
        To enable type support you must set PR_TYPEINFO in ph.flag and 
        implement most of the related callbacks (see "START OF TYPEINFO 
        CALLBACKS" in idp.hpp from the IDA SDK).

    Many thanks to Ilfak Guilfanov @ Hey-Rays

    Versioning:
        2019-02-28 v1.0 (c) Lucien Murray-Pitts - Initial version to support step-over
""" 

#--------------------------------------------------------------------------
#
class m68kstepover_idp_hook_t(idaapi.IDP_Hooks):
    #
    def __init__(self):
        idaapi.IDP_Hooks.__init__(self)

    #--------------------------------------------------------------------------
    #
    # Handling of GDB related events
	
    # Calculate the address of the instruction which will be
    # executed after "step over". The kernel will put a breakpoint there.
    #
    def ev_calc_step_over(self, target, ip):
        # Supports the following instructions
        # 4E80..BF        jsr     (a0) +2
        # 61xx            bsr     (a0) +2
        # 6100 xxyy       bsr     (a0) +4
        # 61ff xxyy xxyy  bsr     (a0) +6
        # idaapi.BADADDR  else    BADADDR
        #
        logging.debug("BSR/JSR, checking instruction")
        ta=ida_pro.ea_pointer_frompointer(target)

        insn = idautils.DecodeInstruction(ip)
        if insn and insn.itype in [idaapi.mc_jsr, idaapi.mc_bsr]:
            over_ea = insn.ea + insn.size
            logging.debug(" - JSR STEP> +%s %s" % (insn.size, hex(over_ea)) )
            ta.assign(over_ea)
        else:
            logging.debug(" - NO STEP> "+hex(ip))
            ta.assign(idaapi.BADADDR)

# Original code before IDA guys helped me write the above!
#
#        if idaapi.get_bytes(ip, 2) >= "\x4E\x80" and idaapi.get_bytes(ip, 2) <= "\x4E\xBF":
#            logging.debug( "JSR STEP> +2 " + hex(ip+2) )
#            ta.assign(ip+2)
#        elif idaapi.get_bytes(ip, 2) == "\x61\x00":
#            logging.debug( "BSR STEP> +4 " + hex(ip+4) )
#            ta.assign(ip+4)
#        elif idaapi.get_bytes(ip, 2) == "\x61\xff":
#            logging.debug( "BSR STEP> +6 " + hex(ip+6) )
#            ta.assign(ip+6)
#        elif idaapi.get_bytes(ip, 1) == "\x61":
#            logging.debug( "BSR STEP> +2 " + hex(ip+2) )
#            ta.assign(ip+2)
#        else:
#            logging.debug( "NO STEP> "+hex(ip) )
#            ta.assign(idaapi.BADADDR)
        return 1

    #--------------------------------------------------------------------------
    #
    # Handling of TYPE related events
	
    # Setup default type libraries.
    # (called after loading a new file into the database). The processor module may load tils, 
    # setup memory model and perform other actions required to set up the type system. This is an optional callback.
    #
    def ev_setup_til(self, abi_names, abi_opts, comp):
        logging.debug("##########ev_setup_til")
        return
		
    # Get all possible ABI names and optional extensions for given compiler
    # abiname/option is a string entirely consisting of letters, digits and underscore
    # 0=not implemented, 1=ok
    #
    def ev_get_abi_info(self, abi_names, abi_opts, comp):
        logging.debug("##########ev_get_abi_info")
        logging.debug(abi_names)
        logging.debug(abi_opts)
        logging.debug(comp)
        return 0

    # Get maximal size of a pointer in bytes.
    #
    #
    def ev_max_ptr_size(self, target, ip):
        logging.debug("##########ev_max_ptr_size")
        return 4

    # Get default enum size.
    # returns sizeof(enum)
    #
    def ev_get_default_enum_size(self, cm):
        logging.debug("##########ev_get_default_enum_size")
        return 4

    # Get register allocation convention for given calling convention.
    # 1, or 0=not implemented
    #
    def ev_get_cc_regs(self, regs, cc):
        logging.debug("##########ev_get_cc_regs")
        return 0

    # Get offset from SP to the first stack argument.
    #
    #
    def ev_get_stkarg_offset(self):
        logging.debug("##########ev_get_stkarg_offset")
        return 0

    # Get SIMD-related types according to given attributes ant/or argument location.
    # returns number of found types, -1-error If name==NULL, initialize all SIMD types
    #
    def ev_get_simd_types(self, out, simd_attrs, argloc, create_tifs):
        logging.debug("##########ev_get_simd_types")
        return 0

    # Calculate number of purged bytes after call.
    # returns number of purged bytes (usually add sp, N)
    # 
    def ev_calc_cdecl_purged_bytes(self, ea):
        logging.debug("##########ev_calc_cdecl_purged_bytes")
        return 0

    # Calculate number of purged bytes by the given function type.
    # 1, or 0 = not implemented
    #
    def ev_calc_purged_bytes(self, p_purged_bytes, fti):
        logging.debug( "##########ev_calc_purged_bytes" )
        logging.debug( fti.size() )
        logging.debug( fti.get_call_method() )
        logging.debug( fti.is_pure() )
        logging.debug( fti.rettype.get_size() )
        
        # Information about each argument IF there are any, NOTE that if the next instruction is unlink then NO purge is done
        ##########
        # NOTES:
        #   https://reverseengineering.stackexchange.com/questions/8870/extracting-arguments-from-ida
        #   https://github.com/gsmk/hexagon/blob/master/pmbase.h
        #   https://github.com/idapython/src/blob/master/pywraps/py_idp.hpp#L608
        #   LOCATION TYPE: https://www.hex-rays.com/products/ida/support/sdkdoc/group___a_l_o_c__.html
        #
        for i in xrange(fti.size()):
                logging.debug( "Arg %d: %s (of type %s, and of location: %s)" % (i, fti[i].name, idaapi.print_tinfo('', 0, 0, idaapi.PRTYPE_1LINE, fti[i].type, '', ''), fti[i].argloc.atype()) )

        # Sort of wrong, there could be shorts or other odd sized things on the stack
        p_purged_bytes = fti.size() * 4
        
        return 1

    # Calculate return value location.
    # 0 not implemented, 1 ok, -1 error
    #
    def ev_calc_retloc(self, retloc, rettype, cc):
        logging.debug( "##########ev_calc_retloc" )
        
        logging.debug( "retloc=" + str(retloc) ) 
        logging.debug( "rettype=" + str(rettype) )
        logging.debug( "cc=" + str(strcc) )
		
        # SEE: add_argregs in emu.cpp 
        # IF just a single reg then reg1, else if a range reg2 START##########END
        retloc.set_reg1(0)
        # IF its a range reg2 START##########END
        #retloc.set_reg2(0,0)
        # IF its a mix, then create a scattered_aloc_t, and push_back multiple part argpart_t
        #retloc.consume_scattered(0)
		
        #retloc.set_stkoff(0)
        #retloc.set_ea(0)
        #retloc.consume_rrel(0) # argloc_t_consume_rrel
        #argloc_t
        return 1


    # Calculate function argument locations.
    # This callback should fill retloc, all arglocs, and stkargs.
    # This callback is never called for ::CM_CC_SPECIAL functions.
    # 0=not implemented, 1=ok, -1=error
    #
    # See: alloc_args from arc emu.cpp in SDK
    def ev_calc_arglocs(self, fti):
        logging.debug(  "##########ev_calc_arglocs" )
        logging.debug(  "is_void() = " + str(fti.rettype.is_void()) )

        # Set the location of each argument
        #    NOTE: LOCATION TYPE: https://www.hex-rays.com/products/ida/support/sdkdoc/group___a_l_o_c__.html
        #
        for i in xrange(fti.size()):
            logging.debug(  "Arg %d: %s (of type %s, and of location: %s)" % (i, fti[i].name, idaapi.print_tinfo('', 0, 0, idaapi.PRTYPE_1LINE, fti[i].type, '', ''), fti[i].argloc.atype()) )
            # Set funciton argument #i location as register number "i"
#            fti[i].argloc.set_reg1(i)
            # Set funciton argument #i location as stack at location "i"*4
            fti[i].argloc.set_stkoff(i*4)
        # Return value is D0
        fti.retloc.set_reg1(0) # D0
        # Set size of stack arguments passed
        fti.stkargs = fti.size() * 4
        return 1

    # Calculate locations of the arguments that correspond to '...'
    # 0=not implemented, -1=error, 1=ok
    # 
    def ev_calc_varglocs(self, ftd, regs, stkargs, nfixed):
        logging.debug(  "##########ev_calc_varglocs" )
        return 0

    # Adjust argloc according to its type/size and platform endianess.
    # 0=not implemented, -1=error, 1=ok
    # 
    def ev_adjust_argloc(self, argloc, optional_type, size):
        logging.debug(  "##########ev_adjust_argloc" )
        return 0

    # Get function arguments which should be converted to pointers when lowering function prototype.
    # The processor module can also modify 'fti' in order to make non-standard conversion
    # 0=not implemented, 1=argnums was filled, 2 argnums was filled and made substantial changes to fti
    def ev_lower_func_type(self, argnums, fti):
        logging.debug(  "##########ev_lower_func_type" )
        return 0

    # Are 2 register arglocs the same?
    # 1=yes, -1=no, 0=not implemented
    # 
    def ev_equal_reglocs(self, a1, a2):
        logging.debug(  "##########ev_equal_reglocs" )
        return 0

    # Use information about a stack argument.
    #  retval <=0 failed, the kernel will create a comment with the
    #  argument name or type for the instruction
    #
    def ev_use_stkarg_type(self, ea, arg):
        logging.debug(  "##########ev_use_stkarg_type" )
        return 0

    # Use information about register argument.
    # 0=not implemented, 1=???
    # 
    def ev_use_regarg_type(self, ea, rargs):
        logging.debug(  "##########ev_use_regarg_type" )
        return 0

    # Use information about callee arguments.
    # 0=not implemented, 1 (and removes handled arguments from fti and rargs)
    # 
    def ev_use_arg_types(self, ea, fti, rargs):
        logging.debug(  "##########ev_use_arg_types" )
        return 0

    # Argument address info is ready.
    # <0: do not save into idb; other values mean "ok to save"
    #
    def ev_arg_addrs_ready(self, caller, n, tif, addrs):
        logging.debug(  "##########ev_arg_addrs_ready" )
        return 1

    # Decorate/undecorate a C symbol name, unlike C the return is the demanged name
    #   ev_decorate_name(self, name, mangle, cc, optional_type) -> PyObject *
    #
    def ev_decorate_name(self, name, mangle, cc, optional_type):
        logging.debug(  "##########ev_decorate_name" )
        logging.debug(  "NAME=" + str(name) )
        logging.debug(  "MANGLE?=" + str(mangle) )
        logging.debug(  "CC=" + str(cc) )
        logging.debug(  "OptionalType=" + str(optional_type) )
        logging.debug(  "procName=" + str(idaapi.get_inf_structure().procName) )
        
        # Compiler Info "compiler_info_t cc;                   ///< Target compiler" from ida.hpp
        logging.debug(  "cc.id=" + str(idaapi.get_inf_structure().cc.id) )
        
        #get_compiler_name
        # FROM: idaapi.gen_decorate_name(name, mangle, cc, type) -> bool 
        gen_decorate_name(name,mangle,cc,optional_type)
        logging.debug( "NAME POST Decoration:" + str(name) )
        return name

        
        
#--------------------------------------------------------------------------
# CTYPE mapping for global ph structure into python
#

# NOTE: Requires MSVCR90.dll, you will need the 2008 redist from MS for this to work right
#import ctypes
from ctypes import *
from binascii import hexlify
from sys import getsizeof
import ida_idp as ida_idp

# Partial structure information for accessing PH directly
#
class structPH(Structure):
        _fields_ = [("version", c_uint),
        ("id", c_uint),   
        ("flag", c_uint)]    


#--------------------------------------------------------------------------
class m68kstepover(idaapi.plugin_t):
    # Processor fix plugin module
    flags = idaapi.PLUGIN_PROC | idaapi.PLUGIN_HIDE
    comment = ""
    wanted_hotkey = ""
    help = "Supports Step-Over for M68K CPU on BSR/JSR, and type information"
    wanted_name = "m68kstepover"
       
    def init(self):
        self.prochook = None
        
        # Enable controlable logging 
        logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(filename)s:%(funcName)s@%(lineno)s: %(message)s')
       
        logging.debug('Module for MC68000 class CPU Step-Overs, and type information')
        
        # Random Debug to analyze PH mapping support
        #
        #print cdll.ida.ph
        #ph = structPH.in_dll(cdll.ida, "ph")
        #ph = structPH()
        #ph.values = cast( cdll.ida.ph, POINTER(structPH) )
		
        # Map ida.dll into this Python module, and lets us change the flags to add TYPEINFO support.
        #
        # For O/S agnostic support this follow advice from: http://www.hexblog.com/?p=695
        #
        idaname = "ida64" if ida_idaapi.__EA64__ else "ida"
        try:
            logging.debug( "Platform for ida library loading: " + str(sys.platform) )
            if sys.platform == "win32":
                ida_dll = ctypes.windll[idaname + ".dll"]
            elif sys.platform == "linux2":
                ida_dll = ctypes.cdll["lib" + idaname + ".so"]
            elif sys.platform == "darwin":
                ida_dll = ctypes.cdll["lib" + idaname + ".dylib"]        
        except:
            logging.error( "Python CTYPES probably couldnt find the dll, OR the MSCVR90.dll or other system DLL is missing" )
            return idaapi.PLUGIN_SKIP
		
        # Access ph global in ida.dll using ctypes
        ph = structPH.in_dll(ida_dll, "ph")
        
        # Log Values via PH direct access
        logging.debug( "- - - - - - - - - - - - - - - - - -" )
        logging.debug( "  ph.version = %d" % ph.version )
        logging.debug( "  ph.id      = %d" % ph.id )
        logging.debug( "  ph.flag    = %d" % ph.flag )

        # Enable TYPEINFO support
        ph.flag = ph.flag | ida_idp.PR_TYPEINFO

        # Official API, confirms PH was correctly accessed
        logging.debug( "- - - - - - - - - - - - - - - - - -" )
        logging.debug( "  ph.version = %d" % idaapi.ph_get_version() )
        logging.debug( "  ph.id      = %d" % idaapi.ph_get_id() )
        logging.debug( "  ph.flag    = %d" % idaapi.ph_get_flag() )

        #print(hexlify(string_at(cdll.ida.ph, getsizeof(G))))
			
        if idaapi.ph_get_id() != idaapi.PLFM_68K: # or idaapi.cvar.inf.filetype != idaapi.f_ELF:
            logging.error("skipped this plugin because CPU being used is not 68000 derivative!")
            return idaapi.PLUGIN_SKIP

        # Connect the IDP hooks
        self.prochook = m68kstepover_idp_hook_t()
        self.prochook.hook()

        logging.debug( "succeeded, stepover support should now work for BSR/JSR!" )
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        pass

    def term(self):
        logging.debug( "ida is terminating, removing step-over support event handling!" )
        if self.prochook:
            self.prochook.unhook()

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


repository: clovme/WTools
file: software/Reverse Cracking/IDA Professional 9.0/python/examples/core/install_user_defined_prefix.py
wanted_name: user defined prefix
comment: This is a user defined prefix sample plugin
created_at: 2025-03-27 00:14:54+00:00
pushed_at: 2025-05-11 10:18:21+00:00
url:
https://github.com/clovme/WTools/blob/f12f01c833a6a6241799a9d01e9117cc28a65dbc/software/Reverse%20Cracking/IDA%20Professional%209.0/python/examples/core/install_user_defined_prefix.py

content:
"""
summary: inserting information into disassembly prefixes

description:
  By default, disassembly line prefixes contain segment + address
  information (e.g., '.text:08047718'), but it is possible to
  "inject" other bits of information in there, thanks to the
  `ida_lines.user_defined_prefix_t` helper type.
"""

import ida_lines
import ida_idaapi

PREFIX = ida_lines.SCOLOR_INV + ' ' + ida_lines.SCOLOR_INV

class my_user_prefix_t(ida_lines.user_defined_prefix_t):
    def get_user_defined_prefix(self, ea, insn, lnnum, indent, line):
        if (ea % 2 == 0) and indent == -1:
            return PREFIX
        else:
            return ""


class prefix_plugin_t(ida_idaapi.plugin_t):
    flags = 0
    comment = "This is a user defined prefix sample plugin"
    help = "This is help"
    wanted_name = "user defined prefix"
    wanted_hotkey = ""

    def __init__(self):
        self.prefix = None

    def init(self):
        self.prefix = my_user_prefix_t(8)
        print("prefix installed")
        return ida_idaapi.PLUGIN_KEEP

    def run(self, arg):
        pass

    def term(self):
        self.prefix = None
        print("prefix uninstalled!")


def PLUGIN_ENTRY():
    return prefix_plugin_t()



repository: techbliss/PyQt_Geo_Locator
file: 7.0+/plugins/Geolocateloader.py
description: IP Geolocator and maps using PyQt
wanted_name: Geo locate
comment: Run me
created_at: 2016-09-11 19:03:49+00:00
pushed_at: 2019-03-10 16:38:33+00:00
url:
https://github.com/techbliss/PyQt_Geo_Locator/blob/466bbedb26ab88559d0b64d5dcd022d628d73031/7.0%2B/plugins/Geolocateloader.py

content:
# Created by: Storm Shadow http://www.techbliss.org
import os
import ida_idaapi, ida_kernwin
import idc
from idc import *
from idaapi import *
import sys
sys.path.insert(0 , idaapi.idadir("plugins\\geo\\icons")) #change path to the icon
sys.path.insert(0 , idaapi.idadir("plugins"))
sys.path.insert(0, idaapi.idadir("plugins\\geo\\"))
import ico
from ico import *
import subprocess

PLUGIN_VERSION = "1.4" #change Plugin Version
IDAVERISONS    = "IDA PRO 7.0+" #change Ida version
AUTHORS        = "Storm Shadow" #change author
DATE           = "2017" #change date
TWITTER        = "Twitter @zadow28" #change social media

def banner():
    banner_options = (PLUGIN_VERSION, AUTHORS, DATE, TWITTER, IDAVERISONS)
    banner_titles = "Geo Locator v%s - (c) %s - %s - %s - %s" % banner_options #change Python editor

# print plugin banner
    print("---[" + banner_titles + "]---\n")

banner()

# 1) Create the handler class
class MyEditorHandler(idaapi.action_handler_t): # the class name must be unique for each plugin nb same as line (53)
    def __init__(self):
        idaapi.action_handler_t.__init__(self)

    # Run editor when invoked.
    def activate(self, ctx):
		g = globals()
		idahome = idaapi.idadir("plugins\\geo") #change to set plguin path where the main plguin script is
		process = subprocess.Popen(['python', idahome + "\\geo.py"])
		#print os.getcwd()		#change for the main plugin script

    def update(self, ctx):
        return idaapi.AST_ENABLE_ALWAYS

class geoeye(idaapi.plugin_t): #change class to unique name for each plugin , also change line (103)
    flags = idaapi.PLUGIN_FIX #different flags this is for plugin visible at startup
    comment = "Run me" #help me text
    help = "Geo helper" #help text
    wanted_name = "Geo locate" #change the plugins name
    wanted_hotkey = "" #the tooltip goes away in menu when setting it here ,DONT DO it! and only is shown in File/Plugins menu


    def locate_menuaction(self): #change for something unique.also change line (90)
        action_desc = idaapi.action_desc_t(
            'my:locateaction',  # The action name. This acts like an ID and must be unique same as line (64), (68)
            'Geo locate!',  # The action text.
            MyEditorHandler(),  # The action handler must be unique , also change line (99)
            'Ctrl+G',  # Optional: the action shortcut DO IT  HERE!
            'Geo locate',  # Optional: the action tooltip (available in menus/toolbar)
            idaapi.load_custom_icon(":/ico/python.png")  # hackish load action icon , if no custom icon use number from 1-150 from internal ida
        )

        # 3) Register the action
        idaapi.register_action(action_desc)

        idaapi.attach_action_to_menu(
            'File/Editor...',  # The relative path of where to add the action
            'my:locateaction',  # The action ID (see line(51))
            idaapi.SETMENU_APP)  # We want to append the action after the 'Manual instruction...

        form = idaapi.get_current_tform()
        idaapi.attach_action_to_popup(form, None, "my:locateaction", None) # The action ID (see line(51)

    def init(self):
        """
        This is called by IDA when it is loading the plugin.
        """
        #self._icon_id_file = idaapi.BADADDR
        # attempt plugin initialization
        try:
            self._install_plugin()

        # failed to initialize or integrate the plugin, log and skip loading
        except Exception as e:
            form = idaapi.get_current_tform()
            pass

        return PLUGIN_KEEP

    def _install_plugin(self):
        """
        Initialize & integrate the plugin into IDA.
        """
        self.locate_menuaction() #same as line (49)
        self._init()

    def term(self):
        pass

    def run(self, arg = 0):
        #we need the calls again if we wanna load it via File/Plugins/editor
        idaapi.msg("Python Geo loacte Loaded to menu \n use Alt+G hot key to quick load ")
        hackish = MyEditorHandler() #must be the same as line (53) change also hackish =  must be unique also line (100)
        hackish.activate(self)  #hackish must the same as line hackish = (99)

def PLUGIN_ENTRY():
    return geoeye() #must be unique for each plugin and same as line 43


repository: VoidSec/DriverBuddyReloaded
file: DriverBuddyReloaded.py
wanted_name: Driver Buddy Reloaded
created_at: 2021-10-22 14:55:04+00:00
pushed_at: 2024-10-25 10:06:20+00:00
description:
Driver Buddy Reloaded is an IDA Pro Python plugin that helps automate some tedious Windows Kernel Drivers reverse engineering tasks

url:
https://github.com/VoidSec/DriverBuddyReloaded/blob/311c9e743f9299914bd5b30a91a5fc80c8dd9cff/DriverBuddyReloaded.py

content:
import os

import idaapi
import idc
from DriverBuddyReloaded import NTSTATUS
from DriverBuddyReloaded import device_name_finder
from DriverBuddyReloaded import dump_pool_tags
from DriverBuddyReloaded import ioctl_decoder
from DriverBuddyReloaded import utils

"""
DriverBuddyReloaded.py: Entry point for IDA python plugin used in Windows driver vulnerability research.
Updated in 2021 by Paolo Stagno aka VoidSec: https://voidsec.com - https://twitter.com/Void_Sec
"""
# needed GLOBALs
driver_name = idaapi.get_root_filename()
path = "{}{}".format(os.getcwd(), os.sep)
ioctl_file_name = "{}-{}-{}-IOCTLs.txt".format(driver_name, utils.today(), utils.timestamp())
analysis_file_name = "{}-{}-{}-DriverBuddyReloaded_autoanalysis.txt".format(driver_name, utils.today(),
                                                                            utils.timestamp())
pool_file_name = "{}-{}-{}-pooltags.txt".format(driver_name, utils.today(), utils.timestamp())


class UiAction(idaapi.action_handler_t):
    """
    Simple wrapper class for creating action handlers which add options to menu's and are triggered via hot keys
    """

    def __init__(self, id, name, tooltip, menuPath, callback, shortcut):
        """
        :param id: id
        :param name: action name
        :param tooltip: action tooltip
        :param menuPath: where to register the action
        :param callback: function to execute
        :param shortcut: hot-keys to register
        """

        idaapi.action_handler_t.__init__(self)
        self.id = id
        self.name = name
        self.tooltip = tooltip
        self.menuPath = menuPath
        self.callback = callback
        self.shortcut = shortcut

    def registerAction(self):
        """
        Register an action which add an options to the menu's
        :return:
        """

        action_desc = idaapi.action_desc_t(
            self.id,
            self.name,
            self,
            self.shortcut,
            self.tooltip,
            0
        )
        if not idaapi.register_action(action_desc):
            return False
        if not idaapi.attach_action_to_menu(self.menuPath, self.id, 0):
            return False
        return True

    def unregisterAction(self):
        """
        Unregister an action which remove an options from the menu's
        :return:
        """

        idaapi.detach_action_from_menu(self.menuPath, self.id)
        idaapi.unregister_action(self.id)

    def activate(self, ctx):
        """
        :param ctx:
        :return:
        """

        self.callback()
        return 1

    def update(self, ctx):
        """
        :param ctx:
        :return:
        """

        return idaapi.AST_ENABLE_ALWAYS


def make_comment(pos, string):
    """
    Creates a comment with contents `string` at address `pos`.
    If the address is already commented append the new comment to the existing comment
    :param pos: position where to create the comment
    :param string: comment to write
    :return:
    """

    current_comment = idc.get_cmt(pos, 0)
    if not current_comment:
        idc.set_cmt(pos, string, 0)
    elif string not in current_comment:
        idc.set_cmt(pos, current_comment + " " + string, 0)


def get_operand_value(addr):
    """
    Returns the value of the second operand to the instruction at `addr` masked to be a 32 bit value
    :param addr: address to get the operand from
    :return:
    """

    return idc.get_operand_value(addr, 1) & 0xffffffff


class IOCTLTracker:
    """
    A simple container to keep track of decoded IOCTL codes and codes marked as invalid
    """

    def __init__(self):
        self.ioctl_locs = set()
        self.ioctls = []

    def add_ioctl(self, addr, value):
        """
        Add discovered IOCTL code to the list
        :param addr: address
        :param value: IOCTL code
        :return:
        """

        self.ioctl_locs.add(addr)
        self.ioctls.append((addr, value))

    def remove_ioctl(self, addr, value):
        """
        Remove discovered IOCTL code to the list
        :param addr: address
        :param value: IOCTL code
        :return:
        """

        self.ioctl_locs.remove(addr)
        self.ioctls.remove((addr, value))

    def print_table(self, ioctls):
        """
        Print table of decoded IOCTL codes and write the result to a file
        :param ioctls: IOCTL to decode
        :return:
        """
        try:
            with open(ioctl_file_name, "w") as IOCTL_file:
                print("\nDriver Buddy Reloaded - IOCTLs\n"
                      "-----------------------------------------------")
                IOCTL_file.write("Driver Buddy Reloaded - IOCTLs\n"
                                 "-----------------------------------------------\n")
                print("%-10s  | %-10s | %-42s | %-10s | %-22s | %s" % (
                    "Address", "IOCTL Code", "Device", "Function", "Method", "Access"))
                IOCTL_file.write("%-10s | %-10s | %-42s | %-10s | %-22s | %s\n" % (
                    "Address", "IOCTL Code", "Device", "Function", "Method", "Access"))
                for (addr, ioctl_code) in ioctls:
                    function = ioctl_decoder.get_function(ioctl_code)
                    device_name, device_code = ioctl_decoder.get_ioctl_code(ioctl_code)
                    method_name, method_code = ioctl_decoder.get_method(ioctl_code)
                    access_name, access_code = ioctl_decoder.get_access(ioctl_code)
                    all_vars = (
                        addr, ioctl_code, device_name, device_code, function, method_name, method_code, access_name,
                        access_code)
                    print("0x%-8X | 0x%-8X | %-31s 0x%-8X | 0x%-8X | %-17s %-4d | %s (%d)" % all_vars)
                    IOCTL_file.write("0x%-8X | 0x%-8X | %-31s 0x%-8X | 0x%-8X | %-17s %-4d | %s (%d)\n" % all_vars)
            print("\n[>] Saved decoded IOCTLs to \"{}{}\"".format(path, ioctl_file_name))
        except IOError as e:
            print("ERROR #{}: {}\nCan't save decoded IOCTLs to \"{}{}\"".format(e.errno, e.strerror, path,
                                                                                ioctl_file_name))
            print("\nDriver Buddy Reloaded - IOCTLs\n"
                  "-----------------------------------------------")
            print("%-10s  | %-10s | %-42s | %-10s | %-22s | %s" % (
                "Address", "IOCTL Code", "Device", "Function", "Method", "Access"))
            for (addr, ioctl_code) in ioctls:
                function = ioctl_decoder.get_function(ioctl_code)
                device_name, device_code = ioctl_decoder.get_ioctl_code(ioctl_code)
                method_name, method_code = ioctl_decoder.get_method(ioctl_code)
                access_name, access_code = ioctl_decoder.get_access(ioctl_code)
                all_vars = (
                    addr, ioctl_code, device_name, device_code, function, method_name, method_code, access_name,
                    access_code)
                print("0x%-8X | 0x%-8X | %-31s 0x%-8X | 0x%-8X | %-17s %-4d | %s (%d)" % all_vars)


def find_all_ioctls():
    """
    From the currently selected address attempts to traverse all blocks inside the current function to find all
    immediate values which are used for a comparison/sub immediately before a jz. Returns a list of address, second operand pairs.
    :return:
    """

    ioctls = []
    # Find the currently selected function and get a list of all of its basic blocks
    addr = idc.get_screen_ea()
    f = idaapi.get_func(addr)
    fc = idaapi.FlowChart(f, flags=idaapi.FC_PREDS)
    for block in fc:
        start = block.start_ea
        end = block.end_ea
        # print("Block: {} - {}".format(start, end))
        for instr in range(start, end):
            # if the penultimate instruction is cmp or sub or mov against an immediate value
            if idc.print_insn_mnem(instr) in ['cmp', 'sub', 'mov'] and idc.get_operand_type(instr, 1) == 5:
                value = get_operand_value(instr)
                # value >= 0x10000 (lower false positives) and is not a known NTSTATUS value; check issue #15
                if value >= 0x10000 and value not in NTSTATUS.ntstatus_values:
                    ioctls.append((instr, value))
                    ioctl_tracker.add_ioctl(instr, value)
    return ioctls


def track_ioctls(ioctls):
    """
    Decode and add an IOCTL code to the global table, generate the comments
    :param ioctls: IOCTL code
    :return:
    """

    global ioctl_tracker
    for addr, ioctl_code in ioctls:
        ioctl_tracker.add_ioctl(addr, ioctl_code)
        define = ioctl_decoder.get_define(ioctl_code)
        make_comment(addr, define)
    ioctl_tracker.print_table(ioctls)


def decode_all_ioctls():
    """
    Attempts to locate all the IOCTLs in a function and decode them all
    :return:
    """

    global ioctl_tracker
    ioctls = find_all_ioctls()
    track_ioctls(ioctls)


def get_position_and_translate():
    """
    Gets the current selected address and decodes the second parameter to the instruction if it exists/is an immediate
    then adds the C define for the code as a comment and prints a summary table of all decoded IOCTL codes.
    :return:
    """

    pos = idc.get_screen_ea()
    if idc.get_operand_type(pos, 1) != 5:  # Check the second operand to the instruction is an immediate
        return

    value = get_operand_value(pos)
    # value >= 0x10000 (lower false positives) and is not a known NTSTATUS value; check issue #15
    if value >= 0x10000 and value not in NTSTATUS.ntstatus_values:
        ioctl_tracker.add_ioctl(pos, value)
        define = ioctl_decoder.get_define(value)
        make_comment(pos, define)
        # Print summary table each time a new IOCTL code is decoded
        ioctls = []
        for inst in ioctl_tracker.ioctl_locs:
            value = get_operand_value(inst)
            ioctls.append((inst, value))
        ioctl_tracker.print_table(ioctls)


class ActionHandler(idaapi.action_handler_t):
    """
    Basic wrapper class to avoid all action handlers needing to implement update identically
    """

    def update(self, ctx):
        """
        :param ctx:
        :return:
        """

        return idaapi.AST_ENABLE_ALWAYS


class DecodeHandler(ActionHandler):
    """
    Wrapper for `get_position_and_translate` used for right-click context menu hook
    """

    def activate(self, ctx):
        """
        :param ctx:
        :return:
        """

        get_position_and_translate()


class DecodeAllHandler(ActionHandler):
    """
    Wrapper for `decode_all_ioctls` used for right-click context menu hook
    """

    def activate(self, ctx):
        """
        :param ctx:
        :return:
        """

        decode_all_ioctls()


class InvalidHandler(ActionHandler):
    """
    Only available when right-clicking on an address marked as an IOCTL code location, removes it from the location list
    and deletes C define comment marking it (but leaves any other comment content at that location intact).
    """

    def activate(self, ctx):
        """
        :param ctx:
        :return:
        """

        pos = idc.get_screen_ea()
        # Get current comment for this instruction and remove the C define from it, if present
        comment = idc.get_cmt(pos, 0)
        code = get_operand_value(pos)
        define = ioctl_decoder.get_define(code)
        comment = comment.replace(define, "")
        idc.set_cmt(pos, comment, 0)
        # Remove the ioctl from the valid list and add it to the invalid list to avoid 'find_all_ioctls' accidently re-indexing it.
        ioctl_tracker.remove_ioctl(pos, code)


def register_dynamic_action(form, popup, description, handler):
    """
    Registers a new item in a popup which will trigger a function when selected
    :param form:
    :param popup:
    :param description:
    :param handler:
    :return:
    """

    # Note the 'None' as action name (1st parameter), that's because the action will be deleted immediately
    # after the context menu is hidden anyway, so there's no need giving it a valid ID.
    action = idaapi.action_desc_t(None, description, handler)
    idaapi.attach_dynamic_action_to_popup(form, popup, action, 'Driver Buddy Reloaded/')


class WinDriverHooks(idaapi.UI_Hooks):
    """
    Installs hook function which is triggered when popup forms are created
    and adds extra menu options if it is the right-click disasm view menu
    """

    def finish_populating_widget_popup(self, form, popup):
        """
        :param form:
        :param popup:
        :return:
        """

        if idaapi.get_widget_type(form) != idaapi.BWN_DISASM:
            return

        pos = idc.get_screen_ea()
        register_dynamic_action(form, popup, 'Decode All IOCTLs in Function', DecodeAllHandler())
        # If the second argument to the current selected instruction is an immediately
        # then give the option to decode it.
        if idc.get_operand_type(pos, 1) == 5:
            register_dynamic_action(form, popup, 'Decode IOCTL', DecodeHandler())
            if pos in ioctl_tracker.ioctl_locs:
                register_dynamic_action(form, popup, 'Invalid IOCTL', InvalidHandler())


class DriverBuddyPlugin(idaapi.plugin_t):
    """
    Main entry class for DriverBuddyPlugin
    """
    flags = idaapi.PLUGIN_UNL
    comment = ("Plugin to aid in Windows driver vulnerability research. " +
               "Automatically tries to find IOCTL handlers, decode IOCTLS, " +
               "flag dangerous C/C++ functions, find Windows imports for privilege escalation, " +
               "dump Pooltags and identify the type of Windows driver.")
    help = ""
    wanted_name = "Driver Buddy Reloaded"
    wanted_hotkey = "Ctrl+Alt+A"

    def init(self):
        """
        Define hooks and shortcut actions
        :return:
        """
        global ioctl_tracker
        ioctl_tracker = IOCTLTracker()
        global hooks
        hooks = WinDriverHooks()
        hooks.hook()
        decode_ioctl = UiAction(
            id="ioctl:decode",
            name="Decode IOCTL",
            tooltip="Decodes the currently selected constant into its IOCTL details.",
            menuPath="",
            shortcut="Ctrl+Alt+D",
            callback=get_position_and_translate
        )
        decode_ioctl.registerAction()
        decode_all_ioctl = UiAction(
            id="ioctl:decode_all",
            name="Decode ALL IOCTLs in a Function",
            tooltip="Decodes ALL IOCTLs in a Function into its IOCTL details.",
            menuPath="",
            shortcut="Ctrl+Alt+F",
            callback=decode_all_ioctls
        )
        decode_all_ioctl.registerAction()
        return idaapi.PLUGIN_KEEP

    def run(self, args):
        """
        Driver Buddy Reloaded Auto-analysis function
        :param args:
        :return:
        """
        try:
            with open(analysis_file_name, "w") as log_file:
                print("\nDriver Buddy Reloaded Auto-analysis\n"
                      "-----------------------------------------------")
                log_file.write("\nDriver Buddy Reloaded Auto-analysis\n"
                               "-----------------------------------------------\n")
                idc.auto_wait()  # Wait for IDA analysis to complete
                file_type = idaapi.get_file_type_name()
                if "portable executable" not in file_type.lower():
                    print("[!] ERR: Loaded file is not a valid PE")
                    log_file.write("[!] ERR: Loaded file is not a valid PE\n")
                else:
                    driver_entry_addr = utils.is_driver()
                    if driver_entry_addr is False:
                        print("[!] ERR: Loaded file is not a Driver")
                        log_file.write("[!] ERR: Loaded file is not a Driver\n")
                    else:
                        print("[+] `DriverEntry` found at: 0x{addr:08x}".format(addr=driver_entry_addr))
                        log_file.write("[+] `DriverEntry` found at: 0x{addr:08x}\n".format(addr=driver_entry_addr))
                        print("[>] Searching for `DeviceNames`...")
                        log_file.write("[>] Searching for `DeviceNames`...\n")
                        device_name_finder.search(log_file)
                        print("[>] Searching for `Pooltags`...")
                        log_file.write("[>] Searching for `Pooltags`...\n")
                        pool = dump_pool_tags.get_all_pooltags()
                        if pool:
                            print(pool)
                            try:
                                with open(pool_file_name, "w") as pool_file:
                                    pool_file.write(pool)
                            except IOError as e:
                                print(
                                    "ERROR #{}: {}\nCan't write pool file to \"{}{}\"".format(e.errno, e.strerror, path,
                                                                                              pool_file_name))
                        if utils.populate_data_structures(log_file) is True:
                            driver_type = utils.get_driver_id(driver_entry_addr, log_file)
                            print("[+] Driver type detected: {}".format(driver_type))
                            log_file.write("[+] Driver type detected: {}\n".format(driver_type))
                            if ioctl_decoder.find_ioctls_dumb(log_file, ioctl_file_name) is False:
                                print("[!] Unable to automatically find any IOCTLs")
                                log_file.write("[!] Unable to automatically find any IOCTLs\n")
                            else:
                                print("\n[>] Saved decoded IOCTLs log file to \"{}{}_dumb.txt\"".format(path,
                                                                                                        ioctl_file_name))
                        else:
                            print("[!] ERR: Unable to enumerate functions")
                            log_file.write("[!] ERR: Unable to enumerate functions\n")
                print("[+] Analysis Completed!\n"
                      "-----------------------------------------------")
                log_file.write("[+] Analysis Completed!\n"
                               "-----------------------------------------------")
            print("\n[>] Saved Autoanalysis log file to \"{}{}\"".format(path, analysis_file_name))
            if pool:
                print("[>] Saved Pooltags file to \"{}{}\"".format(path, pool_file_name))
        except IOError as e:
            print("ERROR #{}: {}\nAutoanalysis aborted, can't write log file to \"{}{}\"".format(e.errno, e.strerror,
                                                                                                 path,
                                                                                                 analysis_file_name))
        return

    def term(self):
        pass


def PLUGIN_ENTRY():
    return DriverBuddyPlugin()


repository: michal-kapala/dubRE
file: plugins/utf16_export.py
description: ML-driven function symbol extraction plugin for IDA Pro
wanted_name: dubRE UTF-16 String Exporter
comment: dubRE UTF-16 String Exporter
created_at: 2023-05-23 15:03:18+00:00
pushed_at: 2024-12-17 22:26:33+00:00
url:
https://github.com/michal-kapala/dubRE/blob/ce7be63623178ed55081e3b3297a4b4f3d49ac6a/plugins/utf16_export.py

content:
import idaapi, idc
import sqlite3
import os, platform

def get_filename():
    """Returns the extensionless IDB file name."""
    path = idc.get_input_file_path()
    return os.path.splitext(os.path.basename(path))[0]

def fix_whitespace(literal):
    """Replaces whitespace characters with corresponding literals."""
    if literal is None:
        return literal

    literal = literal.replace(
    "\\", "\\\\"
    ).replace(
    "\t", "\\t"
    ).replace(
    "\r", "\\r"
    ).replace(
    "\n", "\\n"
    ).replace(
    "\b", "\\b"
    ).replace(
    "\v", "\\v"
    ).replace(
    "\a", "\\a"
    ).replace(
    "\f", "\\f"
    ).replace(
    "\x1b", "\\x1B"
    )

    return literal

def export_strings():
    """Exports missing strings (referenced from `paths`) from IDB to SQLite database (`strings`) in the current directory."""
    filename = get_filename() + '.db'
    trailing_slash = "\\" if platform.system() == "Windows" else "/"
    print("Exporting missing strings to " + os.getcwd() + trailing_slash + filename)
    
    conn = sqlite3.connect(filename)
    conn.text_factory = lambda x: x.decode("utf-8")
    c = conn.cursor()

    c.execute('''CREATE TABLE IF NOT EXISTS strings
                (address integer UNIQUE, literal text UNIQUE)''')
    
    try:
        c.execute('''SELECT * FROM (SELECT string_addr, count(*) from paths GROUP BY string_addr) AS pstrings LEFT JOIN strings ON pstrings.string_addr = strings.address WHERE literal IS NULL''')
        missing = c.fetchall()
    except:
        raise Exception("Could not query the missing strings. Make sure you have a populated `paths` table.")

    exported = 0
    duplicates = 0
    for row in missing:
        string_addr = row[0]

        # Try exporting as UTF-16 string
        literal = fix_whitespace(idc.get_strlit_contents(string_addr, -1, 1))

        if literal is None or literal == "":
            # Try exporting as a UTF-8 string
            literal = fix_whitespace(idc.get_strlit_contents(string_addr, -1, 0))
            if literal is None or literal == "":
                print("Error: could not export string at " + str(hex(string_addr)))
                continue

        try:
            c.execute("INSERT INTO strings VALUES (?,?)", (string_addr, literal))
            exported += 1
        # Skip duplicate errors on UNIQUE db constraint ('column X is not unique')
        except sqlite3.IntegrityError:
            duplicates += 1
            pass

    print("Exported " + str(exported) + " unique strings.")
    print("Skipped " + str(duplicates) + " duplicates.")
    conn.commit()
    conn.close()


class Utf16StringExporterPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_PROC
    comment = "dubRE UTF-16 String Exporter"
    help = "This plugin exports missing UTF-16 and UTF-8 strings to SQLite database; `paths` table is required."
    wanted_name = "dubRE UTF-16 String Exporter"
    wanted_hotkey = "Shift+M"

    def init(self):
        return idaapi.PLUGIN_KEEP

    def term(self):
        pass

    def run(self, arg):
        export_strings()


def PLUGIN_ENTRY():
    return Utf16StringExporterPlugin()


repository: hoangcuongflp/Stuff
file: ida-master/plugins/mipsrop/mipsrop.py
description: because I've orderliness obsession :(
wanted_name: MIPS ROP Finder
comment: MIPS ROP Finder
created_at: 2015-10-12 16:35:32+00:00
pushed_at: 2014-12-22 00:32:06+00:00
url:
https://github.com/hoangcuongflp/Stuff/blob/ee100760a05f524ff73837222d80f1dd12da3cee/ida-master/plugins/mipsrop/mipsrop.py

content:
# IDA plugin for identifying ROP gadgets in Linux MIPS binaries.
#
# Return Oriented Programming in Linux MIPS is more like Jump Oriented Programming; the idea is to
# control enough of the stack/registers in order to control various jumps. Since all instructions
# in MIPS must be 4-byte aligned, you cannot "create" new instructions by returning into the middle
# of existing instructions, as is possible with some other architectures.
#
# In any given MIPS function, various registers are saved onto the stack by necessity:
#
#	o $s0 - $s7
#	o $fp
#	o $ra
#
# These values are restored from the stack before the function returns, thus, during a stack overflow
# one can control some or all of these register values. The subroutine registers ($s*) are of particular
# interest, as they are commonly used by the compiler to store function pointers. By convention, gcc will
# move function pointers into the $t9 register, then call the function using jalr:
#
# 	move $t9, $s0  <-- If we control $s0, we control where the jump is taken
#	jalr $t9
#
# While there are other jumps that are of use, and which this plugin searches for, the premise is the same:
# control the stack/registers, and you control various jumps allowing you to chain various blocks of code
# together.
#
# With a list of controllable jumps such as these, we then just need to search the surrounding instructions
# to see if they perform some operation which may be useful. For example, let's say we need to load the
# value 1 into the $a0 register; in this case, we would want to look for a controllable jump such as this:
#
#	move $t9, $s1
#	jalr $t9
#	li $a0, 1    <-- Remember MIPS has jump delay slots, so this instruction is executed with the jump
#
# If we return to this piece of code (and if we control $s1), we can pre-load $s1 with the address of the
# next ROP gadget; thus, $a0 will be loaded with the value 1 and we can chain this block of code with other
# gadgets in order to perform more complex operations.
#
# This plugin finds all potentially controllable jumps, and then allows you to search for desired instructions
# surrounding these controllable jumps. Example:
#
#	Python> mipsrop.find("li $a0, 1")
#	----------------------------------------------------------------------------------------------------
#	|  Address     |  Action                                              |  Control Jump              |
#	----------------------------------------------------------------------------------------------------
#	|  0x0002F0F8  |  li $a0,1                                            |  jalr  $s4                 |
#	|  0x00057E50  |  li $a0,1                                            |  jalr  $s1                 |
#	----------------------------------------------------------------------------------------------------
#
# The output shows the offset of each ROP gadget, the instruction within the gadget that your search matched,
# and the effective register that is jumped to after that instruction is executed.
#
# The specified instruction can be a full instruction, such as the example above, or a partial instruction.
# Regex is supported for any of the instruction mnemonics or operands; for convenience, the dollar signs in 
# front of register names are automatically escaped.
#
# Craig Heffner
# Tactical Network Solutions

import re
import idc
import idaapi
import idautils

# Global instance of MIPSROPFinder
mipsrop = None

class MIPSInstruction(object):
	'''
	Class for storing info about a specific instruction.
	'''

	def __init__(self, mnem, opnd0=None, opnd1=None, opnd2=None, ea=idc.BADADDR):
		self.mnem = mnem
		self.operands = [opnd0, opnd1, opnd2]
		self.opnd0 = opnd0
		self.opnd1 = opnd1
		self.opnd2 = opnd2
		self.ea = ea

	def __str__(self):
		string = self.mnem + " "

		for op in self.operands:
			if op:
				string += "%s," % op
			else:
				break

		return string[:-1]

class ROPGadget(object):
	'''
	Class for storing information about a specific ROP gadget.
	'''
	
	def __init__(self, control, jump, operation=None, description="ROP gaget"):
		self.h = '-' * 112
		self.control = control
		self.exit = jump
		self.operation = operation
		self.description = description
		
		if self.control.opnd1:
			self.control.register = self.control.opnd1
		else:
			self.control.register = self.control.opnd0

		if self.exit.opnd1:
			self.exit.register = self.exit.opnd1
		else:
			self.exit.register = self.exit.opnd0
		
		if self.operation:
			if self.operation.ea < self.control.ea:
				self.entry = self.operation
			else:
				self.entry = self.control
		else:
			self.operation = self.control
			self.entry = self.control


	def header(self):
		return self.h + "\n|  Address     |  Action                                              |  Control Jump                          |\n" + self.h

	def footer(self):
		return self.h

	def __str__(self):
		return "|  0x%.8X  |  %-50s  |  %-5s %-30s  |" % (self.entry.ea, str(self.operation), self.exit.mnem, self.control.register)

class BowcasterBuilder(object):
	'''
	Class to generate bowcaster code from a list of selected ROP gadgets. WIP.
	'''

	INSIZE = 4
	SEARCH_DEPTH = 25

	def __init__(self, gadgets):
		self.code = []
		self.gadgets = gadgets

	def build_code(self):
		keys = self.gadgets.keys()
		keys.sort()
	
		for key in keys[::-1]:
			last_instruction = False
			ea = self.gadgets[key]
			end_ea = ea + self.SEARCH_DEPTH

			while ea <= end_ea:
				mnem = idc.GetMnem(ea)
				if mnem in ['jr', 'jalr']:
					last_instruction = True
				ea += self.INSIZE

	def print_code(self):
		for line in self.code:
			print line

class MIPSROPFinder(object):
	'''
	Primary ROP finder class.
	'''

	CODE = 2
        DATA = 3
	INSIZE = 4
        SEARCH_DEPTH = 25
	
	def __init__(self):
		self.start = idc.BADADDR
		self.end = idc.BADADDR
		self.system_calls = []
		self.double_jumps = []
		self.controllable_jumps = []
		start = 0
		end = 0

		for (start, end) in self._get_segments(self.CODE):
			self.controllable_jumps += self._find_controllable_jumps(start, end)
			self.system_calls += self._find_system_calls(start, end)
			self.double_jumps += self._find_double_jumps(start, end)
			if self.start == idc.BADADDR:
				self.start = start
		self.end = end
	
		if self.controllable_jumps or self.system_calls:
			print "MIPS ROP Finder activated, found %d controllable jumps between 0x%.8X and 0x%.8X" % (len(self.controllable_jumps), self.start, self.end)
		
	def _get_segments(self, attr):
		segments = []
		start = idc.BADADDR
		end = idc.BADADDR
		seg = idc.FirstSeg()

		while seg != idc.BADADDR:
			if idc.GetSegmentAttr(seg, idc.SEGATTR_TYPE) == attr:
				start = idc.SegStart(seg)
				end = idc.SegEnd(seg)
				segments.append((start, end))
			seg = idc.NextSeg(seg)

		return segments

	def _get_instruction(self, ea):
		return MIPSInstruction(idc.GetMnem(ea), idc.GetOpnd(ea, 0), idc.GetOpnd(ea, 1), idc.GetOpnd(ea, 2), ea)

	def _does_instruction_match(self, ea, instruction, regex=False):
		i = 0
		op_cnt = 0
		op_ok_cnt = 0
		match = False
		ins_size = idaapi.decode_insn(ea)
		mnem = GetMnem(ea)

		if (not instruction.mnem) or (instruction.mnem == mnem) or (regex and re.match(instruction.mnem, mnem)):
			for operand in instruction.operands:
				if operand:
					op_cnt += 1
					op = idc.GetOpnd(ea, i)

					if regex:
						if re.match(operand, op):
							op_ok_cnt += 1
					elif operand == op:
						op_ok_cnt += 1
				i += 1

			if op_cnt == op_ok_cnt:
				match = True

		return match

	def _is_bad_instruction(self, ea, bad_instructions=['j', 'b'], no_clobber=[]):
		bad = False
		mnem = GetMnem(ea)

		if mnem and mnem[0] in bad_instructions:
			bad = True
		else:
			for register in no_clobber:
				if (idaapi.insn_t_get_canon_feature(idaapi.cmd.itype) & idaapi.CF_CHG1) == idaapi.CF_CHG1:
					if idc.GetOpnd(ea, 0) == register:
						bad = True

		return bad

	def _contains_bad_instruction(self, start_ea, end_ea, bad_instructions=['j', 'b'], no_clobber=[]):
		ea = start_ea

		while ea <= end_ea:
			if self._is_bad_instruction(ea, bad_instructions, no_clobber):
				return True
			else:
				ea += self.INSIZE

		return False
		
	def _find_prev_instruction_ea(self, start_ea, instruction, end_ea=0, no_baddies=True, regex=False, dont_overwrite=[]):
		instruction_ea = idc.BADADDR
		ea = start_ea
		baddies = ['j', 'b']

		while ea >= end_ea:
			if self._does_instruction_match(ea, instruction, regex):
				instruction_ea = ea
				break
			elif no_baddies and self._is_bad_instruction(ea, no_clobber=dont_overwrite):
				break

			ea -= self.INSIZE

		return instruction_ea

	def _find_next_instruction_ea(self, start_ea, instruction, end_ea=idc.BADADDR, no_baddies=False, regex=False, dont_overwrite=[]):
		instruction_ea = idc.BADADDR
		ea = start_ea

		while ea <= end_ea:
			if self._does_instruction_match(ea, instruction, regex):
				instruction_ea = ea
				break
			elif no_baddies and self._is_bad_instruction(ea, no_clobber=dont_overwrite):
				break

			ea += self.INSIZE

		return instruction_ea

	def _find_controllable_jumps(self, start_ea, end_ea):
		controllable_jumps = []
		t9_controls = [
			MIPSInstruction("move", "\$t9"),
			MIPSInstruction("addiu", "\$t9", "^\$"),
		]
		t9_jumps = [
			MIPSInstruction("jalr", "\$t9"),
			MIPSInstruction("jr", "\$t9"),
		]
		ra_controls = [
			MIPSInstruction("lw", "\$ra"),
		]
		ra_jumps = [
			# TODO: Search for jumps to registers other than $ra.
			MIPSInstruction("jr", "\$ra"),
		]
		t9_musnt_clobber = ["$t9"]
		ra_musnt_clobber = ["$ra"]

		for possible_control_instruction in t9_controls+ra_controls:
			ea = start_ea
			found = 0

			if possible_control_instruction in t9_controls:
				jumps = t9_jumps
				musnt_clobber = t9_musnt_clobber
			else:
				jumps = ra_jumps
				musnt_clobber = ra_musnt_clobber

			while ea <= end_ea:

				ea = self._find_next_instruction_ea(ea, possible_control_instruction, end_ea, regex=True)
				if ea != idc.BADADDR:
					ins_size = idaapi.decode_insn(ea)

					control_instruction = self._get_instruction(ea)
					control_register = control_instruction.operands[1]
					
					if control_register:
						for jump in jumps:
							jump_ea = self._find_next_instruction_ea(ea+ins_size, jump, end_ea, no_baddies=True, regex=True, dont_overwrite=musnt_clobber)
							if jump_ea != idc.BADADDR:
								jump_instruction = self._get_instruction(jump_ea)
								controllable_jumps.append(ROPGadget(control_instruction, jump_instruction, description="Controllable Jump"))
								ea = jump_ea
					
					ea += ins_size

		return controllable_jumps

	def _find_system_calls(self, start_ea, end_ea):
		system_calls = []
		system_load = MIPSInstruction("la", "$t9", "system")
		stack_arg_zero = MIPSInstruction("addiu", "$a0", "$sp")

		for xref in idautils.XrefsTo(idc.LocByName('system')):
			ea = xref.frm
			if ea >= start_ea and ea <= end_ea and idc.GetMnem(ea)[0] in ['j', 'b']:
				a0_ea = self._find_next_instruction_ea(ea+self.INSIZE, stack_arg_zero, ea+self.INSIZE)
				if a0_ea == idc.BADADDR:
					a0_ea = self._find_prev_instruction_ea(ea, stack_arg_zero, ea-(self.SEARCH_DEPTH*self.INSIZE))
				
				if a0_ea != idc.BADADDR:
					control_ea = self._find_prev_instruction_ea(ea-self.INSIZE, system_load, ea-(self.SEARCH_DEPTH*self.INSIZE))
					if control_ea != idc.BADADDR:
						system_calls.append(ROPGadget(self._get_instruction(control_ea), self._get_instruction(ea), self._get_instruction(a0_ea), description="System call"))

				ea += self.INSIZE
			else:
				break

		return system_calls

	def _find_double_jumps(self, start_ea, end_ea):
		double_jumps = []
		
		for i in range(0, len(self.controllable_jumps)):
			g1 = self.controllable_jumps[i]
			if g1.exit.mnem != 'jalr':
				continue

			for j in range(i+1, len(self.controllable_jumps)):
				g2 = self.controllable_jumps[j]
				distance = (g2.entry.ea - g1.exit.ea)

				if distance > 0 and distance <= (self.SEARCH_DEPTH * self.INSIZE):
					if g1.control.register != g2.control.register:
						if not self._contains_bad_instruction(g1.exit.ea+self.INSIZE, g2.control.ea-self.INSIZE, no_clobber=[g2.control.register]):
							double_jumps.append(g1)
							break

		return double_jumps

	def _find_rop_gadgets(self, gadget):
		gadget_list = []

		for controllable_jump in self.controllable_jumps:
			gadget_ea = idc.BADADDR

			ea = self._find_next_instruction_ea(controllable_jump.entry.ea, gadget, controllable_jump.exit.ea+self.INSIZE, regex=True)
			if ea != idc.BADADDR:
				gadget_ea = ea
			else:
				ea = self._find_prev_instruction_ea(controllable_jump.entry.ea, gadget, controllable_jump.entry.ea-(self.SEARCH_DEPTH*self.INSIZE), no_baddies=True, regex=True, dont_overwrite=[controllable_jump.entry.opnd1])
				if ea != idc.BADADDR:
					gadget_ea = ea
		
			if gadget_ea != idc.BADADDR:
				gadget_list.append(ROPGadget(controllable_jump.entry, controllable_jump.exit, self._get_instruction(gadget_ea)))

		return gadget_list

	def _print_gadgets(self, gadgets):
		if gadgets:
			print gadgets[0].header()

		for gadget in gadgets:
			print str(gadget)

		if gadgets:
			print gadgets[0].footer()
		
		print "Found %d matching gadgets" % (len(gadgets))

	def _get_marked_gadgets(self):
		rop_gadgets = {}

		for i in range(1, 1024):
			marked_pos = idc.GetMarkedPos(i)
			if marked_pos != idc.BADADDR:
				marked_comment = idc.GetMarkComment(i)
				if marked_comment and marked_comment.lower().startswith("rop"):
					rop_gadgets[marked_comment] = marked_pos
			else:
				break

		return rop_gadgets

	def double(self):
		self.doubles()

	def doubles(self):
		'''
		Prints a list of all "double jump" gadgets (useful for function calls).
		'''
		self._print_gadgets(self.double_jumps)

	def stackfinder(self):
		self.stackfinders()

	def stackfinders(self):
		'''
		Prints a list of all gadgets that put a stack address into a register.
		'''
		self.find("addiu .*, $sp")

	def lia0(self):
		'''
		Prints a list of all gadgets that load an immediate value number into $a0 (useful for setting up the argument to sleep).
		'''
		self.find("li $a0")

	def tail(self):
		return self.tails()

	def tails(self):
		'''
		Prints a lits of all tail call gadgets (useful for function calls).
		'''
		return self.iret()

	def iret(self):
		'''
		Prints a lits of all tail gadgets (useful for function calls).
		'''
		tail_gadgets = []

		for gadget in self._find_rop_gadgets(MIPSInstruction("move", "\$t9")):
			if gadget.exit.mnem == 'jr' and gadget.exit.register == '$t9':
				tail_gadgets.append(gadget)

		self._print_gadgets(tail_gadgets)

	def system(self):
		'''
		Prints a list of gadgets that may be used to call system().
		'''
		sys_gadgets = self.system_calls + self._find_rop_gadgets(MIPSInstruction("addiu", "\$a0", "\$sp"))
		self._print_gadgets(sys_gadgets)

	def find(self, instruction_string=""):
		'''
		Locates all potential ROP gadgets that contain the specified instruction.

		@instruction_string - The instruction you need executed. This can be either a:

					o Full instruction    - "li $a0, 1"
					o Partial instruction - "li $a0"
					o Regex instruction   - "li $a0, .*"
		'''
		registers = ['$v', '$s', '$a', '$t', '$k', '$pc', '$fp', '$ra', '$gp', '$at', '$zero']

		comma_split = instruction_string.split(',')
		instruction_parts = comma_split[0].split()
		if len(comma_split) > 1:
			instruction_parts += comma_split[1:]

		for i in range(0, 4):
			if i > len(instruction_parts) - 1:
				instruction_parts.append(None)
			else:
				instruction_parts[i] = instruction_parts[i].strip().strip(',').strip()
				for reg in registers:
					instruction_parts[i] = instruction_parts[i].replace(reg, "\\%s" % reg)

		instruction = MIPSInstruction(instruction_parts[0], instruction_parts[1], instruction_parts[2], instruction_parts[3])
		gadgets = self._find_rop_gadgets(instruction)
		if gadgets:
			self._print_gadgets(gadgets)
		else:
			print "No ROP gadgets found!"

	def summary(self):
		'''
		Prints a summary of your currently marked ROP gadgets, in alphabetical order by the marked name.
		To mark a location as a ROP gadget, simply mark the position in IDA (Alt+M) with any name that starts with "ROP".
		'''
		rop_gadgets = self._get_marked_gadgets()
		summaries = []
		delim_char = "-"
		headings = {
			'name' 		: "Gadget Name",
			'offset'	: "Gadget Offset",
			'summary'	: "Gadget Summary"
		}
		lengths = {
			'name'		: len(headings['name']),
			'offset'	: len(headings['offset']),
			'summary'	: len(headings['summary']),
		}
		total_length = (3 * len(headings)) + 1

		if rop_gadgets:
			gadget_keys = rop_gadgets.keys()
			gadget_keys.sort()

			for marked_comment in gadget_keys:
				if len(marked_comment) > lengths['name']:
					lengths['name'] = len(marked_comment)

				summary = []
				ea = rop_gadgets[marked_comment]
				end_ea = ea + (self.SEARCH_DEPTH * self.INSIZE)

				while ea <= end_ea:
					summary.append(idc.GetDisasm(ea))
					mnem = idc.GetMnem(ea)
					if mnem[0].lower() in ['j', 'b']:
						summary.append(idc.GetDisasm(ea+self.INSIZE))
						break

					ea += self.INSIZE

				if len(summary) == 0:
					summary.append('')

				for line in summary:
					if len(line) > lengths['summary']:
						lengths['summary'] = len(line)

				summaries.append(summary)

			for (heading, size) in lengths.iteritems():
				total_length += size

			delim = delim_char * total_length
			line_fmt = "| %%-%ds | %%-%ds | %%-%ds |" % (lengths['name'], lengths['offset'], lengths['summary'])

			print delim
			print line_fmt % (headings['name'], headings['offset'], headings['summary'])
			print delim
			
			for i in range(0, len(gadget_keys)):
				line_count = 0
				marked_comment = gadget_keys[i]
				offset = "0x%.8X" % rop_gadgets[marked_comment]
				summary = summaries[i]
				
				for line in summary:
					if line_count == 0:
						print line_fmt % (marked_comment, offset, line)
					else:
						print line_fmt % ('', '', line)

					line_count += 1

				print delim

	def build(self):
		'''
		WIP.
		'''
		gadgets = self._get_marked_gadgets()
		bc = BowcasterBuilder(gadgets)
		bc.build_code()

	def help(self):
		'''
		Show help info.
		'''
		delim = "---------------------------------------------------------------------" * 2
		
		print ""
		print "mipsrop.find(instruction_string)"
		print delim
		print self.find.__doc__

		print ""
		print "mipsrop.system()"
		print delim
		print self.system.__doc__

		print ""
		print "mipsrop.doubles()"
		print delim
		print self.doubles.__doc__

		print ""
		print "mipsrop.stackfinders()"
		print delim
		print self.stackfinders.__doc__

		print ""
		print "mipsrop.tails()"
		print delim
		print self.tails.__doc__

		print ""
		print "mipsrop.summary()"
		print delim
		print self.summary.__doc__


class mipsropfinder_t(idaapi.plugin_t):
	flags = 0
	comment = "MIPS ROP Finder"
	help = ""
	wanted_name = "MIPS ROP Finder"
	wanted_hotkey = ""

	def init(self):
		self.menu_context = idaapi.add_menu_item("Search/", "mips rop gadgets", "Alt-1", 0, self.run, (None,))
		return idaapi.PLUGIN_KEEP

	def term(self):
		idaapi.del_menu_item(self.menu_context)
		return None

	def run(self, arg):
		global mipsrop
		mipsrop = MIPSROPFinder()
                
def PLUGIN_ENTRY():
        return mipsropfinder_t()

# DEBUG
#if __name__ == '__main__':
#	mipsrop = MIPSROPFinder()


repository: r0ysue/FridaAndrioidNativeBeginnersBook
file: Chap13/各种插件/findcrypt3/findcrypt-yara-master/findcrypt3.py
description: 《安卓Frida SO逆向分析实战》
wanted_name: Findcrypt
comment: Findcrypt plugin for IDA Pro (using yara framework)
created_at: 2023-05-31 07:37:25+00:00
pushed_at: 2023-10-22 13:17:28+00:00
url:
https://github.com/r0ysue/FridaAndrioidNativeBeginnersBook/blob/6be19b07269fb91f5e054e5485ce327a1bf07a60/Chap13/%E5%90%84%E7%A7%8D%E6%8F%92%E4%BB%B6/findcrypt3/findcrypt-yara-master/findcrypt3.py

content:
# -*- coding: utf-8 -*-

import idaapi
import idautils
import ida_bytes
import ida_diskio
import idc
import operator
import yara
import os
import glob

VERSION = "0.2"
YARARULES_CFGFILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "findcrypt3.rules")

try:
    class Kp_Menu_Context(idaapi.action_handler_t):
        def __init__(self):
            idaapi.action_handler_t.__init__(self)

        @classmethod
        def get_name(self):
            return self.__name__

        @classmethod
        def get_label(self):
            return self.label

        @classmethod
        def register(self, plugin, label):
            self.plugin = plugin
            self.label = label
            instance = self()
            return idaapi.register_action(idaapi.action_desc_t(
                self.get_name(),  # Name. Acts as an ID. Must be unique.
                instance.get_label(),  # Label. That's what users see.
                instance  # Handler. Called when activated, and for updating
            ))

        @classmethod
        def unregister(self):
            """Unregister the action.
            After unregistering the class cannot be used.
            """
            idaapi.unregister_action(self.get_name())

        @classmethod
        def activate(self, ctx):
            # dummy method
            return 1

        @classmethod
        def update(self, ctx):
            if ctx.form_type == idaapi.BWN_DISASM:
                return idaapi.AST_ENABLE_FOR_WIDGET
            return idaapi.AST_DISABLE_FOR_WIDGET

    class Searcher(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.search()
            return 1

except:
    pass


p_initialized = False


class YaraSearchResultChooser(idaapi.Choose):
    def __init__(self, title, items, flags=0, width=None, height=None, embedded=False, modal=False):
        idaapi.Choose.__init__(
            self,
            title,
            [
                ["Address", idaapi.Choose.CHCOL_HEX|10],
                ["Rules file", idaapi.Choose.CHCOL_PLAIN|12],
                ["Name", idaapi.Choose.CHCOL_PLAIN|25],
                ["String", idaapi.Choose.CHCOL_PLAIN|25],
                ["Value", idaapi.Choose.CHCOL_PLAIN|40],
            ],
            flags=flags,
            width=width,
            height=height,
            embedded=embedded)
        self.items = items
        self.selcount = 0
        self.n = len(items)

    def OnClose(self):
        return

    def OnSelectLine(self, n):
        self.selcount += 1
        idc.jumpto(self.items[n][0])

    def OnGetLine(self, n):
        res = self.items[n]
        res = [idc.atoa(res[0]), res[1], res[2], res[3], res[4]]
        return res

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

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

#--------------------------------------------------------------------------
# Plugin
#--------------------------------------------------------------------------
class Findcrypt_Plugin_t(idaapi.plugin_t):
    comment = "Findcrypt plugin for IDA Pro (using yara framework)"
    help = "todo"
    wanted_name = "Findcrypt"
    wanted_hotkey = "Ctrl-Alt-F"
    flags = idaapi.PLUGIN_KEEP

    def init(self):
        global p_initialized

        # register popup menu handlers
        try:
            Searcher.register(self, "Findcrypt")
        except:
            pass

        if p_initialized is False:
            p_initialized = True
            self.user_directory = self.get_user_directory()
            idaapi.register_action(idaapi.action_desc_t(
                "Findcrypt",
                "Find crypto constants",
                Searcher(),
                None,
                None,
                0))
            idaapi.attach_action_to_menu("Search", "Findcrypt", idaapi.SETMENU_APP)
            print("=" * 80)
            print("Findcrypt v{0} by David BERARD, 2017".format(VERSION))
            print("Findcrypt search shortcut key is Ctrl-Alt-F")
            print("Global rules in %s" % YARARULES_CFGFILE)
            print("User-defined rules in %s/*.rules" % self.user_directory)
            print("=" * 80)

        return idaapi.PLUGIN_KEEP

    def term(self):
        pass


    def toVirtualAddress(self, offset, segments):
        va_offset = 0
        for seg in segments:
            if seg[1] <= offset < seg[2]:
                va_offset = seg[0] + (offset - seg[1])
        return va_offset


    def get_user_directory(self):
        user_dir = ida_diskio.get_user_idadir()
        plug_dir = os.path.join(user_dir, "plugins")
        res_dir = os.path.join(plug_dir, "findcrypt-yara")
        if not os.path.exists(res_dir):
            os.makedirs(res_dir, 0o755)
        return res_dir


    def get_rules_files(self):
        rules_filepaths = {"global":YARARULES_CFGFILE}
        for fpath in glob.glob(os.path.join(self.user_directory, "*.rules")):
            name = os.path.basename(fpath)
            rules_filepaths.update({name:fpath})
        return rules_filepaths


    def search(self):
        memory, offsets = self._get_memory()
        rules = yara.compile(filepaths=self.get_rules_files())
        values = self.yarasearch(memory, offsets, rules)
        c = YaraSearchResultChooser("Findcrypt results", values)
        r = c.show()

    def yarasearch(self, memory, offsets, rules):
        print(">>> start yara search")
        values = list()
        matches = rules.match(data=memory)
        for match in matches:
            for string in match.strings:
                name = match.rule
                if name.endswith("_API"):
                    try:
                        name = name + "_" + idc.GetString(self.toVirtualAddress(string[0], offsets))
                    except:
                        pass
                value = [
                    self.toVirtualAddress(string[0], offsets),
                    match.namespace,
                    name + "_" + hex(self.toVirtualAddress(string[0], offsets)).lstrip("0x").rstrip("L").upper(),
                    string[1],
                    repr(string[2]),
                ]
                idaapi.set_name(value[0], name
                             + "_"
                             + hex(self.toVirtualAddress(string[0], offsets)).lstrip("0x").rstrip("L").upper()
                             , 0)
                values.append(value)
        print("<<< end yara search")
        return values

    def _get_memory(self):
        result = bytearray()
        segment_starts = [ea for ea in idautils.Segments()]
        offsets = []
        start_len = 0
        for start in segment_starts:
            end = idc.get_segm_attr(start, idc.SEGATTR_END)
            result += ida_bytes.get_bytes(start, end - start)
            offsets.append((start, start_len, len(result)))
            start_len = len(result)
        return bytes(result), offsets

    def run(self, arg):
        self.search()


# register IDA plugin
def PLUGIN_ENTRY():
    return Findcrypt_Plugin_t()


repository: novafacing/ghidra-decompiler
file: ghidra_9.1.2_PUBLIC/Extensions/IDAPro/Python/6xx/plugins/xmlexp.py
description: My experiments with getting the ghidra decompiler to run on other platforms
wanted_name: XML Exporter
comment: Export database as XML file
created_at: 2020-09-08 18:10:43+00:00
pushed_at: 2021-02-24 20:05:26+00:00
url:
https://github.com/novafacing/ghidra-decompiler/blob/5f0ef539ffa92fbc13175c996b7a5381a263c328/ghidra_9.1.2_PUBLIC/Extensions/IDAPro/Python/6xx/plugins/xmlexp.py

content:
#---------------------------------------------------------------------
# xmlexp.py - IDA XML Exporter plugin
#---------------------------------------------------------------------
"""
Plugin for IDA which exports a XML PROGRAM document file from a database.
"""

import idaapi
import idautils
import idc
import datetime
import sys
import time


XML_EXPORTER_VERSION = "4.1.2"
BASELINE_IDA_VERSION = 620
IDA_SDK_VERSION = idaapi.IDA_SDK_VERSION
BADADDR = idaapi.BADADDR


class XmlExporterPlugin(idaapi.plugin_t):
    """
    XML Exporter plugin class
    """
    flags = 0
    comment = "Export database as XML file"
    help = "Export database as XML <PROGRAM> document"
    wanted_name = "XML Exporter"
    wanted_hotkey = "Ctrl-Shift-x"


    def init(self):
        """
        init function for XML Exporter 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 Exporter plugin (xmlexp.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 Exporter plugin.
        
        Args:
            arg: Integer, non-zero value enables auto-run feature for
                IDA batch (no gui) processing mode. Default is 0.
        """
        xml = XmlExporter(arg)
        try:
            try:
                xml.export_xml()
            except Cancelled:
                idaapi.hide_wait_box()
                msg = "XML Export cancelled!"
                print "\n" + msg
                idaapi.warning(msg)
            except:
                idaapi.hide_wait_box()
                msg = "***** Exception occurred: XML Exporter failed! *****"
                print "\n" + msg + "\n", sys.exc_type, sys.exc_value
                idaapi.warning(msg)
        finally:
            xml.cleanup()
            idaapi.setStat(xml.state)


    def term(self):
        pass


def PLUGIN_ENTRY():
    return XmlExporterPlugin()


class Cancelled(Exception):
    pass


class FileError(Exception):
    pass


class XmlExporter:
    """
    XML Exporter contains methods to export an IDA database as a
        XML PROGRAM document.
    """
    def __init__(self, arg):
        """
        Initializes the XmlExporter attributes

        Args:
            arg: Integer, non-zero value enables auto-run feature for
                IDA batch (no gui) processing mode. Default is 0.
        """
        self.debug = False
        self.autorun = False
        if arg != 0:
            self.autorun = True
        self.state = idaapi.setStat(idaapi.st_Work)
        self.indent_level = 0
        self.seg_addr = False
        self.has_overlays = False
        self.hexrays = False
        self.options = None
        self.xmlfile = 0
        self.elements = {}
        self.counters = []
        self.tags = []
        
        # initialize class variables from database
        self.inf = idaapi.get_inf_structure()
        self.min_ea = self.inf.minEA
        self.max_ea = self.inf.maxEA
        self.cbsize = (idaapi.ph_get_cnbits()+7)/8
        self.processor = str.upper(idaapi.get_idp_name())
        self.batch = idaapi.cvar.batch


    def export_xml(self):
        """
        Exports the IDA database to a XML PROGRAM document file.
        """
        self.check_and_load_decompiler()
        self.display_xml_exporter_version()
        self.display_database_info()
        
        self.get_options()
    
        if (self.autorun == True):
            (self.filename, ext) = os.path.splitext(idaapi.cvar.database_idb)
            self.filename += ".xml"
        else:
            self.filename=idaapi.askfile_c(1, "*.xml",
                                           "Enter name of export xml file:")
            
        if self.filename == None or len(self.filename) == 0:
            raise Cancelled
        self.xmlfile = self.open_file(self.filename, "w")
        
        idaapi.show_wait_box("Exporting XML <PROGRAM> document ....")
        idaapi.msg("\n------------------------------------------------" +
                   "-----------")
        idaapi.msg("\nExporting XML <PROGRAM> document ....")
        begin = time.clock()
        
        self.write_xml_declaration()
        self.export_program()
        
        # export database items based on options
        if (self.options.DataTypes.checked         == True or
            self.options.DataDefinitions.checked   == True or 
            self.options.Functions.checked         == True):
            self.export_datatypes()
        if (self.options.MemorySections.checked    == True or
            self.options.MemoryContent.checked     == True):
            self.export_memory_map()
        if (self.options.RegisterValues.checked    == True):
            self.export_register_values()   
        if (self.options.CodeBlocks.checked        == True):
            self.export_code()  
        if (self.options.DataDefinitions.checked   == True):
            self.export_data()  
        if (self.options.Comments.checked          == True):
            self.export_comments()
            self.export_bookmarks()     
        if (self.options.EntryPoints.checked       == True):
            self.export_program_entry_points()  
        if (self.options.Symbols.checked           == True):
            self.export_symbol_table()  
        if (self.options.Functions.checked         == True):
            self.export_functions() 
        if (self.options.MemoryReferences.checked  == True or 
            self.options.StackReferences.checked   == True or
            self.options.Manual.checked            == True or
            self.options.DataTypes.checked         == True):
            self.export_markup()    
        self.end_element(PROGRAM)
        
        idaapi.msg('\n%35s' % 'Total ')
        self.display_cpu_time(begin)
        idaapi.hide_wait_box()  
        self.display_summary()


    def check_and_load_decompiler(self):
        """
        Checks for the presence of a decompiler plugin for the database.
        
        Note: The decompiler must be loaded by the XML Exporter plugin
            if it is running in batch mode. IDA will load the decompiler
            plugin automatically if not in batch mode.
        
        Note: There was no support for decompiler plugins in IDAPython until
            IDA 6.6, so skip if this is an older version.
        
        Note: Currently the 4 decompiler plugins for the  x86, x64,
            ARM32, and ARM64 are supported.
        """
        if self.batch == 0 or IDA_SDK_VERSION < 660:
            return
        plugin = ''
        if self.processor == 'PC':
            if self.inf.is_64bit():
                plugin = "hexx64"
            elif self.inf.is_32bit():
                plugin = 'hexrays'
        elif self.processor == 'ARM':
            if self.inf.is_64bit():
                plugin = "hexarm64"
            elif self.inf.is_32bit():
                plugin = "hexarm"
        if len(plugin) > 0:
            try:
                idaapi.load_plugin(plugin)
                self.hexrays = idaapi.init_hexrays_plugin()
            except:
                return
    

    def check_char(self, ch):
        """
        Replaces a special XML character with an entity string.
        
        Args:
            ch: String containing the character to check.
            
        Returns:
            String containing either the character or the entity
            substition string.
        """
        if ((ord(ch) < 0x20) and (ord(ch) != 0x09 and
             ord(ch) != 0x0A and ord(ch) != 0x0D)): return ''
        elif ch == '&' :  return '&amp;'
        elif ch == '<' :  return "&lt;"
        elif ch == '>' :  return "&gt;"
        elif ch == '\'' : return "&apos;"
        elif ch == '"' :  return "&quot;"
        elif ch == '\x7F': return ''
        elif ord(ch) > 0x7F: return '&#x' + format((ord(ch),"x")) + ";"
        return ch
    

    def check_for_entities(self, text):
        """
        Checks all characters in a string for special XML characters.
        
        Args:
            text: String to check for special XML characters.
            
        Returns:
            String containing original string with substitutions for
                any special XML characters.
        """
        new = ''
        for c in text:
            new += self.check_char(c)
        return new
    

    def check_if_seg_contents(self, seg):
        """
        Determines if any address in a segment contains a value.
        
        Args:
            seg: IDA segment object
            
        Returns:
            True if any address in a segment contains a value.
            False if no address in a segment contains a value.
        """
        for addr in idautils.Heads(seg.startEA, seg.endEA):
            if idaapi.hasValue(idaapi.getFlags(addr)) == True:
                return True
        return False


    def check_stack_frame(self, sframe):
        """
        Determines if stack frame contains any parameters or local variables.
        
        Args:
            sframe: IDA stack frame for a function.
            
        Returns:
            True if stack frame has parameters or local variables.
            False if stack frame has no parameters or local variables.
        """
        n = sframe.memqty
        for i in range(n):
            member = sframe.get_member(i)
            if member == None:
                continue
            mname = idaapi.get_member_name(member.id)
            if mname != None and len(mname) > 0: 
                if mname != " s" and mname != " r":
                    return True
        return False


    def cleanup(self):
        """
        Frees memory and closes message box and XML file at termination.
        """
        if self.options != None:
            self.options.Free()
        idaapi.hide_wait_box()
        self.close_xmlfile()


    def close_binfile(self):
        """
        Closes the binary data file for the XML Exporter.
        """
        if self.binfile != 0:
            self.binfile.close()
            self.binfile = 0


    def close_tag(self, has_contents=False):
        """
        Closes the start tag for an XML element.
        
        Args:
            has_contents: Boolean indicating if the element has
            sub-elements or text.
        """
        if has_contents:
            self.write_to_xmlfile(">")
            self.indent_level += 1
        else:
            self.write_to_xmlfile(" />")


    def close_xmlfile(self):
        """
        Closes the XML data file for the XML Exporter.
        """
        if self.xmlfile != 0:
            self.xmlfile.close()
            self.xmlfile = 0


    def dbg(self, message):
        """
        Outputs debug message if debug flag is enabled.
        
        Args:
            message: String containing the debug message.
        """
        if (self.debug == True):
            idaapi.msg(message)


    def display_cpu_time(self, start):
        """
        Displays the elapsed CPU time since the start time.
        
        Args:
            start: Floating-point value representing start time in seconds.
        """
        idaapi.msg('CPU time: %6.4f' % (time.clock() - start))
        

    def display_database_info(self):
        """
        Displays information about the database in IDA output window.
        """
        idaapi.msg('\nProcessor Module: ' + self.processor)


    def display_message(self, message):
        """
        Displays an indented message to IDA output window.
        
        Args:
            message: String containing message to display.
        """
        line = "    " + message
        idaapi.msg('\n%-35s' % line)
        

    def display_summary(self):
        """
        Displays summary of exported PROGRAM document in IDA output window.
        """
        fileline = '\nFile: %s' % self.filename
        summary  = "\n--------------------------------------"
        total = 0
        tags = self.tags[:]
        for tag in tags:
            count = self.counters[self.elements[tag]]
            summary += "\n%-22s %8d" % (tag, count)
            total += count
        summary += ("\n%-22s %8d" % ("Total XML Elements:",total))
        idaapi.msg(summary)
        idaapi.msg("\n\nDatabase exported to: %s" % self.filename)
        idaapi.msg("\n--------------------------------------------" +
                   "---------------")
        if self.autorun == False:
            frmt  = "TITLE XML Export Successful!\n"
            frmt += "ICON INFO\n"
            frmt += "AUTOHIDE NONE\n"
            frmt += "HIDECANCEL\n"
            details = '\nSee output window for details...'
            idaapi.info("%s" % (frmt + fileline + details))


    def display_xml_exporter_version(self):
        """
        Displays XML Exporter plugin version info in IDA output window.
        """
        plugin = idaapi.idadir(idaapi.PLG_SUBDIR) + '/xmlexp.py'
        plugintime = time.localtime(os.path.getmtime(plugin))
        # ts = time.strftime('%Y-%m-%d %H:%M:%S', plugintime)
        ts = time.strftime('%b %d %Y %H:%M:%S', plugintime)
        version = "\nXML Exporter Version " + XML_EXPORTER_VERSION
        version += " : SDK " + str(IDA_SDK_VERSION)
        version += " : Python : "+ ts + '\n'
        idaapi.msg(version)
    

    def end_element(self, tag, newline=True):
        """
        Writes the element end tag to the XML file.
        
        Args:
            tag: String containing the element name.
            newline: Boolean indicating if end tag should go on new line.
        """
        self.indent_level -= 1
        if newline:
            start = '\n' + ("    " * self.indent_level)
        else:
            start = ''
        self.write_to_xmlfile(start  + "</" + tag + ">")


    '''
    # BIT_MASK not currently supported for ENUM
    def export_bitmask(self, eid, mask):
        """
        Exports an enum bitmask member as BIT_MASK element.
        
        Args:
            eid: Integer representing the IDA enum id
            mask: Integer representing the IDA enum mask value
        """
        name = idaapi.get_bmask_name(eid, mask)
        if name == None:
            return
        self.start_element(BIT_MASK)
        self.write_attribute(NAME, name)
        self.write_numeric_attribute(VALUE, mask)
        regcmt = idaapi.get_bmask_cmt(eid, mask, False)
        rptcmt = idaapi.get_bmask_cmt(eid, mask, True)
        has_comment =  regcmt != None or rptcmt != None
        self.close_tag(has_comment)
        if regcmt != None and len(regcmt) > 0:
            self.export_regular_cmt(regcmt)
        if rptcmt != None and len(rptcmt) > 0:
            self.export_repeatable_cmt(rptcmt)
        if (has_comment):
            self.end_element(BIT_MASK)
    '''


    def export_bookmarks(self):
        """
        Exports marked location descriptions as BOOKMARK elements.
        """
        found = False
        timer = time.clock()
        if IDA_SDK_VERSION > 690:
            import ida_moves
            import ida_pro
        for slot in range(1,1025):
            if IDA_SDK_VERSION <= 690:
                address = idc.GetMarkedPos(slot)
                description = idc.GetMarkComment(slot)
            else:
                curloc = ida_moves.curloc()
                intp = ida_pro.int_pointer()
                intp.assign(slot)
                address = curloc.markedpos(intp)
                description = curloc.markdesc(slot)
            if address == BADADDR:
                continue
            if description == None:
                continue
            if found == False:
                found = True
                self.update_status(BOOKMARKS)
                self.start_element(BOOKMARKS, True)
            self.start_element(BOOKMARK)
            self.write_address_attribute(ADDRESS, address)
            self.write_attribute(DESCRIPTION, description)
            self.close_tag()
        if found:
            self.end_element(BOOKMARKS)
            self.display_cpu_time(timer)


    def export_c_comments(self):
        """
        Exports block and end-of-line comments entered in the decompiler
        interface.
        """
        if IDA_SDK_VERSION < 660 or self.hexrays == False:
            return
        functions = idautils.Functions()
        if functions == None:
            return
        for addr in functions:
            try:
                if idaapi.is_spec_ea(addr):
                    continue
                ccmts = idaapi.restore_user_cmts(addr)
                if ccmts == None:
                    continue
                p = idaapi.user_cmts_begin(ccmts)
                while p != idaapi.user_cmts_end(ccmts):
                    cmk = idaapi.user_cmts_first(p)
                    cmv = idaapi.user_cmts_second(p)
                    if cmk.itp < (idaapi.ITP_COLON+1):
                        self.export_comment(cmk.ea, "end-of-line", cmv.c_str())
                    else:
                        self.export_comment(cmk.ea, "pre", cmv.c_str())
                    p=idaapi.user_cmts_next(p)
                idaapi.user_cmts_free(ccmts)
            except:
                continue


    def export_code(self):
        """
        Exports the address ranges of code sequences as CODE_BLOCK(s)
        with START and END address attributes.
        """
        addr = self.min_ea
        if idaapi.isCode(idaapi.getFlags(addr)) == False:
            addr = idaapi.nextthat(addr, self.max_ea, idaapi.isCode)
        if (addr == BADADDR):
            return
        self.update_status(CODE)
        timer = time.clock()
        data = idaapi.nextthat(addr, self.inf.maxEA, isData)
        unknown = idaapi.next_unknown(addr, self.inf.maxEA)
        self.start_element(CODE, True)
        while (addr != BADADDR):
            start = addr
            end = min(data, unknown)
            if (end == BADADDR):
                if (idaapi.getseg(start).endEA < self.inf.maxEA):
                    codeend = idaapi.getseg(start).endEA - 1
                    addr = idaapi.getseg(idaapi.nextaddr(codeend)).startEA
                    if idaapi.isCode(idaapi.getFlags(addr)) == False:
                        addr = idaapi.nextthat(addr, self.max_ea,
                                               idaapi.isCode)
                else:
                    codeend = self.max_ea - 1
                    addr = BADADDR
            else:
                if (idaapi.getseg(start).endEA < end):
                    codeend = idaapi.getseg(start).endEA - 1
                    addr = idaapi.getseg(idaapi.nextaddr(codeend)).startEA
                    if idaapi.isCode(idaapi.getFlags(addr)) == False:
                        addr = idaapi.nextthat(addr, self.max_ea,
                                               idaapi.isCode)
                else:
                    codeend = idaapi.get_item_end(idaapi.prevthat(end,
                                                start, idaapi.isCode)) - 1
                    addr = idaapi.nextthat(end, self.max_ea,
                                           idaapi.isCode)
                if (data < addr):
                    data = idaapi.nextthat(addr, self.max_ea,
                                           idaapi.isData)
                if (unknown < addr):
                    unknown = idaapi.next_unknown(addr, self.max_ea)
            self.start_element(CODE_BLOCK)
            self.write_address_attribute(START, start)
            self.write_address_attribute(END, codeend)
            self.close_tag()
        self.end_element(CODE)
        self.display_cpu_time(timer)


    def export_comment(self, addr, cmt_type, cmt):
        """
        Exports a <COMMENT> element with ADDRESS and TYPE attributes.
        The comment is exported as the element text (parsed character data).
        
        Args:
            addr: Integers representing address of comment.
            cmt_type: String indicating the comment type.
            cmt: String containing the comment.
        """
        self.start_element(COMMENT)
        self.write_address_attribute(ADDRESS, addr)
        self.write_attribute(TYPE, cmt_type)
        self.close_tag(True)
        # tag_remove seems to be losing last character
        # work around is to add a space
        cmt_text = idaapi.tag_remove(cmt + ' ')
        self.write_text(cmt_text)
        self.end_element(COMMENT, False)


    def export_comments(self):
        """
        Exports all comments in the IDA database as <COMMENT> elements.
        """
        addr = self.min_ea
        if idaapi.has_cmt(idaapi.getFlags(addr)) == False:
            addr = idaapi.nextthat(addr, self.max_ea, idaapi.has_cmt)
        if (addr == BADADDR):
            return
        self.update_status(COMMENTS)
        timer = time.clock()
        self.start_element(COMMENTS, True)
        while (addr != BADADDR):
            cmt = idaapi.get_cmt(addr, False)
            if (cmt != None):
                self.export_comment(addr, "end-of-line", cmt)
            cmt = idaapi.get_cmt(addr, True)
            if (cmt != None):
                self.export_comment(addr, "repeatable", cmt)
            addr = idaapi.nextthat(addr, self.max_ea, idaapi.has_cmt)
        addr = self.min_ea
        if idaapi.hasExtra(idaapi.getFlags(addr)) == False:
            addr = idaapi.nextthat(addr, self.max_ea, idaapi.hasExtra)
        while (addr != BADADDR):
            if IDA_SDK_VERSION < 640:
                extra = idaapi.ExtraGet(addr, idaapi.E_PREV)
            else:
                extra = idaapi.get_extra_cmt(addr, idaapi.E_PREV)
            if (extra != None):
                self.export_extra_comment(addr, "pre", idaapi.E_PREV)
            if IDA_SDK_VERSION < 640:
                extra = idaapi.ExtraGet(addr, idaapi.E_NEXT)
            else:
                extra = idaapi.get_extra_cmt(addr, idaapi.E_NEXT)
            if (extra != None):
                self.export_extra_comment(addr, "post", idaapi.E_NEXT)
            addr = idaapi.nextthat(addr, self.max_ea, idaapi.hasExtra)
        self.export_c_comments()
        self.end_element(COMMENTS)
        self.display_cpu_time(timer)


    def export_data(self):
        """
        Exports the data items in the database as <DEFINED_DATA> elements.
        """
        addr = self.min_ea
        if idaapi.isData(idaapi.getFlags(addr)) == False:
            addr = idaapi.nextthat(addr, self.max_ea, idaapi.isData)
        if (addr == BADADDR):
            return
        timer = time.clock()
        self.update_status(DATA)
        self.start_element(DATA, True)
        while (addr != BADADDR):
            f = idaapi.getFlags(addr)
            if idaapi.isAlign(f) == True:
                addr = idaapi.nextthat(addr, self.max_ea, idaapi.isData)
                continue
            dtype = self.get_datatype(addr)
            size = idaapi.get_item_size(addr)
            ti = idaapi.opinfo_t()
            if IDA_SDK_VERSION < 640:
                msize = idaapi.get_data_elsize(addr, f, ti)
            else:
                msize = idaapi.get_data_type_size(f, ti)
            if idaapi.isStruct(f) == True:
                s = idaapi.get_struc_id(dtype)
                msize = idaapi.get_struc_size(s)
                if msize == 0:
                    msize = 1
            if idaapi.isASCII(f) == False and size != msize:
                dtype = "%s[%d]" % (dtype, size/msize)
            self.start_element(DEFINED_DATA)
            self.write_address_attribute(ADDRESS, addr)
            self.write_attribute(DATATYPE, dtype)
            self.write_numeric_attribute(SIZE, size*self.cbsize)
            #TODO consider using GetTrueNameEx and Demangle
            demangled = idaapi.get_demangled_name(BADADDR, addr,
                            DEMANGLED_TYPEINFO, self.inf.demnames, True)
            outbuf = ''
            if IDA_SDK_VERSION < 660:
                outbuf = idaapi.print_type(addr, False)
            else:
                outbuf = idaapi.print_type(addr, outbuf)
            if demangled == "'string'":
                demangled == None
            has_typeinfo = demangled != None or (outbuf != None and
                                                len(outbuf) > 0)
            #TODO export_data: add DISPLAY_SETTINGS
            self.close_tag(has_typeinfo)
            if has_typeinfo == True:
                if demangled != None:
                    self.export_typeinfo_cmt(demangled)
                else:
                    #TODO export_data - check this
                    self.export_typeinfo_cmt(outbuf[:-1])
                self.end_element(DEFINED_DATA)
            addr = idaapi.nextthat(addr, self.max_ea, idaapi.isData)
        self.end_element(DATA)
        self.display_cpu_time(timer)
        

    def export_datatypes(self):
        """
        Exports the structures and enums in IDA database.
        """
        # skip if no structures/unions to export
        if idaapi.get_struc_qty() == 0: return
        self.update_status(DATATYPES)
        timer = time.clock()
        self.start_element(DATATYPES, True)
        self.export_structures()
        self.export_enums()
        self.end_element(DATATYPES)
        self.display_cpu_time(timer)


    def export_enum_member(self, cid, bf, mask, radix, signness):
        """
        Exports a member of an enum.

        This function can only be called by IDA versions newer than 6.3 
        
        Args:
            cid: Integer representing id of enum member
            bf: Boolean indicates if a bitfield
            mask: Integer representing bitmask if bitfield
            radix: Integer representing numeric display format
            signness: Boolean indicating if signed value 
        """
        if IDA_SDK_VERSION < 640:
            return
        cname = idaapi.get_const_name(cid)
        if cname == None or len(cname) == 0:
            return
        regcmt = idaapi.get_const_cmt(cid, False)
        rptcmt = idaapi.get_const_cmt(cid, True)
        has_comment =  regcmt != None or rptcmt != None
        self.start_element(ENUM_ENTRY)
        self.write_attribute(NAME, cname)
        value = idaapi.get_const_value(cid)
        self.write_numeric_attribute(VALUE, value, radix, signness)
        # BIT_MASK attribute not currently supported for ENUM_ENTRY
        #if bf == True:
        #    self.write_numeric_attribute(BIT_MASK, mask)
        self.close_tag(has_comment)
        if regcmt != None and len(regcmt) > 0:
            self.export_regular_cmt(regcmt)
        if rptcmt != None and len(rptcmt) > 0:
            self.export_repeatable_cmt(rptcmt)
        if (has_comment):
            self.end_element(ENUM_ENTRY)


    def export_enum_member_legacy(self, cid, bf, mask, radix, signness):
        """
        Exports a member of an enum.
        
        This function can only be called by IDA versions older than 6.3 

        Args:
            cid: Integer representing id of enum member
            bf: Boolean indicates if a bitfield
            mask: Integer representing bitmask if bitfield
            radix: Integer representing numeric display format
            signness: Boolean indicating if signed value 
        """
        if IDA_SDK_VERSION > 630:
            return
        cname = idaapi.get_enum_member_name(cid)
        if cname == None or len(cname) == 0:
            return
        regcmt = idaapi.get_enum_member_cmt(cid, False)
        rptcmt = idaapi.get_enum_member_cmt(cid, True)
        has_comment =  regcmt != None or rptcmt != None
        self.start_element(ENUM_ENTRY)
        self.write_attribute(NAME, cname)
        value = idaapi.get_enum_member_value(cid)
        self.write_numeric_attribute(VALUE, value, radix, signness)
        # BIT_MASK attribute not currently supported for ENUM_ENTRY
        #if bf == True:
        #    self.write_numeric_attribute(BIT_MASK, mask)
        self.close_tag(has_comment)
        if regcmt != None and len(regcmt) > 0:
            self.export_regular_cmt(regcmt)
        if rptcmt != None and len(rptcmt) > 0:
            self.export_repeatable_cmt(rptcmt)
        if (has_comment):
            self.end_element(ENUM_ENTRY)


    def export_enum_members(self, eid, bf, eflags):
        """
        Exports the members of an enum.

        This function can only be called by IDA versions newer than 6.3 
        
        Args:
            eid: Integer representing id of enum
            bf: Boolean indicates if a bitfield
            eflags: Integer representing the enum flags
        """
        if IDA_SDK_VERSION < 640:
            return
        mask=0xFFFFFFFF
        if bf == True:
            mask = idaapi.get_first_bmask(eid)
        first = True
        for n in range(idaapi.get_enum_size(eid)):
            if (first == True):
                value = idaapi.get_first_const(eid, mask)
                first = False
            else:
                value = idaapi.get_next_const(eid, value, mask)
            (cid, serial) = idaapi.get_first_serial_const(eid, value, mask)
            main_cid = cid
            while cid != idaapi.BADNODE:
                self.export_enum_member(cid, bf, mask,
                                   idaapi.getRadix(eflags, 0),
                                   self.is_signed_data(eflags))
                last_value = idaapi.get_last_const(eid, mask)
                if value == last_value:
                    # BIT_MASK not currently supported for ENUM
                    #self.export_bitmask(eid, mask)
                    mask = idaapi.get_next_bmask(eid, mask)
                    first = True
                (cid, serial) = idaapi.get_next_serial_const(main_cid)


    def export_enum_members_legacy(self, eid, bf, eflags):
        """
        Exports the members of an enum.
        
        This function can only be called by IDA versions older than 6.3 
        
        Args:
            eid: Integer representing id of enum
            bf: Boolean indicates if a bitfield
            eflags: Integer representing the enum flags
        """
        if IDA_SDK_VERSION > 630:
            return
        mask=0xFFFFFFFF
        if bf == True:
            mask = idaapi.get_first_bmask(eid)
        first = True
        for n in range(idaapi.get_enum_size(eid)):
            if (first == True):
                value = idaapi.get_first_enum_member(eid, mask)
                first = False
            else:
                value = idaapi.get_next_enum_member(eid, value, mask)
            (cid, serial) = idaapi.get_first_serial_enum_member(eid, value, mask)
            main_cid = cid
            while cid != idaapi.BADNODE:
                self.export_enum_member_legacy(cid, bf, mask,
                                   idaapi.getRadix(eflags, 0),
                                   self.is_signed_data(eflags))
                last_value = idaapi.get_last_enum_member(eid, mask)
                if value == last_value:
                    # BIT_MASK not currently supported for ENUM
                    #self.export_bitmask(eid, mask)
                    mask = idaapi.get_next_bmask(eid, mask)
                    first = True
                (cid, serial) = idaapi.get_next_serial_enum_member(main_cid)


    def export_enum_reference(self, addr, op):
        """
        Exports the enum reference for an operand at an address.
        
        Args:
            addr: Integer representing the instruction address.
            op: Integer representing the operand index (0-based)
        """
        (eid, serial) = idaapi.get_enum_id(addr, op)
        idaapi.decode_insn(addr)
        value = idaapi.cmd.Operands[op].value
        cid = idaapi.BADNODE
        last = idaapi.get_last_bmask(eid)
        if idaapi.is_bf(eid) == True:
            last = idaapi.get_last_bmask(eid)
            mask = idaapi.get_first_bmask(eid)
            while  (cid == idaapi.BADNODE):
                if IDA_SDK_VERSION < 640:
                    cid = idaapi.get_enum_member(eid, (value & mask), 0, mask)
                else:
                    cid = idaapi.get_const(eid, (value & mask), 0, mask)
                if cid != idaapi.BADNODE or mask == last:
                    break
                mask = idaapi.get_next_bmask(eid, mask)
        else:
            if IDA_SDK_VERSION < 640:
                cid = idaapi.get_enum_member(eid, value, 0, last)
            else:
               cid = idaapi.get_const(eid, value, 0, last)
        if (cid == idaapi.BADNODE):
            return
        self.start_element(EQUATE_REFERENCE)
        self.write_address_attribute(ADDRESS, addr)
        self.write_numeric_attribute(OPERAND_INDEX, op, 10)
        if IDA_SDK_VERSION < 640:
            self.write_numeric_attribute(VALUE, idaapi.get_enum_member_value(cid))
            cname = idaapi.get_enum_member_name(cid)
        else:
            self.write_numeric_attribute(VALUE, idaapi.get_const_value(cid))
            cname = idaapi.get_const_name(cid)
        if cname != None and len(cname) > 0:
            self.write_attribute(NAME, cname)
        # BIT_MASK feature not currently supported
        #if idaapi.is_bf(eid) == True:
        #    self.write_numeric_attribute("BIT_MASK", mask);
        self.close_tag()
        

    def export_enum_references(self, addr):
        """
        Finds and exports enum references at an address.
        
        Args:
            addr: Integer representing the instruction address.
        """
        f = idaapi.getFlags(addr)
        for op in range(2):
            if idaapi.isEnum(f, op) == True:
                self.export_enum_reference(addr, op)
                

    def export_enums(self):
        """
        Exports enumerations.
        """
        num_enums = idaapi.get_enum_qty()
        if (num_enums == 0):
            return
        for i in range(num_enums):
            self.start_element(ENUM)
            eid = idaapi.getn_enum(i)
            ename = idaapi.get_enum_name(eid)
            if (ename == None or len(ename) == 0):
                continue
            self.write_attribute(NAME, ename)
            ewidth = idaapi.get_enum_width(eid)
            if ewidth != 0 and ewidth <= 7:
                self.write_numeric_attribute(SIZE, 1 << (ewidth-1), 10)
            eflags = idaapi.get_enum_flag(eid)
            bf = idaapi.is_bf(eid)
            # BIT_FIELD attribute not supported for ENUM
            #if bf == True:
            #    self.write_attribute(BIT_FIELD, "yes")
            regcmt = idaapi.get_enum_cmt(eid, False)
            rptcmt = idaapi.get_enum_cmt(eid, True)
            has_children = ((idaapi.get_enum_size(eid) > 0) or
                            (regcmt != None) or (rptcmt != None) or
                            (idaapi.getRadix(eflags, 0) != 16) or
                            (self.is_signed_data(eflags) == True))
            self.close_tag(has_children)
            if (idaapi.getRadix(eflags, 0) != 16 or
                self.is_signed_data(eflags) == True):
                self.start_element(DISPLAY_SETTINGS)
                if idaapi.getRadix(eflags, 0) != 16:
                    self.write_attribute(FORMAT, self.get_format(eflags))
                if self.is_signed_data(eflags) == True:
                    self.write_attribute(SIGNED, "yes")
                self.close_tag()
            if regcmt != None:
                self.export_regular_cmt(regcmt)
            if rptcmt != None:
                self.export_repeatable_cmt(rptcmt)
            if IDA_SDK_VERSION < 640:
                self.export_enum_members_legacy(eid, bf, eflags)
            else:
                self.export_enum_members(eid, bf, eflags)
            if (has_children):
                self.end_element(ENUM)


    def export_extra_comment(self, addr, cmt_type, extra):
        """
        Exports pre- and post- comments for an address.
        
        Args:
            addr: Integer representing the instruction address.
            cmt_type: String indicating comment type
            extra: Integer representing extra comment index
        """
        cmt = ''
        if IDA_SDK_VERSION < 640:
            nextline = idaapi.ExtraGet(addr, extra)
        else:
            nextline = idaapi.get_extra_cmt(addr, extra)
        while (nextline != None):
            # workaround for tag_remove bug is to add space
            cmt += idaapi.tag_remove(nextline + ' ')
            extra += 1
            if IDA_SDK_VERSION < 640:
                nextline = idaapi.ExtraGet(addr, extra)
            else:
                nextline = idaapi.get_extra_cmt(addr, extra)
            if (nextline != None):
                cmt += '\n' 
        self.export_comment(addr, cmt_type, cmt)


    def export_functions(self):
        """
        Exports information about all functions. 
        """
        functions = idautils.Functions()
        if functions == None:
            return
        self.update_status(FUNCTIONS)
        timer = time.clock()
        self.start_element(FUNCTIONS, True)
        for addr in functions:
            function = idaapi.get_func(addr)
            if idaapi.is_spec_ea(function.startEA) == True:
                continue
            self.start_element(FUNCTION)
            self.write_address_attribute(ENTRY_POINT, function.startEA)
            if idaapi.has_user_name(idaapi.getFlags(addr)) == True:
                name = self.get_symbol_name(addr)
                if name != None and len(name) > 0:
                    self.write_attribute(NAME, name)
            if function.flags & idaapi.FUNC_LIB != 0:
                self.write_attribute(LIBRARY_FUNCTION, "y")
            self.close_tag(True)
            fchunks = idautils.Chunks(addr)
            for (startEA, endEA) in fchunks:
                self.start_element(ADDRESS_RANGE)
                self.write_address_attribute(START, startEA)
                self.write_address_attribute(END, endEA-1)
                self.close_tag()
            regcmt = idaapi.get_func_cmt(function, False)
            if regcmt != None:
                self.export_regular_cmt(regcmt)
            rptcmt = idaapi.get_func_cmt(function, True)
            if rptcmt != None:
                self.export_repeatable_cmt(rptcmt)
            demangled = idaapi.get_demangled_name(BADADDR, addr,
                                            DEMANGLED_TYPEINFO,
                                            self.inf.demnames, True)
            if demangled != None and demangled == "'string'":
                demangled = None
            outbuf = ''
            if IDA_SDK_VERSION < 660:
                outbuf = idaapi.print_type(addr, False)
            else:
                outbuf = idaapi.print_type(addr, outbuf)
            has_typeinfo = (demangled != None or (outbuf != None and
                            len(outbuf) > 0))
            if demangled != None:
                self.export_typeinfo_cmt(demangled)
            elif has_typeinfo == True:
                self.export_typeinfo_cmt(outbuf[:-1])
            self.export_stack_frame(function)
            self.end_element(FUNCTION)
        self.end_element(FUNCTIONS)
        self.display_cpu_time(timer)


    def export_manual_instruction(self, addr):
        """
        Exports user-entered "manual instruction" at an address.
        
        Args:
            addr: Integer representing instruction address.
        """
        text = idaapi.get_manual_insn(addr)
        if text == None or len(text) == 0:
            return
        self.start_element(MANUAL_INSTRUCTION)
        self.write_address_attribute(ADDRESS, addr)
        self.close_tag(True)
        self.write_text(text)
        self.end_element(MANUAL_INSTRUCTION, False)
        

    def export_manual_operand(self, addr):
        """
        Exports user-entered "manual operands" at an address.
        
        Args:
            addr: Integer representing instruction address.
        """
        for op in range(idaapi.UA_MAXOP):
            if idaapi.is_forced_operand(addr, op) == True:
                text = idaapi.get_forced_operand(addr, op)
                if text != None and len(text) > 0:
                    self.start_element(MANUAL_OPERAND)
                    self.write_address_attribute(ADDRESS, addr)
                    self.write_numeric_attribute(OPERAND_INDEX, op, 10)
                    self.close_tag(True)
                    self.write_text(text)
                    self.end_element(MANUAL_OPERAND, False)


    def export_markup(self):
        """
        Exports markup for instructions and data items including references
        and manual instructions and operands.
        """
        self.update_status(MARKUP)
        timer = time.clock()
        self.start_element(MARKUP, True)
        addr = self.min_ea
        while addr != BADADDR:
            f = idaapi.getFlags(addr)
            if self.options.MemoryReferences.checked == True:
                if idaapi.hasRef(f) == True:
                    self.export_user_memory_reference(addr)
                if idaapi.isOff(f, idaapi.OPND_ALL) == True:
                    self.export_memory_references(addr)
            if (self.options.Functions.checked == True and
               self.options.StackReferences.checked == True and
               idaapi.isStkvar(f, idaapi.OPND_ALL) == True):
               self.export_stack_reference(addr)
            if (self.options.DataTypes.checked == True and
                    idaapi.isEnum(f, idaapi.OPND_ALL) == True):
                self.export_enum_references(addr)
            if self.options.Manual.checked == True:
                if idaapi.isFop(f, idaapi.OPND_ALL)   == True:
                    self.export_manual_operand(addr)
                if idaapi.is_manual_insn(addr) == True:
                    self.export_manual_instruction(addr)
            addr = idaapi.next_head(addr, self.max_ea)
        self.end_element(MARKUP)
        self.display_cpu_time(timer)


    def export_members(self, s):
        """
        Exports the members of a structure or union.
        
        Args:
            s: IDA structure/union instance
        """
        nmembers = s.memqty
        for n in range(nmembers):
            m = s.get_member(n)
            offset = m.soff
            if s.is_union() == True:
                offset = 0
            self.start_element(MEMBER)
            self.write_numeric_attribute(OFFSET, offset)
            mname = idaapi.get_member_name(m.id)
            if len(mname) > 0:
                self.write_attribute(NAME, mname)
            dtype = self.get_member_type(m)
            if idaapi.is_varmember(m) == True:
                msize = 0
                size  = 0
            else:
                mtibuf = idaapi.opinfo_t()
                mti = idaapi.retrieve_member_info(m, mtibuf)
                if IDA_SDK_VERSION < 640:
                    msize = idaapi.get_type_size0(None, dtype)
                    if msize == None or msize == 0:
                        msize = idaapi.get_member_size(m)
                else:
                    msize = idaapi.get_data_type_size(m.flag, mtibuf)
                size = idaapi.get_member_size(m)
                osize = size
                if size < msize: size = msize
            if (size != msize):
                arraytype = self.get_member_type(m)
                dtype = "%s[%d]" % (arraytype, size/msize)
            self.write_attribute(DATATYPE, dtype)
            self.write_numeric_attribute(SIZE, size*self.cbsize)
            regcmt = idaapi.get_member_cmt(m.id, False)
            rptcmt = idaapi.get_member_cmt(m.id, True)
            hascmt = regcmt != None or rptcmt != None
            self.close_tag(hascmt)
            if (hascmt):
                if regcmt != None:
                    self.export_regular_cmt(regcmt)
                if rptcmt != None:
                    self.export_repeatable_cmt(rptcmt)
                self.end_element(MEMBER)


    def export_memory_contents(self, binfilename, binfile, start, end):
        """
        Exports the binary memory contents in the database.
        
        A MEMORY_CONTENTS element is generated for each contiguous address
        range where each address in the range contains a value.
        The binary values are store in a separate file (not the XML file),
        and the MEMORY_CONTENTS element identifies the file and the
        offset in the file where the address range is located.

        Args:
            binfilename: String containing the absolute filepath
            binfile: IDA file instance for binary file
            start: Integer representing the starting address
            end: Integer representing the ending address
        """
        length = 0
        startaddr = start
        for addr in range(start, end):
            # reset start address when length == 0
            if (length == 0):
                startaddr = addr
            has_val = idaapi.hasValue(idaapi.getFlags(addr))
            if has_val == True:
                length += self.cbsize
            next_address = idaapi.nextaddr(addr)
            if ((has_val == False) or (next_address != addr+1) or
                    (next_address == end)):
                if length > 0:
                    offset = binfile.tell()
                    idaapi.base2file(binfile.get_fp(), offset, startaddr,
                                     startaddr+length)
                    self.start_element(MEMORY_CONTENTS)
                    self.write_address_attribute(START_ADDR, startaddr)
                    self.write_attribute(FILE_NAME, binfilename)
                    self.write_numeric_attribute(FILE_OFFSET, offset)
                    self.write_numeric_attribute(LENGTH, length)
                    self.close_tag(False)
                    length=0


    def export_memory_map(self):
        """
        Exports information about all memory blocks in the database.
        
        A MEMORY_SECTION is generated for each block (segment). If the
        memory block is initialized (has values), the contents are exported
        using the MEMORY_CONTENTS element.
        """
        nsegs = idaapi.get_segm_qty()
        if (nsegs == 0):
            return
        self.update_status(MEMORY_MAP)
        timer = time.clock();
        binfilename = ''
        if (self.options.MemoryContent.checked == True):
            (binfilename, ext) = os.path.splitext(self.filename)
            binfilename += ".bytes"
            self.binfile = idaapi.qfile_t()
            self.binfile.open(binfilename,'wb');
        self.start_element(MEMORY_MAP, True)
        for i in range(nsegs):
            self.export_memory_section(idaapi.getnseg(i), binfilename)
        self.end_element(MEMORY_MAP)
        if (self.options.MemoryContent.checked == True):
            self.close_binfile()
        self.display_cpu_time(timer)


    def export_memory_reference(self, addr, op):
        """
        Exports the memory reference for operand at the address.
        
        Args:
            addr: Integer representing the instruction address.
            op: Integer representing the operand index (0-based)
        """
        f = idaapi.getFlags(addr)
        ri = idaapi.refinfo_t()
        if idaapi.get_refinfo(addr, op, ri) == 1: 
            if ri.target != BADADDR:
                target = ri.target
            elif idaapi.isCode(f) == True:
                idaapi.decode_insn(addr)
                target = idaapi.cmd.Operands[op].value - ri.tdelta + ri.base
            elif idaapi.isData(f) == True:
                target = self.get_data_value(addr) - ri.tdelta + ri.base;
            else:
                return
        else:
            return
        if idaapi.isEnabled(target) == False:
            return
        self.start_element(MEMORY_REFERENCE)
        self.write_address_attribute(ADDRESS, addr)
        self.write_numeric_attribute(OPERAND_INDEX, op, 10)
        self.write_address_attribute(TO_ADDRESS, target)
        self.write_attribute(PRIMARY, "y")
        self.close_tag()
        

    def export_memory_references(self, addr):
        """
        Exports the memory references for any operands at the address.
        
        Args:
            addr: Integer representing the instruction address.
        """
        f = idaapi.getFlags(addr)
        for op in range(idaapi.UA_MAXOP):
            if idaapi.isOff(f, op) == True and (idaapi.isData(f) == True or
                    (idaapi.isCode(f) == True and
                    self.is_imm_op(addr, op) == True)):
                self.export_memory_reference(addr, op)
    

    def export_memory_section(self, seg, binfilename):
        """
        Exports segment information as a MEMORY_SECTIONS element.
        
        Args:
            seg: IDA segment instance
            binfilename: String containing absolute filepath for binary file.
        """
        segname = idaapi.get_segm_name(seg)
        self.start_element(MEMORY_SECTION)
        self.write_attribute(NAME, segname)
        self.write_address_attribute(START_ADDR, seg.startEA)
        length = (seg.endEA - seg.startEA)*self.cbsize
        self.write_numeric_attribute(LENGTH, length)
        perms = ""
        if (seg.perm != 0):
            if (seg.perm & idaapi.SEGPERM_READ  != 0):
                perms += 'r'
            if (seg.perm & idaapi.SEGPERM_WRITE != 0):
                perms += 'w' 
            if (seg.perm & idaapi.SEGPERM_EXEC  != 0):
                perms += 'x'
            if (len(perms) > 0):
                self.write_attribute(PERMISSIONS, perms)
        has_contents = (self.options.MemoryContent.checked == True and
                       self.check_if_seg_contents(seg) == True)
        self.close_tag(has_contents)
        if (has_contents == True):
            self.export_memory_contents(os.path.basename(binfilename),
                                      self.binfile, seg.startEA, seg.endEA)
            self.end_element(MEMORY_SECTION)


    def export_program(self):
        """
        Exports basic information about the program as the PROGRAM,
        INFO_SOURCE, PROCESSOR, and COMPILER elements.
        """
        # output the PROGRAM element
        self.update_status(PROGRAM);
        timer = time.clock()
        self.start_element(PROGRAM)
        self.write_attribute(NAME, idaapi.get_root_filename())
        self.write_attribute(EXE_PATH, idaapi.get_input_file_path())
        etype = idaapi.get_file_type_name()
        if (len(etype) > 0):
            self.write_attribute(EXE_FORMAT, etype)
        # check for presence of INPUT_MD5 netnode
        md5 = idaapi.netnode(INPUT_MD5)
        if md5 == idaapi.BADNODE:
            input_md5 = idc.GetInputMD5()
        else:
            input_md5 = md5.supval(idaapi.RIDX_MD5)
        if input_md5 != None:
            self.write_attribute(INPUT_MD5,input_md5)
        self.close_tag(True)
    
        # output the INFO_SOURCE element
        self.start_element(INFO_SOURCE)
        if IDA_SDK_VERSION < 670:
            tool  = 'IDA-Pro XML plugin (Python) SDK '
            tool +=  str(IDA_SDK_VERSION)
        else:
            tool  = 'IDA-Pro ' + idaapi.get_kernel_version()
            tool += ' XML plugin (Python) SDK ' + str(IDA_SDK_VERSION)
        self.write_attribute(TOOL, tool)
        user = os.getenv("USERNAME", "UNKNOWN")
        if (user == "UNKNOWN"):
            user = os.getenv("USER", "UNKNOWN")
        self.write_attribute(USER, user)
        self.write_attribute(FILE, idaapi.cvar.database_idb)
        ts = datetime.datetime.now().strftime("%Y-%m-%d %H:%M")
        self.write_attribute(TIMESTAMP, ts)
        self.close_tag()
    
        # output the PROCESSOR element
        self.start_element(PROCESSOR)
        self.write_attribute(NAME, self.inf.procName)
        if (self.inf.mf):
            byte_order ="big"
        else:
            byte_order ="little"
        self.write_attribute(ENDIAN, byte_order)
        self.seg_addr = False
        bitness = 1
        model_warning = False
        nsegs = idaapi.get_segm_qty()
        if (nsegs > 0):
            bitness = idaapi.getnseg(0).bitness
            for i in range(1,nsegs):
                seg = idaapi.getnseg(i)
                if (seg.bitness != bitness):
                    model_warning = True
                if (seg.bitness > bitness):
                    bitness = seg.bitness
        addr_model = "32-bit"
        if (bitness == 0):
            addr_model = "16-bit"
        elif (bitness == 2):
            addr_model = "64-bit"
        self.write_attribute(ADDRESS_MODEL, addr_model)
        self.close_tag()
        if (model_warning):
            idaapi.msg("WARNING: Segments do not have same " +
                       "addressing model!\n")
        if (idaapi.ph.id == idaapi.PLFM_386 and bitness == 0):
            self.seg_addr = True
        # find any overlayed memory before processing addressable items
        self.find_overlay_memory()
    
        # output compiler info
        self.start_element(COMPILER)
        self.write_attribute(NAME, idaapi.get_compiler_name(self.inf.cc.id))
        self.close_tag()
        self.display_cpu_time(timer)


    def export_program_entry_points(self):
        """
        Exports entry points for the program.
        """
        nepts = idaapi.get_entry_qty()
        if (nepts  == 0):
            return
        self.update_status(PROGRAM_ENTRY_POINTS)
        timer = time.clock()
        self.start_element(PROGRAM_ENTRY_POINTS, True)
        for i in range(nepts):
            self.start_element(PROGRAM_ENTRY_POINT)
            addr = idaapi.get_entry(idaapi.get_entry_ordinal(i))
            self.write_address_attribute(ADDRESS, addr)
            self.close_tag()
        self.end_element(PROGRAM_ENTRY_POINTS)
        self.display_cpu_time(timer)


    def export_register_values(self):
        """
        Exports segment register value ranges.
        """
        # should use legacy function if IDA > 660
        if IDA_SDK_VERSION < 670:
            self.export_register_values_legacy()
            return
        first = idaapi.ph_get_regFirstSreg()
        last  = idaapi.ph_get_regLastSreg()+1
        has_segregareas = False
        for j in range(first, last):
            nsegregareas = idaapi.get_srareas_qty2(j)
            if nsegregareas != 0:
                has_segregareas = True
                break;
        if has_segregareas == False:
            return
        self.update_status(REGISTER_VALUES)
        timer = time.clock();
        self.start_element(REGISTER_VALUES, True)
        sr = idaapi.segreg_area_t()
        for j in range(first, last):
            nsegregareas = idaapi.get_srareas_qty2(j)
            if nsegregareas == 0:
                continue
            for i in range(nsegregareas):
                success = idaapi.getn_srarea2(sr, j, i)
                if success == False:
                    continue
                value = sr.val
                if value == idaapi.BADSEL:
                    continue
                regname = idaapi.ph.regnames[j]
                if regname == None:
                    continue
                if regname.lower() == "cs":
                    continue
                if (idaapi.ph.id == idaapi.PLFM_TMS and
                    regname.lower() == "ds"):
                    continue
                self.start_element(REGISTER_VALUE_RANGE)
                self.write_attribute(REGISTER, idaapi.ph.regnames[j])
                self.write_numeric_attribute(VALUE, value)
                self.write_address_attribute(START_ADDRESS, sr.startEA)
                length = (sr.endEA - sr.startEA) * self.cbsize
                self.write_numeric_attribute(LENGTH, length)
                self.close_tag()
        self.end_element(REGISTER_VALUES)
        self.display_cpu_time(timer)


    def export_register_values_legacy(self):
        """
        Exports segment register value ranges.
        """
        # should use export_register_values() if IDA > 660
        if IDA_SDK_VERSION > 660:
            return
        if IDA_SDK_VERSION < 650:
            n = 0
            srarea = idaapi.getnSRarea(n)
            while type(srarea) == idaapi.segreg_t:
                n += 1
                srarea = idaapi.getnSRarea(n)
            nsegregareas = n
        else:
            nsegregareas = idaapi.get_srareas_qty()
        if (nsegregareas == 0):
            return
        self.update_status(REGISTER_VALUES)
        timer = time.clock();
        self.start_element(REGISTER_VALUES, True)
        for i in range(nsegregareas):
            if IDA_SDK_VERSION < 650:
                sr = idaapi.getnSRarea(i)
            else:
                sr = idaapi.getn_srarea(i)
            first = idaapi.ph_get_regFirstSreg()
            last  = idaapi.ph_get_regLastSreg()+1
            for j in range(first, last):
                value = idaapi.getSR(sr.startEA, j)
                if value == idaapi.BADSEL:
                    continue
                regname = idaapi.ph.regnames[j]
                if regname == None:
                    continue
                if regname.lower() == "cs":
                    continue
                if (idaapi.ph.id == idaapi.PLFM_TMS and
                    regname.lower() == "ds"):
                    continue
                self.start_element(REGISTER_VALUE_RANGE)
                self.write_attribute(REGISTER, idaapi.ph.regnames[j])
                self.write_numeric_attribute(VALUE, value)
                self.write_address_attribute(START_ADDRESS, sr.startEA)
                length = (sr.endEA - sr.startEA) * self.cbsize
                self.write_numeric_attribute(LENGTH, length)
                self.close_tag()
        self.end_element(REGISTER_VALUES)
        self.display_cpu_time(timer)


    def export_regular_cmt(self, cmt):
        """
        Exports the regular comment for an item.
        
        Args:
            cmt: String containing the regular comment.
        """
        self.write_comment_element(REGULAR_CMT, cmt)
    

    def export_repeatable_cmt(self, cmt):
        """
        Exports the repeatable comment for an item.
        
        Args:
            cmt: String containing the repeatable comment.
        """
        self.write_comment_element(REPEATABLE_CMT, cmt)
    

    def export_stack_frame(self, function):
        """
        Export information about a function stack frame including
        variables allocated on the stack.
        
        Args:
            function: IDA function instance
        """
        sframe = idaapi.get_struc(function.frame)
        if sframe == None or sframe.memqty <= 0:
            return
        self.start_element(STACK_FRAME)
        self.write_numeric_attribute(LOCAL_VAR_SIZE, function.frsize)
        self.write_numeric_attribute(REGISTER_SAVE_SIZE, function.frregs)
        retsize = idaapi.get_frame_retsize(function)
        self.write_numeric_attribute(RETURN_ADDR_SIZE, retsize)
        self.write_numeric_attribute(BYTES_PURGED, function.argsize)
        has_stack_vars = self.check_stack_frame(sframe)
        self.close_tag(has_stack_vars)
        if has_stack_vars == True:
            self.export_stack_vars(function, sframe)
            self.end_element(STACK_FRAME)

    def export_stack_reference(self, addr):
        """
        Exports references to stack variables at the address.
        
        Args:
            addr: Integer containing instruction address.
        """
        f = idaapi.getFlags(addr)
        for op in range(idaapi.UA_MAXOP):
            if idaapi.isCode(f) == True and idaapi.isStkvar(f, op) == True:
                idaapi.decode_insn(addr)
                opnd = idaapi.cmd.Operands[op]
                optype = idaapi.op_t_get_type(opnd)
                if optype == idaapi.o_void:
                    continue
                SV = idaapi.get_stkvar(opnd, idaapi.op_t_get_addr(opnd))
                if SV == None:
                    continue
                (sv, actval) = SV
                function = idaapi.get_func(addr)
                self.start_element(STACK_REFERENCE)
                self.write_address_attribute(ADDRESS, addr)
                self.write_numeric_attribute(OPERAND_INDEX, op, 10)
                offset = opnd.addr
                spoff = offset - function.frregs
                if offset > 0x7FFFFFFF:
                    offset -= 0x100000000
                if spoff > 0x7FFFFFFF:
                    spoff  -= 0x100000000
                self.write_numeric_attribute(STACK_PTR_OFFSET,
                                             spoff,
                                             16, True)
                if (function.flags & FUNC_FRAME) != 0:
                    self.write_numeric_attribute(FRAME_PTR_OFFSET,
                                                offset,
                                                16, True)
                self.close_tag()

    def export_stack_vars(self, function, sframe):
        """
        Exports the stack variables (parameters and locals) in a stack frame.
        
        Args:
            function: IDA function instance.
            sframe: IDA stack frame instance.
        """
        for i in range(sframe.memqty):
            member = sframe.get_member(i)
            if member == None:
                continue
            mname = idaapi.get_member_name(member.id)
            if mname == None or len(mname) < 0:
                continue
            if mname == " s" or mname == " r":
                continue
            spoff = member.soff - function.frsize - function.frregs
            froff = member.soff - function.frsize
            self.start_element(STACK_VAR)
            self.write_numeric_attribute(STACK_PTR_OFFSET, spoff, 16, True)
            if function.flags & idaapi.FUNC_FRAME != 0:
                self.write_numeric_attribute(FRAME_PTR_OFFSET, froff, 16, True)
            pre = mname[0:4]
            if pre != "var_" and pre != "arg_":
                self.write_attribute(NAME, mname)
            f = member.flag
            size = idaapi.get_member_size(member)
            mtype = self.get_member_type(member)
            msize = size
            if idaapi.isStruct(f) == True:
                msize = idaapi.get_struc_size(idaapi.get_struc_id(mtype))
            elif idaapi.isASCII(f) == False:
                mtibuf = idaapi.opinfo_t()
                mti = idaapi.retrieve_member_info(member, mtibuf)
                if IDA_SDK_VERSION < 640:
                    msize = idaapi.get_type_size0(None, mtype)
                    if msize == None or msize == 0:
                        msize = idaapi.get_member_size(member)
                else:
                    msize = idaapi.get_data_type_size(f, mtibuf)
            if size < msize: size = msize
            if (idaapi.isASCII(f) == False and idaapi.isAlign(f) == False
                and size != msize):
                mtype = "%s[%d]" % (mtype, size/msize)
            self.write_attribute(DATATYPE, mtype)
            self.write_numeric_attribute(SIZE, size*self.cbsize)
            regcmt = idaapi.get_member_cmt(member.id, False)
            rptcmt = idaapi.get_member_cmt(member.id, True)
            if regcmt != None:
                regcmt  = idaapi.tag_remove(regcmt + " ", 0)
            if rptcmt != None:
                rptrcmt = idaapi.tag_remove(rptcmt + " ", 0)
            has_regcmt = regcmt != None and len(regcmt) > 0
            has_rptcmt = rptcmt != None and len(rptcmt) > 0
            has_content = has_regcmt or has_rptcmt
            self.close_tag(has_content)
            if has_content == True:
                if has_regcmt == True:
                    self.export_regular_cmt(regcmt)
                if has_rptcmt == True:
                    self.export_repeatable_cmt(rptcmt)
                self.end_element(STACK_VAR)


    def export_structures(self):
        """
        Exports information about all structures and unions.
        """
        structs = idautils.Structs()
        for struct in structs:
            (sid, idx, sname) = struct
            s = idaapi.get_struc(idx)
            stype = STRUCTURE
            if s.is_union() == True:
                stype = UNION
            self.start_element(stype)
            self.write_attribute(NAME, sname)
            size = idaapi.get_struc_size(idx)*self.cbsize
            self.write_numeric_attribute(SIZE, size)
            if s.is_varstr() == True:
                self.write_attribute(VARIABLE_LENGTH, "y")
            regcmt = idaapi.get_struc_cmt(idx, False)
            rptcmt = idaapi.get_struc_cmt(idx, True)
            has_contents = regcmt != None or rptcmt != None or s.memqty > 0
            self.close_tag(has_contents)
            if (has_contents):
                if regcmt != None:
                    self.export_regular_cmt(regcmt)
                if rptcmt != None:
                    self.export_repeatable_cmt(rptcmt)
                if s.memqty > 0:
                    self.export_members(s)
                self.end_element(stype)
    

    def export_symbol(self, addr, name, stype=""):
        """
        Exports name for an address as a SYMBOL element. If the name is a
        demangled name, add the mangled name as the MANGLED attribute.
        
        Args:
            addr: Integer representing the symbol address.
            name: String containing the symbol name.
            stype: String indicating symbol type (global or local)
        """
        self.start_element(SYMBOL)
        self.write_address_attribute(ADDRESS, addr)
        self.write_attribute(NAME, name)
        self.write_attribute(TYPE, stype)
        mangled = idaapi.get_name(BADADDR, addr)
        if name != None and mangled != name:
            self.write_attribute("MANGLED", mangled)
        self.close_tag()
        

    def export_symbol_table(self):
        """
        Exports user-defined and non-default names as SYMBOL elements.
        """
        addr = self.min_ea
        if idaapi.has_any_name(idaapi.getFlags(addr)) == False:
            addr = idaapi.nextthat(addr, self.max_ea, idaapi.has_any_name)
        if addr == BADADDR:
            return
        self.update_status(SYMBOL_TABLE)
        self.start_element(SYMBOL_TABLE, True)
        timer = time.clock()
        while addr != BADADDR:
            # only export meaningful names (user and auto)
            f = idaapi.getFlags(addr)
            if (idaapi.has_user_name(f) == True or
                idaapi.has_auto_name(f) == True):
                # check for global name
                name = self.get_symbol_name(addr)
                if name != None and len(name) > 0:
                    self.export_symbol(addr, name)
                # check for local name
                if ((IDA_SDK_VERSION > 620 and
                    idaapi.has_lname(addr)) or
                    idaapi.get_aflags(addr) & idaapi.AFL_LNAME):
                    if IDA_SDK_VERSION < 680:
                        name = idaapi.get_name(addr, addr)
                    else:
                        name = idaapi.get_ea_name(addr, idaapi.GN_LOCAL)
                    if name != None and len(name) > 0:
                        self.export_symbol(addr, name, 'local')
            # get next address with any name
            addr = idaapi.nextthat(addr, self.max_ea,
                                   idaapi.has_any_name)
        self.end_element(SYMBOL_TABLE)
        self.display_cpu_time(timer)
        

    def export_typeinfo_cmt(self, cmt):
        """
        Exports comment containing type information for data and functions.
        
        Args:
            cmt: String containing type info.
        """
        # older versions of IDAPython returned a '\n' at end of cmt
        if cmt[-1] == '\n':
            cmt = cmt[0:-1]
        self.write_comment_element(TYPEINFO_CMT, cmt)
        

    def export_user_memory_reference(self, addr):
        """
        Exports a user-specified memory reference at the address.
        
        Args:
            addr: Integer representing the instruction address.
        """
        for xref in idautils.XrefsTo(addr, idaapi.XREF_FAR):
            if xref.user == 1:
                self.start_element(MEMORY_REFERENCE)
                self.write_address_attribute(ADDRESS, xref.frm)
                self.write_address_attribute(TO_ADDRESS, xref.to)
                self.write_attribute(USER_DEFINED, "y")
                self.close_tag()


    def find_overlay_memory(self):
        """
        Determines if any memory blocks (segments) are overlays.
        
        A segment is an overlay if it translates to the same logical
        address as another segment. This is rare, but may occur, for
        example when a processor has a small logical address space
        (i.e. a 16-bit address is limited to 64K) and multiple physical
        segments are mapped into the same logical segment.
        """
        self.overlay = dict()
        self.has_overlays = False;
        nsegs = idaapi.get_segm_qty()
        if nsegs == 0:
            return
        s = idaapi.getnseg(0)
        start = self.translate_address(s.startEA)
        self.overlay[start] = False
        for i in range(1, nsegs):
            s = idaapi.getnseg(i)
            space = self.get_space_name(s.startEA)
            saddr = self.translate_address(s.startEA)
            eaddr = self.translate_address(s.endEA-1)
            is_overlay = False
            for j in range(i):
                s2 = idaapi.getnseg(j)
                space2 = self.get_space_name(s2.startEA)
                if space == space2:
                    start = self.translate_address(s2.startEA)
                    end   = self.translate_address(s2.endEA-1)
                    if ((saddr >= start and saddr <= end) or 
                        (eaddr >= start and eaddr <= end)):
                        is_overlay = True
                        self.has_overlays = True
                        break
            self.overlay[saddr] = is_overlay


    def get_address_string(self, addr):
        """
        Returns a string representing the address.
        
        The representation is typically a hex string of the address,
        but may include a segment or space name prefixe based on the
        processor or architecture.
        
        Args:
            addr: Integer representing a program address.
        """
        temp = "0x%X" % (addr - idaapi.get_segm_base(idaapi.getseg(addr)))
        space = self.get_space_name(addr)
        if space != None:
            temp = "%s:%04X" % (space,
                            addr - idaapi.get_segm_base(idaapi.getseg(addr)))
        else:
            if (idaapi.ph_get_id() == idaapi.PLFM_386 and
                idaapi.getseg(addr).bitness == 0):
                base = idaapi.get_segm_para(idaapi.getseg(addr))
                temp = "%04X:%04X" % (base, addr - (base << 4))
        if idaapi.ph_get_id() == idaapi.PLFM_C166:
            temp = "0x%X" % addr
        if self.has_overlays == True and self.is_overlay(addr) == True:
            oname = idaapi.get_segm_name(idaapi.getseg(addr))
            if len(oname) > 0:
                temp = oname + "::" + temp
        return temp


    def get_data_value(self, addr):
        """
        Returns the data item value at an address based on its size.
        
        Args:
            addr: Integer representing a program address.
        """
        size = idaapi.get_item_size(addr)*self.cbsize
        if size == 1:   return idaapi.get_byte(addr)
        if size == 2:   return idaapi.get_16bit(addr)
        if size == 4:   return idaapi.get_32bit(addr)
        if size == 8:   return idaapi.get_64bit(addr)
        return 0
    

    def get_datatype(self, addr):
        """
        Returns the datatype at an address.
        
        The type could be a basic type (byte, word, dword, etc.),
        a structure, an array, a pointer, or a string type.
        
        Args:
            addr: Integer representing a program address.
        """
        f = idaapi.getFlags(addr)
        t = self.get_type(f)
        if idaapi.isStruct(f) == True:
            opndbuf = idaapi.opinfo_t()
            opnd = idaapi.get_opinfo(addr, 0, f, opndbuf)
            return idaapi.get_struc_name(opnd.tid)
        if idaapi.isASCII(f) == True:
            str_type = idc.GetStringType(addr)
            if str_type == idaapi.ASCSTR_TERMCHR:   return "string"
            if str_type == idaapi.ASCSTR_PASCAL:    return "string1"
            if str_type == idaapi.ASCSTR_LEN2:      return "string2"
            if str_type == idaapi.ASCSTR_LEN4:      return "string4"
            if str_type == idaapi.ASCSTR_UNICODE:   return "unicode"
            if str_type == idaapi.ASCSTR_ULEN2:     return "unicode2"
            if str_type == idaapi.ASCSTR_ULEN4:     return "unicode4"
            return "string"
        if idaapi.isOff0(f) == True: return "pointer"
        return t


    def get_format(self, flags):
        """
        Returns the display format of a data item based on its flags.
        
        Args:
            flags: Integer representing IDA item flags
            
        Returns:
            String representing IDA display format.
        """
        if idaapi.isChar0(flags): return "char"
        radix = idaapi.getRadix(flags, 0)
        if radix == 2:  return "binary"
        if radix == 8:  return "octal"
        if radix == 10: return "decimal"
        return "hex" # default
    

    def get_member_type(self, m):
        """
        Returns the datatype of a structure member.
        
        Args:
            m: IDA member instance.
            
        Returns:
            String representing member datatype.
        """
        f = m.flag
        t = self.get_type(f)
        if idaapi.isOff0(f) == True:
            t = "pointer"
        if idaapi.isStruct(f) == False:
            return t
        s = idaapi.get_sptr(m)
        if (s == None):
            return t
        sname = idaapi.get_struc_name(s.id)
        if (sname == None):
            return t
        return sname


    def get_options(self):
        """
        Displays the options menu and retrieves the option settings. 
        """
        fmt =  "HELP\n"
        fmt += "XML plugin (Python)\n"
        fmt += "IDA SDK: "+ str(IDA_SDK_VERSION) + "\n"
        fmt += "\n"
        fmt += "The XML interface provides a dump of the IDA-Pro database as "
        fmt += "a XML \"PROGRAM\" document. The XML PROGRAM document contains "
        fmt += "information from the idb file in a readable text format, and "
        fmt += "can be viewed with a text editor or web browser.\n\n"
        fmt += "ENDHELP\n"
        fmt += "Export as XML PROGRAM document...."
        fmt += "\n <##Options##Memory Sections:{MemorySections}>"
        fmt += "\n <Memory Content:{MemoryContent}>"
        fmt += "\n <Segment Register Value Ranges:{RegisterValues}>"
        fmt += "\n <Data Types:{DataTypes}>"
        fmt += "\n <Code Blocks:{CodeBlocks}>"
        fmt += "\n <Data Definitions:{DataDefinitions}>"
        fmt += "\n <Comments:{Comments}>"
        fmt += "\n <Entry Points:{EntryPoints}>"
        fmt += "\n <Symbols:{Symbols}>"
        fmt += "\n <Functions:{Functions}>"
        fmt += "\n <Memory References:{MemoryReferences}>"
        fmt += "\n <Stack References:{StackReferences}>"
        fmt += "\n <Manual Instructions/Operands:{Manual}>{cGroup1}>"
        fmt += "\n\n"
        
        Opts = { 'cGroup1': idaapi.Form.ChkGroupControl
            (
                (
                "MemorySections",
                "MemoryContent",
                "RegisterValues",
                "DataTypes",
                "CodeBlocks",
                "DataDefinitions",
                "Comments",
                "EntryPoints",
                "Symbols",
                "Functions",
                "MemoryReferences",
                "StackReferences",
                "Manual"
                )
            )
        }
        
        self.options = idaapi.Form(fmt, Opts)
        self.options.Compile()
        
        self.options.MemorySections.checked   = True
        self.options.MemoryContent.checked    = True
        self.options.DataTypes.checked        = True
        self.options.RegisterValues.checked   = True
        self.options.CodeBlocks.checked       = True
        self.options.DataDefinitions.checked  = True
        self.options.Symbols.checked          = True
        self.options.EntryPoints.checked      = True
        self.options.Functions.checked        = True
        self.options.Comments.checked         = True
        self.options.MemoryReferences.checked = True
        self.options.StackReferences.checked  = False
        self.options.Manual.checked           = True
    
        if (self.autorun == False):
            ok = self.options.Execute()
            if (ok == 0):
                raise Cancelled


    def get_space_name(self, addr):
        """
        Returns the memory space name associated with an address.
        
        Args:
            addr: Integer representing a program address.
            
        Returns:
            String containg the memory space name.
            None if single address space architecture.

        Used for Harvard architectures (Intel 8051 and TMS, add others
        as needed). 
        """
        pid = idaapi.ph_get_id()
        stype = idaapi.segtype(addr)
        if pid == idaapi.PLFM_8051:
            if stype == idaapi.SEG_CODE:
                return "CODE"
            else:
                if stype == idaapi.SEG_IMEM:
                    iaddr = addr - idaapi.get_segm_base(idaapi.getseg(addr))
                    if iaddr < 0x80:
                        return "INTMEM"
                    else:
                        return "SFR"
                else:
                    return "EXTMEM"
        if pid == idaapi.PLFM_TMS:
            if stype == idaapi.SEG_CODE:
                return "CODE"
            else:
                return "DATA"
        return None


    def get_symbol_name(self, ea):
        """
        Returns the symbol name for the address.
        
        Args:
            ea: Integer representing the symbol address.
            
        Returns:
            String containing the symbol name.
        
        The demangled name will be returned if it exists, otherwise the
        displayed name is returned. Spaces (' ') will be replaced with '_'.
        """
        flags = idaapi.getFlags(ea)
        name = idaapi.get_demangled_name(BADADDR, ea, DEMANGLED_FORM,
                                         self.inf.demnames, False)
        if name == None or len(name) == 0 or name == "`string'":
            name = idaapi.get_name(BADADDR, ea)
        if name != None:
            name = name.replace(" ","_")
        return name


    def get_type(self, flags):
        """
        Returns a datatype string based on the item flags.
        
        Args:
            flags: IDA item flags.
            
        Returns:
            String representing item datatype.
        """
        if (self.cbsize == 2):
            if idaapi.isByte(flags) == True: return "word"
            if idaapi.isWord(flags) == True: return "dword"
        if idaapi.isByte(flags)     == True: return "byte"
        if idaapi.isWord(flags)     == True: return "word"
        if idaapi.isDwrd(flags)     == True: return "dword"
        if idaapi.isQwrd(flags)     == True: return "qword"
        if idaapi.isOwrd(flags)     == True: return "oword"
        if idaapi.isTbyt(flags)     == True: return "tbyte"
        if idaapi.isFloat(flags)    == True: return "float"
        if idaapi.isDouble(flags)   == True: return "double"
        if idaapi.isPackReal(flags) == True: return "packed"
        if idaapi.isASCII(flags)    == True: return "ascii"
        if idaapi.isStruct(flags)   == True: return "structure"
        if idaapi.isAlign(flags)    == True: return "align"
        return "unknown"


    def is_imm_op(self, addr, op):
        """
        Returns true if instruction operand at address is an immediate value.
        
        Args:
            addr: Integer representing instruction address.
            op: Integer representing operand index (0-based).
            
        Returns:
            True if instruction operand at address is an immediate value.
            False otherwise.
        """
        idaapi.decode_insn(addr)
        if (idaapi.cmd.Operands[op].type == idaapi.o_imm):
            return True
        return False
        

    def is_overlay(self, addr):
        """
        Checks if memory block (segment) is an overlay.
        
        Args:
            addr: Integer representing a program address.
            
        Returns:
            True if memory block (segment) is an overlay.
        """
        if idaapi.ph_get_id() == idaapi.PLFM_C166:
            return False
        s = idaapi.getseg(addr)
        if s.startEA in self.overlay:
            return self.overlay[s.startEA]
        return False
    
    def is_signed_data(self, flags):
        return (flags & idaapi.FF_SIGN) != 0


    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 start_element(self, tag, close=False):
        """
        Outputs the start of a new element on a new indented line.
        
        Args:
            tag: String representing the element tag
            close: Boolean indicating if tag is should be closed.
        """
        if idaapi.wasBreak() == True:
            raise Cancelled
        self.write_to_xmlfile("\n" + ("    " * self.indent_level) + "<" + tag)
        if (close):
            self.close_tag(True)
        # Add a counter for this tag; increment counter if it already exists.
        if (tag in self.elements) == False:
            self.elements[tag] = len(self.elements)
            self.counters.append(1)
            self.tags.append(tag)
        else:
            self.counters[self.elements[tag]] += 1
        

    def translate_address(self, addr):
        """
        Returns the translated logical address.

        The logical address is adjusted for the segment base address.
        For 16-bit segmented memory, return the 20-bit address.
        
        Args:
            addr: Integer representing a program address.
            
        Returns:
            Integer representing the logical address.
        """
        if self.seg_addr == False:
            return addr - idaapi.get_segm_base(idaapi.getseg(addr))
        base = idaapi.get_segm_para(idaapi.getseg(addr))
        return (base << 16) + (addr - (base << 4))
    

    def update_status(self, tag):
        """
        Displays the processing status in the IDA window.
        
        Args:
            tag: String representing XML element tag
        """
        status = "Exporting <" + tag + ">"
        idaapi.msg('\n%-35s' % status)
        idaapi.hide_wait_box()
        idaapi.show_wait_box(status)


    def write_address_attribute(self, name, addr):
        """
        Outputs an address attribute for an element.
        
        Args:
            name: String representing attribute name.
            addr: Integer representing a program address.
        """
        self.write_attribute(name, self.get_address_string(addr))
    

    def write_attribute(self, name, value):
        """
        Outputs an attribute (name and value) for an element.
        
        Args:
            name: String representing attribute name.
            value: String representing attribute value.
        """
        if name == None or value == None:
            return
        if (len(name) == 0) or (len(value) == 0):
            return
        attr = " " + name + '="' + self.check_for_entities(value) + '"'
        self.write_to_xmlfile(attr)
        

    def write_comment_element(self, name, cmt):
        """
        Outputs the tag and text for a comment element.
        Comment elements can be REGULAR_CMT, REPEATABLE_CMT, or TYPEINFO_CMT.
        
        Args:
            name: String representing the comment element name.
            cmt: String containing the comment.
        """
        self.start_element(name, True)
        self.write_text(cmt)
        self.end_element(name, False)
        

    def write_numeric_attribute(self, name, value, base=16, signedhex=False):
        """
        Outputs a numeric value attribute (name and value) for an element.
        
        Args:
            name: String representing the attribute name.
            value: Integer representing the attribute value.
            base: Integer representing numeric base to use for value.
            signedhex: Boolean indicating if hex representation of
                value is signed.
        """
        if base == 10:
            temp = "%d" % value
        else:
            if signedhex == True and value < 0:
                temp = "-0x%X" % abs(value)
            else:
                temp = "0x%X" % value
        self.write_attribute(name, temp)


    def write_text(self, text):
        """
        Outputs the parsed character text for an element.
        The text is checked for special characters.
        
        Args:
            text: String representing the element text.
        """
        self.write_to_xmlfile(self.check_for_entities(text))
    

    def write_to_xmlfile(self, buf):
        """
        Writes the buffer to the XML file.
        
        Args:
            buf: String containg data to write to XML file.
        """
        self.xmlfile.write(buf)
        self.dbg(buf)
    

    def write_xml_declaration(self):
        """
        Writes the XML Declarations at the start of the XML file.
        """
        self.dbg("\n")
        xml_declaration  = "<?xml version=\"1.0\" standalone=\"yes\"?>"
        xml_declaration += "\n<?program_dtd version=\"1\"?>\n"
        self.write_to_xmlfile(xml_declaration)

# Global constants
# mangled name inhibit flags are not exposed in idaapi
# inhibit flags for symbol names
# DEMANGLE_FORM (MNG_SHORT_FORM | MNG_NOBASEDT | MNG_NOCALLC | MNG_NOCSVOL)
DEMANGLED_FORM = 0x0ea3ffe7
# inhibit flags for typeinfo cmts
# DEMANGLED_TYPEINFO (MNG_LONG_FORM)
DEMANGLED_TYPEINFO = 0x06400007

# 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'


