# plugins

repository: carsond135/malwareanalysis
file: tools/IDA_Pro_7.7/IDA_Pro_7.7/plugins/BinCAT/python/idabincat/bcplugin.py
description: NKU恶意代码分析
wanted_name: BinCAT
comment: Interface to the BinCAT analyzer
created_at: 2023-06-25 19:16:29+00:00
pushed_at: 2022-12-18 09:34:14+00:00
url:
https://github.com/carsond135/malwareanalysis/blob/45c8bf862a54239ff1e7c91313eccf9cf7d4a35b/tools/IDA_Pro_7.7/IDA_Pro_7.7/plugins/BinCAT/python/idabincat/bcplugin.py

content:
# -*- coding: utf-8 -*-
# version IDA 6.9
# runs the "bincat" command from ida
"""
    This file is part of BinCAT.
    Copyright 2014-2017 - Airbus Group

    BinCAT is free software: you can redistribute it and/or modify
    it under the terms of the GNU Affero General Public License as published by
    the Free Software Foundation, either version 3 of the License, or (at your
    option) any later version.

    BinCAT is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU Affero General Public License for more details.

    You should have received a copy of the GNU Affero General Public License
    along with BinCAT.  If not, see <http://www.gnu.org/licenses/>.
"""

from collections.abc import MutableSequence
import hashlib
import logging
import os
import shutil
import sys
import tempfile
import traceback
import zlib
try:
    from configparser import NoSectionError
except ImportError:
    from ConfigParser import NoSectionError

# Ugly but IDA Python Linux doesn't have it !
try:
    import distutils.spawn
    no_spawn = False
except ImportError:
    no_spawn = True
try:
    import requests
except ImportError:
    # log message will be displayed later
    pass

import idaapi
from idaapi import NW_OPENIDB, NW_CLOSEIDB, NW_TERMIDA, NW_REMOVE
import idabincat.netnode
import idabincat.npkgen
from idabincat.plugin_options import PluginOptions
from idabincat.analyzer_conf import AnalyzerConfig, AnalyzerConfigurations, ConfigHelpers
from idabincat.gui import GUI, taint_color
import pybincat

from PyQt5 import QtCore
# used in idaapi
from PyQt5 import QtWidgets

# Logging
logging.basicConfig(level=logging.DEBUG)
bc_log = logging.getLogger('bincat.plugin')
bc_log.setLevel(logging.INFO)


def dedup_loglines(loglines, max=None):
    res = []
    staging = None
    n = 0

    def flush_staging():
        if n > 0:
            res.append(staging)
            if max and len(res) >= max:
                return True
        if n == 2:
            res.append(staging)
            if max and len(res) >= max:
                return True
        if n > 2:
            res.append("== (same line repeated %i times) ==" % n)
            if max and len(res) >= max:
                return True

    while True:
        if not loglines:
            flush_staging()
            break
        l = loglines.pop()
        if l == staging:
            n += 1
        else:
            if flush_staging():
                break
            staging = l
            n = 1
    res.reverse()
    return res


class AnalyzerUnavailable(Exception):
    pass


class BincatPlugin(idaapi.plugin_t):
    # variables required by IDA
    flags = 0  # normal plugin
    wanted_name = "BinCAT"
    wanted_hotkey = "Ctrl-Shift-B"
    comment = "Interface to the BinCAT analyzer"
    help = ""
    initialized = False

    def __init__(self):
        super(BincatPlugin, self).__init__()
        self.state = None

    # IDA API methods: init, run, term
    def init(self):
        info = idaapi.get_inf_structure()
        procname = info.procname
        if procname != 'metapc' and procname != 'ARM' and procname != 'ARMB' and procname != 'PPC':
            bc_log.info("CPU '%s' is not supported, not loading BinCAT", procname)
            return idaapi.PLUGIN_SKIP
        try:
            from pybincat import cfa as cfa_module
            global cfa_module
        except:
            bc_log.warning(
                "Failed to load 'pybincat.cfa' python module\n%s",
                repr(sys.exc_info()))
            return idaapi.PLUGIN_SKIP
        PluginOptions.init()

        # add plugin's 'bin' dir to PATH
        userdir = idaapi.get_user_idadir()
        bin_path = os.path.join(userdir, "plugins", "idabincat", "bin")
        if os.path.isdir(bin_path):
            path_env_sep = ';' if os.name == 'nt' else ':'
            os.environ['PATH'] += path_env_sep+bin_path
        if no_spawn:
            bc_exe = None
        else:
            # Check if bincat is available
            bc_exe = distutils.spawn.find_executable('bincat')
        if bc_exe is None:
            if no_spawn:
                bc_exe = os.path.join(bin_path, "bincat")
            else:
                bc_exe = distutils.spawn.find_executable('bincat')
        if bc_exe is None and no_spawn is False:
            bc_log.warning('Could not find bincat binary, will not be able to run local analysis')

        if PluginOptions.get("autostart") != "True":
            # will initialize later
            return idaapi.PLUGIN_OK

        bc_log.info("Autostarting")

        self.state = State()
        self.initialized = True
        bc_log.info("IDABinCAT ready.")
        return idaapi.PLUGIN_KEEP

    def run(self, args):
        if self.initialized:
            return
        self.state = State()
        # XXX move
        self.initialized = True
        self.state.gui.show_windows()
        bc_log.info("IDABinCAT ready.")

    def term(self):
        if self.state:
            bc_log.debug("Terminating BinCAT")
            self.state.clear_background()
            self.state.gui.term()
            self.state.gui = None
            self.state = None


class Analyzer(object):
    def __init__(self, path, finish_cb):
        self.path = path
        self.finish_cb = finish_cb

    def generate_tnpk(self, fname=None, destfname=None):
        """
        Generates TNPK file for provided fname. If None, generate one for the
        binary that is currently being analyzed in IDA, using IDA-provided
        headers.

        Returns file path to generated tnpk (string), or None if generation was
        not successful.
        """
        return None

    @property
    def initfname(self):
        return os.path.join(self.path, "init.ini")

    @property
    def outfname(self):
        return os.path.join(self.path, "out.ini")

    @property
    def cfainfname(self):
        return os.path.join(self.path, "cfain.marshal")

    @property
    def cfaoutfname(self):
        return os.path.join(self.path, "cfaout.marshal")

    @property
    def logfname(self):
        return os.path.join(self.path, "analyzer.log")

class LocalAnalyzerTimer(object):
    """
    IDA timer used to kill the BinCAT analyzer if the user
    cancels the analysis
    """
    def __init__(self, qprocess):
        self.interval = 500  # ms
        self.qprocess = qprocess
        self.timer = idaapi.register_timer(self.interval, self)
        if self.timer is None:
            raise RuntimeError("Failed to register timer")

    def __call__(self):
        if idaapi.user_cancelled():
            bc_log.info("Analysis cancelled by user, terminating process")
            self.qprocess.terminate()
            return -1
        else:
            return self.interval

    def destroy(self):
        idaapi.unregister_timer(self.timer)


class LocalAnalyzer(Analyzer, QtCore.QProcess):
    """
    Runs BinCAT locally using QProcess.
    """

    def __init__(self, *args, **kwargs):
        QtCore.QProcess.__init__(self)
        Analyzer.__init__(self, *args, **kwargs)
        # Qprocess signal handlers
        self.error.connect(self.procanalyzer_on_error)
        self.stateChanged.connect(self.procanalyzer_on_state_change)
        self.started.connect(self.procanalyzer_on_start)
        self.finished.connect(self.procanalyzer_on_finish)

    def generate_tnpk(self, fname=None, destfname=None):
        if fname:
            imports_data = open(fname, 'r').read()
        else:
            imports_data = ""
        try:
            npk_fname = idabincat.npkgen.NpkGen().generate_tnpk(
                imports_data=imports_data, destfname=destfname)
            return npk_fname
        except idabincat.npkgen.NpkGenException:
            bc_log.warning("Could not compile header file, "
                           "types from IDB will not be used for type propagation")
            return

    def run(self):
        idaapi.show_wait_box("Running analysis")

        # Create a timer to allow for cancellation
        self.timer = LocalAnalyzerTimer(self)

        cmdline = ["bincat",
                   [self.initfname, self.outfname, self.logfname]]
        # start the process
        bc_log.debug(
            "Analyzer cmdline: [%s %s]", cmdline[0], " ".join(cmdline[1]))
        try:
            self.start(*cmdline)
        except Exception as e:
            bc_log.error("BinCAT failed to launch the analyzer.py")
            bc_log.warning("Exception: %s\n%s", str(e), traceback.format_exc())
        else:
            bc_log.info("Analyzer started.")

    def procanalyzer_on_error(self, err):
        errors = ["Failed to start", "Crashed", "TimedOut", "Read Error",
                  "Write Error", "Unknown Error"]
        try:
            errtxt = errors[err]
        except IndexError:
            errtxt = "Unspecified error %s" % err
        bc_log.error("Analyzer error: %s", errtxt)
        self.process_output()

    def procanalyzer_on_state_change(self, new_state):
        states = ["Not running", "Starting", "Running"]
        bc_log.debug("Analyzer new state: %s", states[new_state])

    def procanalyzer_on_start(self):
        bc_log.info("Analyzer: starting process")

    def procanalyzer_on_finish(self):
        bc_log.info("Analyzer process finished")
        exitcode = self.exitCode()
        if exitcode != 0:
            bc_log.error("analyzer returned exit code=%i", exitcode)
        try:
            self.process_output()
        except Exception as e:
            idaapi.hide_wait_box()
            bc_log.error("Caught exception, hiding wait box", exc_info=True)

    def process_output(self):
        """
        Try to process analyzer output.
        """
        self.timer.destroy()
        bc_log.info("---- stdout ----------------")
        bc_log.info(str(self.readAllStandardOutput()))
        bc_log.info("---- stderr ----------------")
        bc_log.info(str(self.readAllStandardError()))
        bc_log.info("---- logfile ---------------")
        if os.path.exists(self.logfname):
            with open(self.logfname, 'r') as f:
                log_lines = f.readlines()
            log_lines = dedup_loglines(log_lines, max=100)
            if len(log_lines) > 100:
                bc_log.info("---- Only the last 100 log lines (deduped) are displayed here ---")
                bc_log.info("---- See full log in %s ---" % self.logfname)
            for line in log_lines:
                bc_log.info(line.rstrip())

        bc_log.info("====== end of logfile ======")
        idaapi.hide_wait_box()
        self.finish_cb(self.outfname, self.logfname, self.cfaoutfname)


class WebAnalyzer(Analyzer):
    API_VERSION = "1.2"

    def __init__(self, *args, **kwargs):
        Analyzer.__init__(self, *args, **kwargs)
        self.server_url = PluginOptions.get("server_url").rstrip("/")
        self.reachable_server = False
        self.check_version()  # raises exception if server is unreachable
        self.reachable_server = True

    def check_version(self):
        try:
            version_req = requests.get(self.server_url + "/version")
            srv_api_version = str(version_req.text)
        except requests.ConnectionError as e:
            raise AnalyzerUnavailable(
                "BinCAT server at %s could not be reached : %s" % (self.server_url, str(e)))
        if srv_api_version != WebAnalyzer.API_VERSION:
            raise AnalyzerUnavailable(
                "API mismatch: this plugin supports version %s, while server "
                "supports version %s." % (WebAnalyzer.API_VERSION,
                                          srv_api_version))
        return True

    def generate_tnpk(self, fname=None, destfname=None):
        if not self.reachable_server:
            return
        if not fname:
            headers_data = idabincat.npkgen.NpkGen().get_header_data()
            fname = os.path.join(self.path, "ida_generated_headers.h")
            with open(fname, 'w') as f:
                f.write(headers_data)
        sha256 = self.sha256_digest(fname)
        if not self.upload_file(fname, sha256):
            return
        npk_res = requests.post(self.server_url + "/convert_to_tnpk/" + sha256)
        if npk_res.status_code != 200:
            bc_log.warning("Error while compiling file to tnpk "
                           "on BinCAT analysis server, types from IDB will "
                           "not be used for type propagation")
            return
        res = npk_res.json()
        if 'status' not in res or res['status'] != 'ok':
            return
        sha256 = res['sha256']
        npk_contents = self.download_file(sha256)
        npk_fname = os.path.join(self.path, "headers.npk")
        with open(npk_fname, 'w') as f:
            f.write(npk_contents)
        return npk_fname

    def run(self):
        if not self.reachable_server:
            return
        if 'requests' not in sys.modules:
            bc_log.error("python module 'requests' could not be imported, "
                         "so remote BinCAT cannot be used.")
            return
        # create temporary AnalyzerConfig to replace referenced file names with
        # sha256 of their contents
        with open(self.initfname, 'r') as f:
            temp_config = AnalyzerConfig.load_from_str(f.read())
        # patch filepath - set filepath to sha256, upload file
        sha256 = self.sha256_digest(temp_config.binary_filepath)
        if not self.upload_file(temp_config.binary_filepath, sha256):
            return
        temp_config.binary_filepath = sha256
        # patch [imports] headers - replace with sha256, upload files
        files = temp_config.headers_files.split(',')
        h_shalist = []
        for f in files:
            if not f:
                continue
            try:
                f_sha256 = self.sha256_digest(f)
                h_shalist.append('"%s"' % f_sha256)
            except IOError as e:
                bc_log.error("Could not open file %s" % f, exc_info=1)
                return
            if not self.upload_file(f, f_sha256):
                bc_log.error("Could not upload file %s" % f, exc_info=1)
                return
        temp_config.headers_files = ','.join(h_shalist)

        # patch in_marshalled_cfa_file - replace with file contents sha256
        if os.path.exists(self.cfainfname):
            cfa_sha256 = self.sha256_digest(self.cfainfname)
            temp_config.in_marshalled_cfa_file = '"%s"' % cfa_sha256
        else:
            temp_config.in_marshalled_cfa_file = '"no-input-file"'
        temp_config.out_marshalled_cfa_file = '"cfaout.marshal"'
        # write patched config file
        init_ini_str = str(temp_config)
        # --- Run analysis
        run_res = requests.post(
            self.server_url + "/analyze",
            files={'init.ini': ('init.ini', init_ini_str)})
        if run_res.status_code != 200:
            bc_log.error("Error while uploading analysis configuration file "
                         "to BinCAT analysis server (%r)" % run_res.text)
            return
        files = run_res.json()
        if files["errorcode"]:
            bc_log.error("Error while analyzing file. Bincat output is:\n"
                         "----------------\n%s\n----------------"
                         % self.download_file(files["stdout.txt"]))
            if not files["out.ini"]:  # try to parse out.ini if it exists
                return
        with open(self.outfname, 'w') as outfp:

            outfp.write(self.download_file(files["out.ini"]))
        analyzer_log_contents = self.download_file(files["analyzer.log"])
        with open(self.logfname, 'w') as logfp:
            logfp.write(analyzer_log_contents)
        bc_log.info("---- stdout+stderr ----------------")
        bc_log.info(self.download_file(files["stdout.txt"]))
        bc_log.info("---- logfile ---------------")
        log_lines = analyzer_log_contents.split('\n')

        log_lines = dedup_loglines(log_lines, max=100)
        if len(log_lines) > 100:
            bc_log.info("---- Only the last 100 log lines (deduped) are displayed here ---")
            bc_log.info("---- See full log in %s ---" % self.logfname)
        for line in log_lines:
            bc_log.info(line.rstrip())

        bc_log.info("----------------------------")
        if "cfaout.marshal" in files:
            # might be absent (ex. when analysis failed, or not requested)
            with open(self.cfaoutfname, 'w') as cfaoutfp:
                cfaoutfp.write(self.download_file(files["cfaout.marshal"]))
        self.finish_cb(self.outfname, self.logfname, self.cfaoutfname)

    def download_file(self, fname):
        r = requests.get(self.server_url + '/download/' + fname + '/zlib')
        try:
            return zlib.decompress(r.content)
        except Exception as e:
            bc_log.error("Error uncompressing downloaded file [%s] (%s)" %
                         (fname, e))

    def sha256_digest(self, path):
        with open(path, 'rb') as f:
            h = hashlib.new('sha256')
            h.update(f.read())
            sha256 = h.hexdigest().lower()
        return sha256

    def upload_file(self, path, sha256):
        # check whether file has already been uploaded
        try:
            check_res = requests.head(
                self.server_url + "/download/%s" % sha256)
        except requests.exceptions.ConnectionError as e:
            bc_log.error("Error when contacting the BinCAT server: %s", e)
            return
        if check_res.status_code == 404:
            # Confirmation dialog before uploading file?
            msgBox = QtWidgets.QMessageBox()
            msgBox.setText("Do you really want to upload %s to %s?" %
                           (path, self.server_url))
            msgBox.setStandardButtons(QtWidgets.QMessageBox.Yes |
                                      QtWidgets.QMessageBox.No)
            msgBox.setIcon(QtWidgets.QMessageBox.Question)
            ret = msgBox.exec_()
            if ret == QtWidgets.QMessageBox.No:
                bc_log.info("Upload aborted.")
                return
            # upload
            upload_res = requests.put(
                self.server_url + "/add",
                files={'file': ('file', open(path, 'rb').read())})
            if upload_res.status_code != 200:
                bc_log.error("Error while uploading file %s "
                             "to BinCAT analysis server.", path)
                return
        return True


class State(object):
    """
    Container for (static) plugin state related data & methods.
    """
    def __init__(self):
        self.current_ea = None
        self.cfa = None
        self.current_node = None
        #: list of node ids (int)
        self.current_node_ids = []
        #: cfa.Unrel object or None
        self.current_unrel = None
        #: last run config
        self.current_config = None
        #: config to be edited
        self.edit_config = None
        #: Analyzer instance - protects against merciless garbage collector
        self.analyzer = None
        self.hooks = None
        self.netnode = idabincat.netnode.Netnode("$ com.bincat.bcplugin")
        #: acts as a List of (eip (int), "register name", "taint mask")
        #: XXX store in IDB?
        self.overrides = CallbackWrappedList()
        #: list of [@ or function names (str)]
        self.nops = CallbackWrappedList()
        #: list of [@ or function names (str), arg_nb, ret_val]
        self.skips = CallbackWrappedList()
        #: list of (name, config)
        self.configurations = AnalyzerConfigurations(self)
        # XXX store in idb after encoding?
        self.last_cfaout_marshal = None
        #: filepath to last dumped remapped binary
        self.remapped_bin_path = None
        self.remap_binary = True
        self.remapped_sections = None
        # for debugging purposes, to interact with this object from the console
        global bc_state
        bc_state = self

        self.gui = GUI(self)
        if PluginOptions.get("load_from_idb") == "True":
            self.load_from_idb()

    def new_analyzer(self, *args, **kwargs):
        """
        returns current Analyzer class (web or local)
        """
        if (PluginOptions.get("web_analyzer") == "True" and
                PluginOptions.get("server_url") != ""):
            if 'requests' not in sys.modules:
                bc_log.error("Trying to launch a remote analysis without the 'requests' module !")
                raise AnalyzerUnavailable()
            return WebAnalyzer(*args, **kwargs)
        else:
            return LocalAnalyzer(*args, **kwargs)

    def load_from_idb(self):
        if "out.ini" in self.netnode and "analyzer.log" in self.netnode:
            bc_log.info("Loading analysis results from idb")
            path = tempfile.mkdtemp(suffix='bincat')
            outfname = os.path.join(path, "out.ini")
            logfname = os.path.join(path, "analyzer.log")
            with open(outfname, 'w') as outfp:
                outfp.write(self.netnode["out.ini"])
            with open(logfname, 'w') as logfp:
                logfp.write(self.netnode["analyzer.log"])
            if "current_ea" in self.netnode:
                ea = self.netnode["current_ea"]
            else:
                ea = None
            self.analysis_finish_cb(outfname, logfname, cfaoutfname=None,
                                    ea=ea)
        if "remapped_bin_path" in self.netnode:
            fname = self.netnode["remapped_bin_path"]
            if os.path.isfile(fname):
                self.remapped_bin_path = fname
                if "remapped_sections" in self.netnode:
                    self.remapped_sections = self.netnode["remapped_sections"]
            else:
                bc_log.warning("cannot find remapped binary %s", fname)
        if "remap_binary" in self.netnode:
            self.remap_binary = self.netnode["remap_binary"]

    def clear_background(self):
        """
        reset background color for previous analysis
        """
        if self.cfa:
            color = idaapi.calc_bg_color(idaapi.NIF_BG_COLOR)
            for v in self.cfa.addr_nodes:
                ea = v.value
                idaapi.set_item_color(ea, color)

    def analysis_finish_cb(self, outfname, logfname, cfaoutfname, ea=None):
        idaapi.show_wait_box("HIDECANCEL\nParsing BinCAT analysis results")
        bc_log.debug("Parsing analyzer result file")
        # Here we can't check for user_cancelled because the UI is
        # unresponsive when parsing.
        try:
            cfa = cfa_module.CFA.parse(outfname, logs=logfname)
        except (pybincat.PyBinCATException, NoSectionError):
            idaapi.hide_wait_box()
            bc_log.error("Could not parse result file")
            return None
        self.clear_background()
        self.cfa = cfa
        if cfa:
            # XXX add user preference for saving to idb? in that case, store
            # reference to marshalled cfa elsewhere
            bc_log.info("Storing analysis results to idb...")
            with open(outfname, 'r') as f:
                self.netnode["out.ini"] = f.read()
            with open(logfname, 'r') as f:
                self.netnode["analyzer.log"] = f.read()
            if self.remapped_bin_path:
                self.netnode["remapped_bin_path"] = self.remapped_bin_path
                if self.remapped_sections:
                    self.netnode["remapped_sections"] = self.remapped_sections
            self.netnode["remap_binary"] = self.remap_binary
            if cfaoutfname is not None and os.path.isfile(cfaoutfname):
                with open(cfaoutfname, 'rb') as f:
                    self.last_cfaout_marshal = f.read()
            bc_log.info("Analysis results have been stored idb.")
        else:
            bc_log.info("Empty or unparseable result file.")
        bc_log.debug("----------------------------")
        idaapi.replace_wait_box("Updating IDB with BinCAT results")
        # Update current RVA to start address (nodeid = 0)
        # by default, use current ea - e.g, in case there is no results (cfa is
        # None) or no node 0 (happens in backward mode)
        current_ea = self.current_ea
        if ea is not None:
            current_ea = ea
        else:
            try:
                node0 = cfa['0']
                if node0:
                    current_ea = node0.address.value
            except (KeyError, TypeError):
                # no cfa is None, or no node0
                pass
        try:
            self.set_current_ea(current_ea, force=True)
        except TypeError as e:
            bc_log.warn("Could not load results from IDB")
            bc_log.warn("------ BEGIN EXCEPTION -----")
            bc_log.exception(e)
            bc_log.warn("------ END EXCEPTION -----")
            idaapi.hide_wait_box()
            return None
        self.netnode["current_ea"] = current_ea
        if not cfa:
            return
        for addr, nodeids in list(cfa.addr_nodes.items()):
            if hasattr(idaapi, "user_cancelled") and idaapi.user_cancelled() > 0:
                bc_log.info("User cancelled!")
                idaapi.hide_wait_box()
                return None
            ea = addr.value
            tainted = False
            taint_id = 1
            for n_id in nodeids:
                # is it tainted?
                # find child nodes
                node = cfa[n_id]
                if node.tainted:
                    tainted = True
                    if node.taintsrc:
                        # Take the first one
                        try:
                            taint_id = int(node.taintsrc[0].split("-")[1])
                        except IndexError: # Taint is TOP
                            taint_id = 0
                    break

            if tainted:
                idaapi.set_item_color(ea, taint_color(taint_id))
            else:
                idaapi.set_item_color(ea, 0xF0F0F0)
        idaapi.hide_wait_box()
        self.gui.focus_registers()

    def set_current_node(self, node_id, unrel_id=None):
        if self.cfa:
            node = self.cfa[node_id]
            if node:
                self.set_current_ea(
                    node.address.value,
                    force=True, node_id=node_id, unrel_id=unrel_id)

    def set_current_ea(self, ea, force=False, node_id=None, unrel_id=None):
        """
        :param ea: int or long
        """
        if not (force or ea != self.current_ea):
            return
        self.gui.before_change_ea()
        self.current_ea = ea
        empty_node = True
        if self.cfa:
            node_ids = self.cfa.node_id_from_addr(ea)
            if node_ids:
                if node_id in node_ids:
                    self.current_node = self.cfa[node_id]
                else:
                    self.current_node = self.cfa[node_ids[0]]
                self.current_node_ids = node_ids
                # find unrel_id
                if not unrel_id:
                    unrel_id = self.current_node.default_unrel_id()
                if unrel_id and unrel_id in self.current_node.unrels:
                    self.current_unrel = self.current_node.unrels[unrel_id]
                    empty_node = False
        if empty_node:
            self.current_node = None
            self.current_node_ids = []
            self.current_unrel = None

        self.gui.after_change_ea()

    def guess_filepath(self):
        filepath = self.current_config.binary_filepath
        if os.path.isfile(filepath):
            return filepath
        filepath = ConfigHelpers.guess_file_path()
        if os.path.isfile(filepath):
            return filepath
        # give up
        return None

    def start_analysis(self):
        """
        Creates new temporary dir. File structure:
        input files: init.ini, cfain.marshal
        output files: out.ini, cfaout.marshal
        """
        binary_filepath = self.guess_filepath()
        if not binary_filepath:
            bc_log.error(
                "File %s does not exit. Please fix path in configuration.",
                self.current_config.binary_filepath)
            return
        bc_log.debug("Using %s as source binary path", binary_filepath)
        self.current_config.binary_filepath = binary_filepath

        path = tempfile.mkdtemp(suffix='bincat')

        # instance variable: we don't want the garbage collector to delete the
        # *Analyzer instance, killing an unlucky QProcess in the process
        try:
            self.analyzer = self.new_analyzer(path, self.analysis_finish_cb)
        except AnalyzerUnavailable:
            bc_log.error("Analyzer is unavailable", exc_info=True)
            return

        bc_log.info("Current analyzer path: %s", path)

        # Update overrides
        self.current_config.update_overrides(self.overrides, self.nops,
                                             self.skips)
        analysis_method = self.current_config.analysis_method
        if analysis_method in ("forward_cfa", "backward"):
            if self.last_cfaout_marshal is None:
                bc_log.error("No marshalled CFA has been recorded - run a "
                             "forward analysis first.")
                return
            with open(self.analyzer.cfainfname, 'w') as f:
                f.write(self.last_cfaout_marshal)
        # Set correct file names
        # Note: bincat_native expects a filename for in_marshalled_cfa_file -
        # may not exist if analysis mode is forward_binary
        self.current_config.set_cfa_options('true', self.analyzer.cfainfname,
                                            self.analyzer.cfaoutfname)
        bc_log.debug("Generating .no files...")

        headers_filenames = self.current_config.headers_files.split(',')
        # compile .c files for libs, if there are any
        bc_log.debug("Initial header files: %r", headers_filenames)
        new_headers_filenames = []
        for f in headers_filenames:
            f = f.strip()
            if not f:
                continue
            if f.endswith('.c'):
                # generate .npk from .c file
                if not os.path.isfile(f):
                    bc_log.warning(
                        "header file %s could not be found, continuing", f)
                    continue
                new_npk_fname = f[:-2] + '.no'
                headers_filenames.remove(f)
                if not os.path.isfile(new_npk_fname):
                    # compile
                    self.analyzer.generate_tnpk(fname=f,
                                                destfname=new_npk_fname)
                    if not os.path.isfile(new_npk_fname):
                        bc_log.warning(
                            ".no file containing type data for the headers "
                            "file %s could not be generated, continuing", f)
                        continue
                    f = new_npk_fname
            # Relative paths are copied
            elif f.endswith('.no') and os.path.isfile(f):
                if f[0] != os.path.sep:
                    temp_npk_fname = os.path.join(path, os.path.basename(f))
                    shutil.copyfile(f, temp_npk_fname)
                    f = temp_npk_fname
            else:
                bc_log.warning(
                    "Header file %s does not exist or does not match expected "
                    "extensions (.c, .no), ignoring.", f)
                continue

            new_headers_filenames.append(f)
        headers_filenames = new_headers_filenames
        # generate npk file for the binary being analyzed (unless it has
        # already been generated)
        if not any(
                [s.endswith('pre-processed.no') for s in headers_filenames]):
            npk_filename = self.analyzer.generate_tnpk()
            if not npk_filename:
                bc_log.warning(
                    ".no file containing type data for the file being "
                    "analyzed could not be generated, continuing. The "
                    "ida-generated header could be invalid.")
            else:
                headers_filenames.append(npk_filename)
            bc_log.debug("Final npk files: %r", headers_filenames)
        self.current_config.headers_files = ','.join(headers_filenames)

        self.current_config.write(self.analyzer.initfname)
        self.analyzer.run()

    def re_run(self):
        """
        Re-run analysis, taking new overrides settings into account
        """
        if self.current_config is None:
            # XXX upload all required files in Web Analyzer
            # XXX Store all required files in IDB
            bc_log.error(
                "You have to run the analysis first using the 'Analyze From "
                "here (Ctrl-Shift-A)' menu - reloading previous results from "
                "IDB is not yet supported.")
        else:
            self.start_analysis()

    def add_or_replace_override(self, ea, htext, mask):
        for idx, (e, h, _) in enumerate(self.overrides):
            if e == ea and h == htext:
                del self.overrides[idx]
        self.overrides.append((ea, htext, mask))


class CallbackWrappedList(MutableSequence):
    """
    Acts as a List object, wraps write access with calls to properly invalidate
    models associated with View GUI objects.
    Should store only immutable objects.
    """
    def __init__(self):
        self._data = []
        #: list of functions to be called prior to updating list
        self.pre_callbacks = []
        #: list of functions to be called after updating list
        self.post_callbacks = []
        #: cache

    def register_callbacks(self, pre_cb, post_cb):
        if pre_cb:
            self.pre_callbacks.append(pre_cb)
        if post_cb:
            self.post_callbacks.append(post_cb)

    def __getitem__(self, index):
        # no need for wrappers - values should be immutable
        return self._data[index]

    def _callback_wrap(f):
        def wrap(self, *args, **kwargs):
            for cb in self.pre_callbacks:
                cb()
            f(self, *args, **kwargs)
            for cb in self.post_callbacks:
                cb()
        return wrap

    @_callback_wrap
    def __setitem__(self, index, value):
        self._data[index] = tuple(value)

    @_callback_wrap
    def __delitem__(self, item):
        del self._data[item]

    @_callback_wrap
    def clear(self):
        self._data = []

    def __len__(self):
        return len(self._data)

    @_callback_wrap
    def insert(self, index, value):
        self._data.insert(index, tuple(value))


def PLUGIN_ENTRY():
    return BincatPlugin()


repository: PyEGo/exp-gist
file: hard-gists-me/c6936d219db8e6635d25/snippet.py
description: Experiment results on hard-gists dataset
wanted_name: simpa64
comment: Simplifier
created_at: 2021-04-18 11:16:06+00:00
pushed_at: 2022-11-18 17:22:11+00:00
url:
https://github.com/PyEGo/exp-gist/blob/bb8212e435ab525bb8ebdee2f620f850133a9b99/hard-gists-me/c6936d219db8e6635d25/snippet.py

content:
#  AArch64 mov simplifier IDA plugin
#
#  Copyright (c) 2015 xerub
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# based on Rolf Rolles x86 deobfuscator http://www.msreverseengineering.com

import idaapi
import idc

ARM64_MOVE_I = idaapi.ARM_mov

def dump_cmd(cmd):
	print "cs = %lx" % cmd.cs
	print "ip = %lx" % cmd.ip
	print "ea = %lx" % cmd.ea
	print "itype = %lx" % cmd.itype
	print "size = %lx" % cmd.size
	print "auxpref = %lx" % cmd.auxpref
	print "segpref = %lx" % cmd.segpref
	print "insnpref = %lx" % cmd.insnpref
	print "flags = %lx" % cmd.flags

def dump_op(op):
	print "n = %lx" % op.n
	print "type = %lx" % op.type
	print "offb = %lx" % op.offb
	print "offo = %lx" % op.offo
	print "flags = %lx" % op.flags
	print "dtyp = %lx" % op.dtyp
	print "reg = %lx" % op.reg
	print "phrase = %lx" % op.phrase
	print "value = %lx" % op.value
	print "addr = %lx" % op.addr
	print "specval = %lx" % op.specval
	print "specflag1 = %lx" % op.specflag1
	print "specflag2 = %lx" % op.specflag2
	print "specflag3 = %lx" % op.specflag3
	print "specflag4 = %lx" % op.specflag4

def HighestSetBit(N, imm):
	i = N - 1
	while i >= 0:
		if imm & (1 << i):
			return i
		i -= 1
	return -1

def ZeroExtendOnes(M, N):				# zero extend M ones to N width
	return (1 << M) - 1

def RORZeroExtendOnes(M, N, R):
	val = ZeroExtendOnes(M, N)
	return ((val >> R) & ((1 << (N - R)) - 1)) | ((val & ((1 << R) - 1)) << (N - R))

def Replicate(val, bits):
	ret = val
	shift = bits
	while shift < 64:				# XXX actually, it is either 32 or 64
		ret |= (val << shift)
		shift += bits
	return ret

def DecodeBitMasks(immN, imms, immr, immediate):
	len = HighestSetBit(7, (immN << 6) | (~imms & 0x3F))
	if len < 1:
		return None
	levels = ZeroExtendOnes(len, 6)
	if immediate and (imms & levels) == levels:
		return None
	S = imms & levels
	R = immr & levels
	esize = 1 << len
	return Replicate(RORZeroExtendOnes(S + 1, esize, R), esize)

def DecodeMov(opcode, total, first):
	# opc
	o = (opcode >> 29) & 3
	# constant
	k = (opcode >> 23) & 0x3F

	if k == 0x24 and o == 1:			# MOV (bitmask imm) <=> ORR (immediate)
		# sf
		s = (opcode >> 31) & 1
		# N
		N = (opcode >> 22) & 1
		if s == 0 and N != 0:
			return None
		# rn
		rn = (opcode >> 5) & 0x1F
		if rn == 31:
			imms = (opcode >> 10) & 0x3F
			immr = (opcode >> 16) & 0x3F
			return DecodeBitMasks(N, imms, immr, True)
	elif k == 0x25:					# MOVN/MOVZ/MOVK
		# sf
		s = (opcode >> 31) & 1
		# hw
		h = (opcode >> 21) & 3
		# imm16
		i = (opcode >> 5) & 0xFFFF
		if s == 0 and h > 1:
			return None
		h *= 16
		i <<= h
		if o == 0:				# MOVN
			return ~i
		elif o == 2:				# MOVZ
			return i
		elif o == 3 and not first:		# MOVK
			return (total & ~(0xFFFF << h)) | i
	elif (k | 1) == 0x23 and not first:		# ADD (immediate)
		# shift
		h = (opcode >> 22) & 3
		if h > 1:
			return None
		# rn
		rd = opcode & 0x1F
		rn = (opcode >> 5) & 0x1F
		if rd != rn:
			return None
		# imm12
		i = (opcode >> 10) & 0xFFF
		h *= 12
		i <<= h
		if o & 2:				# SUB
			return total - i
		else:					# ADD
			return total + i

	return None

def check_mov_sequence(ea):
	oldea = ea
	reg = -1
	total = 0
	is64 = False
	while idaapi.getseg(ea).use64():
		d = idaapi.get_long(ea)
		# reg
		r = d & 0x1F
		if reg >= 0 and reg != r:
			break
		newval = DecodeMov(d, total, reg < 0)
		if newval is None:
			break
		if reg >= 0 and idaapi.get_first_fcref_to(ea) != idaapi.BADADDR:
			break
		if (d >> 31) & 1:
			is64 = True
		total = newval
		reg = r
		ea += 4
	return ea - oldea, reg, is64, total

def is_my_mov(cmd):
	if cmd.itype == ARM64_MOVE_I and cmd.flags == idaapi.INSN_MACRO and cmd.size > 4:
		return True
	return False

def check_ubfm_shift(ea):
	if idaapi.getseg(ea).use64():
		opcode = idaapi.get_long(ea)
		# opc
		o = (opcode >> 29) & 3
		# constant
		k = (opcode >> 23) & 0x3F
		if (o & 1) == 0 and k == 0x26:
			# sf
			s = (opcode >> 31) & 1
			# N
			N = (opcode >> 22) & 1
			if s == N:
				# imm
				imms = (opcode >> 10) & 0x3F
				immr = (opcode >> 16) & 0x3F
				mask = 0x1F | ((s & N) << 5)
				if imms == mask:
					return idaapi.ARM_lsr if o else idaapi.ARM_asr, opcode, s, immr
				elif immr == imms + 1:
					return idaapi.ARM_lsl if o else idaapi.ARM_null, opcode, s, mask - imms
	return idaapi.ARM_null, 0, 0, 0

class simpA64Hook(idaapi.IDP_Hooks):
	def __init__(self):
		idaapi.IDP_Hooks.__init__(self)
		self.n = idaapi.netnode("$ A64 Simplifier",0,1)

	def custom_ana(self):
		len, reg, is64, imm = check_mov_sequence(idaapi.cmd.ea)
		if len > 4:
			#print "0x%x: MOV/MOVK %c%d, #0x%x" % (idaapi.cmd.ea, 'X' if is64 else 'W', reg, imm)
			#dump_cmd(idaapi.cmd)
			#dump_op(idaapi.cmd.Op1)
			#dump_op(idaapi.cmd.Op2)
			idaapi.cmd.itype = ARM64_MOVE_I
			idaapi.cmd.segpref = 14 # ARM Condition = ALways
			idaapi.cmd.Op1.type = idaapi.o_reg
			idaapi.cmd.Op1.dtyp = idaapi.dt_qword if is64 else idaapi.dt_dword
			idaapi.cmd.Op1.reg = reg + 129 # Use Wn/Xn registers instead of Rn
			idaapi.cmd.Op2.type = idaapi.o_imm
			idaapi.cmd.Op2.dtyp = idaapi.dt_qword if is64 else idaapi.dt_dword
			idaapi.cmd.Op2.value = imm
			idaapi.cmd.flags = idaapi.INSN_MACRO
			idaapi.cmd.size = len
			return True
		insn, regs, is64, shift = check_ubfm_shift(idaapi.cmd.ea)
		if insn != idaapi.ARM_null:
			idaapi.cmd.itype = insn
			idaapi.cmd.segpref = 14
			idaapi.cmd.Op1.type = idaapi.o_reg
			idaapi.cmd.Op1.dtyp = idaapi.dt_qword if is64 else idaapi.dt_dword
			idaapi.cmd.Op1.reg = (regs & 0x1F) + 129
			idaapi.cmd.Op2.type = idaapi.o_reg
			idaapi.cmd.Op2.dtyp = idaapi.dt_qword if is64 else idaapi.dt_dword
			idaapi.cmd.Op2.reg = ((regs >> 5) & 0x1F) + 129
			idaapi.cmd.Op3.type = idaapi.o_imm
			idaapi.cmd.Op3.dtyp = idaapi.dt_qword if is64 else idaapi.dt_dword
			idaapi.cmd.Op3.value = shift
			idaapi.cmd.size = 4
			return True
		return False

	def custom_mnem(self): # totally optional
		if is_my_mov(idaapi.cmd):
			return "MOVE"
		return None

#	def custom_out(self): # XXX ida would just append .EQ
#		if is_my_mov(idaapi.cmd):
#			buf = idaapi.init_output_buffer(1024)
#			idaapi.OutMnem(16, "")
#			idaapi.out_one_operand(0)
#			idaapi.out_symbol(',')
#			idaapi.OutChar(' ')
#			idaapi.out_one_operand(1)
#			idaapi.term_output_buffer()
#			idaapi.MakeLine(buf)
#			return True
#		return False

#	def custom_outop(self, op): # XXX ida would just use Rn
#		if is_my_mov(idaapi.cmd) and op.type == idaapi.o_reg:
#			idaapi.out_register("%c%d" % ('X' if op.dtyp == idaapi.dt_qword else 'W', op.reg))
#			return True
#		return False

class simpa64_t(idaapi.plugin_t):
	flags = idaapi.PLUGIN_PROC
	comment = "Simplifier"
	wanted_hotkey = "Alt-Z"
	help = "Runs transparently"
	wanted_name = "simpa64"
	hook = None
	enabled = 1

	def init(self):
		self.hook = None
		if idaapi.ph_get_id() != idaapi.PLFM_ARM or idaapi.BADADDR <= 0xFFFFFFFF:
			return idaapi.PLUGIN_SKIP

		self.hook = simpA64Hook()
		flag = self.hook.n.altval(0)
		if flag:
			self.enabled = flag - 1
		print "%s is %sabled" % (self.wanted_name, "en" if self.enabled else "dis")
		if self.enabled:
			self.hook.hook()
		return idaapi.PLUGIN_KEEP

	def run(self, arg):
		print "%sabling %s" % ("dis" if self.enabled else "en", self.wanted_name)
		if self.enabled:
			self.hook.unhook()
		else:
			self.hook.hook()
		self.enabled = self.enabled ^ 1
		self.hook.n.altset(0, self.enabled + 1)
		idc.Refresh()

	def term(self):
		if self.hook:
			self.hook.unhook()

def PLUGIN_ENTRY():
	return simpa64_t()


repository: PyEGo/exp-gist
file: hard-gists-ego/c6936d219db8e6635d25/snippet.py
description: Experiment results on hard-gists dataset
wanted_name: simpa64
comment: Simplifier
created_at: 2021-04-18 11:16:06+00:00
pushed_at: 2022-11-18 17:22:11+00:00
url:
https://github.com/PyEGo/exp-gist/blob/bb8212e435ab525bb8ebdee2f620f850133a9b99/hard-gists-ego/c6936d219db8e6635d25/snippet.py

content:
#  AArch64 mov simplifier IDA plugin
#
#  Copyright (c) 2015 xerub
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# based on Rolf Rolles x86 deobfuscator http://www.msreverseengineering.com

import idaapi
import idc

ARM64_MOVE_I = idaapi.ARM_mov

def dump_cmd(cmd):
	print "cs = %lx" % cmd.cs
	print "ip = %lx" % cmd.ip
	print "ea = %lx" % cmd.ea
	print "itype = %lx" % cmd.itype
	print "size = %lx" % cmd.size
	print "auxpref = %lx" % cmd.auxpref
	print "segpref = %lx" % cmd.segpref
	print "insnpref = %lx" % cmd.insnpref
	print "flags = %lx" % cmd.flags

def dump_op(op):
	print "n = %lx" % op.n
	print "type = %lx" % op.type
	print "offb = %lx" % op.offb
	print "offo = %lx" % op.offo
	print "flags = %lx" % op.flags
	print "dtyp = %lx" % op.dtyp
	print "reg = %lx" % op.reg
	print "phrase = %lx" % op.phrase
	print "value = %lx" % op.value
	print "addr = %lx" % op.addr
	print "specval = %lx" % op.specval
	print "specflag1 = %lx" % op.specflag1
	print "specflag2 = %lx" % op.specflag2
	print "specflag3 = %lx" % op.specflag3
	print "specflag4 = %lx" % op.specflag4

def HighestSetBit(N, imm):
	i = N - 1
	while i >= 0:
		if imm & (1 << i):
			return i
		i -= 1
	return -1

def ZeroExtendOnes(M, N):				# zero extend M ones to N width
	return (1 << M) - 1

def RORZeroExtendOnes(M, N, R):
	val = ZeroExtendOnes(M, N)
	return ((val >> R) & ((1 << (N - R)) - 1)) | ((val & ((1 << R) - 1)) << (N - R))

def Replicate(val, bits):
	ret = val
	shift = bits
	while shift < 64:				# XXX actually, it is either 32 or 64
		ret |= (val << shift)
		shift += bits
	return ret

def DecodeBitMasks(immN, imms, immr, immediate):
	len = HighestSetBit(7, (immN << 6) | (~imms & 0x3F))
	if len < 1:
		return None
	levels = ZeroExtendOnes(len, 6)
	if immediate and (imms & levels) == levels:
		return None
	S = imms & levels
	R = immr & levels
	esize = 1 << len
	return Replicate(RORZeroExtendOnes(S + 1, esize, R), esize)

def DecodeMov(opcode, total, first):
	# opc
	o = (opcode >> 29) & 3
	# constant
	k = (opcode >> 23) & 0x3F

	if k == 0x24 and o == 1:			# MOV (bitmask imm) <=> ORR (immediate)
		# sf
		s = (opcode >> 31) & 1
		# N
		N = (opcode >> 22) & 1
		if s == 0 and N != 0:
			return None
		# rn
		rn = (opcode >> 5) & 0x1F
		if rn == 31:
			imms = (opcode >> 10) & 0x3F
			immr = (opcode >> 16) & 0x3F
			return DecodeBitMasks(N, imms, immr, True)
	elif k == 0x25:					# MOVN/MOVZ/MOVK
		# sf
		s = (opcode >> 31) & 1
		# hw
		h = (opcode >> 21) & 3
		# imm16
		i = (opcode >> 5) & 0xFFFF
		if s == 0 and h > 1:
			return None
		h *= 16
		i <<= h
		if o == 0:				# MOVN
			return ~i
		elif o == 2:				# MOVZ
			return i
		elif o == 3 and not first:		# MOVK
			return (total & ~(0xFFFF << h)) | i
	elif (k | 1) == 0x23 and not first:		# ADD (immediate)
		# shift
		h = (opcode >> 22) & 3
		if h > 1:
			return None
		# rn
		rd = opcode & 0x1F
		rn = (opcode >> 5) & 0x1F
		if rd != rn:
			return None
		# imm12
		i = (opcode >> 10) & 0xFFF
		h *= 12
		i <<= h
		if o & 2:				# SUB
			return total - i
		else:					# ADD
			return total + i

	return None

def check_mov_sequence(ea):
	oldea = ea
	reg = -1
	total = 0
	is64 = False
	while idaapi.getseg(ea).use64():
		d = idaapi.get_long(ea)
		# reg
		r = d & 0x1F
		if reg >= 0 and reg != r:
			break
		newval = DecodeMov(d, total, reg < 0)
		if newval is None:
			break
		if reg >= 0 and idaapi.get_first_fcref_to(ea) != idaapi.BADADDR:
			break
		if (d >> 31) & 1:
			is64 = True
		total = newval
		reg = r
		ea += 4
	return ea - oldea, reg, is64, total

def is_my_mov(cmd):
	if cmd.itype == ARM64_MOVE_I and cmd.flags == idaapi.INSN_MACRO and cmd.size > 4:
		return True
	return False

def check_ubfm_shift(ea):
	if idaapi.getseg(ea).use64():
		opcode = idaapi.get_long(ea)
		# opc
		o = (opcode >> 29) & 3
		# constant
		k = (opcode >> 23) & 0x3F
		if (o & 1) == 0 and k == 0x26:
			# sf
			s = (opcode >> 31) & 1
			# N
			N = (opcode >> 22) & 1
			if s == N:
				# imm
				imms = (opcode >> 10) & 0x3F
				immr = (opcode >> 16) & 0x3F
				mask = 0x1F | ((s & N) << 5)
				if imms == mask:
					return idaapi.ARM_lsr if o else idaapi.ARM_asr, opcode, s, immr
				elif immr == imms + 1:
					return idaapi.ARM_lsl if o else idaapi.ARM_null, opcode, s, mask - imms
	return idaapi.ARM_null, 0, 0, 0

class simpA64Hook(idaapi.IDP_Hooks):
	def __init__(self):
		idaapi.IDP_Hooks.__init__(self)
		self.n = idaapi.netnode("$ A64 Simplifier",0,1)

	def custom_ana(self):
		len, reg, is64, imm = check_mov_sequence(idaapi.cmd.ea)
		if len > 4:
			#print "0x%x: MOV/MOVK %c%d, #0x%x" % (idaapi.cmd.ea, 'X' if is64 else 'W', reg, imm)
			#dump_cmd(idaapi.cmd)
			#dump_op(idaapi.cmd.Op1)
			#dump_op(idaapi.cmd.Op2)
			idaapi.cmd.itype = ARM64_MOVE_I
			idaapi.cmd.segpref = 14 # ARM Condition = ALways
			idaapi.cmd.Op1.type = idaapi.o_reg
			idaapi.cmd.Op1.dtyp = idaapi.dt_qword if is64 else idaapi.dt_dword
			idaapi.cmd.Op1.reg = reg + 129 # Use Wn/Xn registers instead of Rn
			idaapi.cmd.Op2.type = idaapi.o_imm
			idaapi.cmd.Op2.dtyp = idaapi.dt_qword if is64 else idaapi.dt_dword
			idaapi.cmd.Op2.value = imm
			idaapi.cmd.flags = idaapi.INSN_MACRO
			idaapi.cmd.size = len
			return True
		insn, regs, is64, shift = check_ubfm_shift(idaapi.cmd.ea)
		if insn != idaapi.ARM_null:
			idaapi.cmd.itype = insn
			idaapi.cmd.segpref = 14
			idaapi.cmd.Op1.type = idaapi.o_reg
			idaapi.cmd.Op1.dtyp = idaapi.dt_qword if is64 else idaapi.dt_dword
			idaapi.cmd.Op1.reg = (regs & 0x1F) + 129
			idaapi.cmd.Op2.type = idaapi.o_reg
			idaapi.cmd.Op2.dtyp = idaapi.dt_qword if is64 else idaapi.dt_dword
			idaapi.cmd.Op2.reg = ((regs >> 5) & 0x1F) + 129
			idaapi.cmd.Op3.type = idaapi.o_imm
			idaapi.cmd.Op3.dtyp = idaapi.dt_qword if is64 else idaapi.dt_dword
			idaapi.cmd.Op3.value = shift
			idaapi.cmd.size = 4
			return True
		return False

	def custom_mnem(self): # totally optional
		if is_my_mov(idaapi.cmd):
			return "MOVE"
		return None

#	def custom_out(self): # XXX ida would just append .EQ
#		if is_my_mov(idaapi.cmd):
#			buf = idaapi.init_output_buffer(1024)
#			idaapi.OutMnem(16, "")
#			idaapi.out_one_operand(0)
#			idaapi.out_symbol(',')
#			idaapi.OutChar(' ')
#			idaapi.out_one_operand(1)
#			idaapi.term_output_buffer()
#			idaapi.MakeLine(buf)
#			return True
#		return False

#	def custom_outop(self, op): # XXX ida would just use Rn
#		if is_my_mov(idaapi.cmd) and op.type == idaapi.o_reg:
#			idaapi.out_register("%c%d" % ('X' if op.dtyp == idaapi.dt_qword else 'W', op.reg))
#			return True
#		return False

class simpa64_t(idaapi.plugin_t):
	flags = idaapi.PLUGIN_PROC
	comment = "Simplifier"
	wanted_hotkey = "Alt-Z"
	help = "Runs transparently"
	wanted_name = "simpa64"
	hook = None
	enabled = 1

	def init(self):
		self.hook = None
		if idaapi.ph_get_id() != idaapi.PLFM_ARM or idaapi.BADADDR <= 0xFFFFFFFF:
			return idaapi.PLUGIN_SKIP

		self.hook = simpA64Hook()
		flag = self.hook.n.altval(0)
		if flag:
			self.enabled = flag - 1
		print "%s is %sabled" % (self.wanted_name, "en" if self.enabled else "dis")
		if self.enabled:
			self.hook.hook()
		return idaapi.PLUGIN_KEEP

	def run(self, arg):
		print "%sabling %s" % ("dis" if self.enabled else "en", self.wanted_name)
		if self.enabled:
			self.hook.unhook()
		else:
			self.hook.hook()
		self.enabled = self.enabled ^ 1
		self.hook.n.altset(0, self.enabled + 1)
		idc.Refresh()

	def term(self):
		if self.hook:
			self.hook.unhook()

def PLUGIN_ENTRY():
	return simpa64_t()


repository: PyEGo/exp-gist
file: hard-gists-reqs/c6936d219db8e6635d25/snippet.py
description: Experiment results on hard-gists dataset
wanted_name: simpa64
comment: Simplifier
created_at: 2021-04-18 11:16:06+00:00
pushed_at: 2022-11-18 17:22:11+00:00
url:
https://github.com/PyEGo/exp-gist/blob/bb8212e435ab525bb8ebdee2f620f850133a9b99/hard-gists-reqs/c6936d219db8e6635d25/snippet.py


repository: PyEGo/exp-gist
file: hard-gists-ego-s2/c6936d219db8e6635d25/snippet.py
description: Experiment results on hard-gists dataset
wanted_name: simpa64
comment: Simplifier
created_at: 2021-04-18 11:16:06+00:00
pushed_at: 2022-11-18 17:22:11+00:00
url:
https://github.com/PyEGo/exp-gist/blob/bb8212e435ab525bb8ebdee2f620f850133a9b99/hard-gists-ego-s2/c6936d219db8e6635d25/snippet.py


repository: vAlerainTech/vAlerain-Ark
file: cmake-build-debug/IDA_Pro_7.7/python/examples/hexrays/vds6.py
wanted_name: Hex-Rays space remover (IDAPython)
comment: Sample plugin6 for Hex-Rays decompiler
created_at: 2024-03-25 14:26:36+00:00
pushed_at: 2024-07-27 05:07:46+00:00
url:
https://github.com/vAlerainTech/vAlerain-Ark/blob/df8bcafeda091d30ceade7bb2c7c73b5ae351a38/cmake-build-debug/IDA_Pro_7.7/python/examples/hexrays/vds6.py

content:
"""
summary: superficially modify the decompilation output

description:
  modifies the decompilation output in a superficial manner,
  by removing some white spaces

  Note: this is rather crude, not quite "pythonic" code.
"""

from __future__ import print_function

import idautils
import idc
import ida_idaapi
import ida_hexrays
import ida_lines

def dbg(msg):
    #print(msg)
    pass

def is_cident_char(c):
    return c.isalnum() or c == '_'

def my_tag_skipcodes(l, storage):
    n = ida_lines.tag_skipcodes(l)
    dbg("Skipping %d chars ('%s')" % (n, l[0:n]))
    storage.append(l[0:n])
    return l[n:]

def remove_spaces(sl):
    dbg("*" * 80)
    l = sl.line
    out = []

    def push(c):
        dbg("Appending '%s'" % c)
        out.append(c)

    # skip initial spaces, do not compress them
    while True:
        l = my_tag_skipcodes(l, out)
        if not l:
            break
        c = l[0]
        if not c.isspace():
            break
        push(c)
        l = l[1:]

    # remove all spaces except in string and char constants
    delim = None # if not None, then we are skipping until 'delim'
    last = None # last seen character
    while True:
        # go until comments
        if l.startswith("//"):
            push(l)
            break
        dbg("-" * 60)
        nchars = ida_lines.tag_advance(l, 1)
        push(l[0:nchars])
        l = l[nchars:]
        l = my_tag_skipcodes(l, out)
        if not l:
            break
        c = l[0]
        dbg("c: '%s', last: '%s', l: '%s'" % (c, last, l))
        if delim:
            # we're inside a literal.
            if c == delim:
                delim = None # literal ended
        elif c == '"' or c == "'":
            delim = c # string/char literal started
        elif c.isspace():
            end = l.lstrip()
            nptr = my_tag_skipcodes(end, out)
            dbg("end: '%s', nptr: '%s'" % (end, nptr))
            # do not concatenate idents
            if not is_cident_char(last) or not is_cident_char(nptr[0]):
                l = end
                c = l[0] if l else ''
                dbg("new l: '%s'" % l)
        last = l[0] if l else ''

    sl.line = "".join(out)


class vds6_hooks_t(ida_hexrays.Hexrays_Hooks):
    def func_printed(self, cfunc):
        for sl in cfunc.get_pseudocode():
            remove_spaces(sl);
        return 0

# a plugin interface, boilerplate code
class my_plugin_t(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_HIDE
    wanted_name = "Hex-Rays space remover (IDAPython)"
    wanted_hotkey = ""
    comment = "Sample plugin6 for Hex-Rays decompiler"
    help = ""
    def init(self):
        if ida_hexrays.init_hexrays_plugin():
            self.vds6_hooks = vds6_hooks_t()
            self.vds6_hooks.hook()
            return ida_idaapi.PLUGIN_KEEP # keep us in the memory
    def term(self):
        self.vds6_hooks.unhook()
    def run(self, arg):
        pass

def PLUGIN_ENTRY():
    return my_plugin_t()


repository: BraveCattle/ctf
file: tools/ida_pro/IDA_Pro_7.7/python/examples/hexrays/vds6.py
wanted_name: Hex-Rays space remover (IDAPython)
comment: Sample plugin6 for Hex-Rays decompiler
created_at: 2024-01-23 08:17:51+00:00
pushed_at: 2024-10-23 09:31:30+00:00
url:
https://github.com/BraveCattle/ctf/blob/5858d0da5c048eb8468ff8be1d63fbc50d61b601/tools/ida_pro/IDA_Pro_7.7/python/examples/hexrays/vds6.py

content:
"""
summary: superficially modify the decompilation output

description:
  modifies the decompilation output in a superficial manner,
  by removing some white spaces

  Note: this is rather crude, not quite "pythonic" code.
"""

from __future__ import print_function

import idautils
import idc
import ida_idaapi
import ida_hexrays
import ida_lines

def dbg(msg):
    #print(msg)
    pass

def is_cident_char(c):
    return c.isalnum() or c == '_'

def my_tag_skipcodes(l, storage):
    n = ida_lines.tag_skipcodes(l)
    dbg("Skipping %d chars ('%s')" % (n, l[0:n]))
    storage.append(l[0:n])
    return l[n:]

def remove_spaces(sl):
    dbg("*" * 80)
    l = sl.line
    out = []

    def push(c):
        dbg("Appending '%s'" % c)
        out.append(c)

    # skip initial spaces, do not compress them
    while True:
        l = my_tag_skipcodes(l, out)
        if not l:
            break
        c = l[0]
        if not c.isspace():
            break
        push(c)
        l = l[1:]

    # remove all spaces except in string and char constants
    delim = None # if not None, then we are skipping until 'delim'
    last = None # last seen character
    while True:
        # go until comments
        if l.startswith("//"):
            push(l)
            break
        dbg("-" * 60)
        nchars = ida_lines.tag_advance(l, 1)
        push(l[0:nchars])
        l = l[nchars:]
        l = my_tag_skipcodes(l, out)
        if not l:
            break
        c = l[0]
        dbg("c: '%s', last: '%s', l: '%s'" % (c, last, l))
        if delim:
            # we're inside a literal.
            if c == delim:
                delim = None # literal ended
        elif c == '"' or c == "'":
            delim = c # string/char literal started
        elif c.isspace():
            end = l.lstrip()
            nptr = my_tag_skipcodes(end, out)
            dbg("end: '%s', nptr: '%s'" % (end, nptr))
            # do not concatenate idents
            if not is_cident_char(last) or not is_cident_char(nptr[0]):
                l = end
                c = l[0] if l else ''
                dbg("new l: '%s'" % l)
        last = l[0] if l else ''

    sl.line = "".join(out)


class vds6_hooks_t(ida_hexrays.Hexrays_Hooks):
    def func_printed(self, cfunc):
        for sl in cfunc.get_pseudocode():
            remove_spaces(sl);
        return 0

# a plugin interface, boilerplate code
class my_plugin_t(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_HIDE
    wanted_name = "Hex-Rays space remover (IDAPython)"
    wanted_hotkey = ""
    comment = "Sample plugin6 for Hex-Rays decompiler"
    help = ""
    def init(self):
        if ida_hexrays.init_hexrays_plugin():
            self.vds6_hooks = vds6_hooks_t()
            self.vds6_hooks.hook()
            return ida_idaapi.PLUGIN_KEEP # keep us in the memory
    def term(self):
        self.vds6_hooks.unhook()
    def run(self, arg):
        pass

def PLUGIN_ENTRY():
    return my_plugin_t()


repository: carsond135/malwareanalysis
file: tools/IDA_Pro_7.7/IDA_Pro_7.7/python/examples/hexrays/vds6.py
description: NKU恶意代码分析
wanted_name: Hex-Rays space remover (IDAPython)
comment: Sample plugin6 for Hex-Rays decompiler
created_at: 2023-06-25 19:16:29+00:00
pushed_at: 2022-12-18 09:34:14+00:00
url:
https://github.com/carsond135/malwareanalysis/blob/45c8bf862a54239ff1e7c91313eccf9cf7d4a35b/tools/IDA_Pro_7.7/IDA_Pro_7.7/python/examples/hexrays/vds6.py


repository: Varshith-JV-1410/IDA-pro-MCP-server
file: ida-mcp-server.py
wanted_name: IDA MCP Server
comment: IDA Model Context Protocol Server
created_at: 2025-05-03 07:08:43+00:00
pushed_at: 2025-05-06 09:40:58+00:00
url:
https://github.com/Varshith-JV-1410/IDA-pro-MCP-server/blob/dd34e41f54621bc7c21735b8f1b4b8b42bae8085/ida-mcp-server.py

content:
import glob
import json
import os
import ida_bytes
import ida_ua
import ida_funcs
import ida_hexrays
import ida_name
import ida_segment
import idautils
import idc
import ida_idaapi
import ida_kernwin
import idaapi
import threading

from typing import Dict, List, Optional, Any, Tuple
from datetime import datetime
from functools import wraps
from starlette.middleware import Middleware
from starlette.middleware.cors import CORSMiddleware
from starlette.applications import Starlette
from starlette.requests import Request
from starlette.routing import Mount, Route
from starlette.responses import Response
from mcp.server import Server
from mcp.server.sse import SseServerTransport
from mcp.server import FastMCP
import uvicorn


# Initialize FastMCP server for IDA tools
mcp = FastMCP("IDA MCP Server", port=3000)

# 封裝函數執行在主線程的裝飾器
def execute_on_main_thread(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        result = []
        exception = []
        
        def run_function():
            try:
                result.append(f(*args, **kwargs))
            except Exception as e:
                exception.append(e)
            return 0
        
        ida_kernwin.execute_sync(run_function, ida_kernwin.MFF_FAST)
        
        if exception:
            raise exception[0]
        return result[0]
    return wrapper


@mcp.tool()
@execute_on_main_thread
def get_bytes(ea: int, size: int) -> List[int]:
    """Get bytes at specified address.

    Args:
        ea: Effective address to read from
        size: Number of bytes to read
    """
    try:
        result = [ida_bytes.get_byte(ea + i) for i in range(size)]
        return result
    except Exception as e:
        print(f"Error in get_bytes: {str(e)}")
        return {"error": str(e)}


@mcp.tool()
@execute_on_main_thread
def get_disasm(ea: int) -> str:
    """Get disassembly at specified address.

    Args:
        ea: Effective address to disassemble
    """
    return idc.generate_disasm_line(ea, 0)


@mcp.tool()
@execute_on_main_thread
def get_decompiled_func(ea: int) -> Dict[str, Any]:
    """Get decompiled pseudocode of function containing address.

    Args:
        ea: Effective address within the function
    """
    try:
        func = ida_funcs.get_func(ea)
        if not func:
            return {"error": "No function found at address"}

        decompiler = ida_hexrays.decompile(func.start_ea)
        if not decompiler:
            return {"error": "Failed to decompile function"}

        return {"code": str(decompiler)}
    except Exception as e:
        return {"error": str(e)}


@mcp.tool()
@execute_on_main_thread
def get_function_name(ea: int) -> str:
    """Get function name at specified address.

    Args:
        ea: Effective address of the function
    """
    return ida_name.get_name(ea)


@mcp.tool()
@execute_on_main_thread
def get_segments() -> List[Dict[str, Any]]:
    """Get all segments information.

    @return: List of segments (start, end, name, class, perm, bitness, align, comb, type, sel, flags)
    """
    segments = []
    n = 0
    seg = ida_segment.getnseg(n)
    while seg:
        segments.append(
            {
                "start": seg.start_ea,
                "end": seg.end_ea,
                "name": ida_segment.get_segm_name(seg),
                "class": ida_segment.get_segm_class(seg),
                "perm": seg.perm,
                "bitness": seg.bitness,
                "align": seg.align,
                "comb": seg.comb,
                "type": seg.type,
                "sel": seg.sel,
                "flags": seg.flags,
            }
        )
        n += 1
        seg = ida_segment.getnseg(n)
    return segments


@mcp.tool()
@execute_on_main_thread
def get_functions() -> List[Dict[str, Any]]:
    """Get all functions in the binary."""
    functions = []
    for func_ea in idautils.Functions():
        func_name = ida_name.get_name(func_ea)
        functions.append({"address": func_ea, "name": func_name})
    return functions


@mcp.tool()
@execute_on_main_thread
def get_xrefs_to(ea: int) -> List[Dict[str, Any]]:
    """Get all cross references to specified address.

    Args:
        ea: Effective address to find references to
    """
    xrefs = []
    for xref in idautils.XrefsTo(ea, 0):
        xrefs.append({"from": xref.frm, "type": xref.type})
    return xrefs


@mcp.tool()
@execute_on_main_thread
def get_imports() -> dict[str, list[tuple[int, str, int]]]:
    """Get all imports in the binary.

    Args:
        None

    Returns:
        A dictionary where the keys are module names and the values are lists of tuples.
        Each tuple contains the address of the imported function, the name of the function,
        and the ordinal value of the function.
    """
    tree = {}
    nimps = idaapi.get_import_module_qty()

    for i in range(0, nimps):
        name = idaapi.get_import_module_name(i)
        if not name:
            continue
        # Create a list for imported names
        items = []

        def imports_names_cb(ea, name, ord):
            items.append((ea, "" if not name else name, ord))
            # True -> Continue enumeration
            return True

        # Enum imported entries in this module
        idaapi.enum_import_names(i, imports_names_cb)

        if name not in tree:
            tree[name] = []
        tree[name].extend(items)

    return tree


@mcp.tool()
@execute_on_main_thread
def get_exports() -> List[Tuple[int, int, int, str]]:
    """Get all exports in the binary.

    @return: List of tuples (index, ordinal, ea, name)
    """
    return list(idautils.Entries())


@mcp.tool()
@execute_on_main_thread
def get_entry_point() -> int:
    """Get the entry point of the binary."""
    try:
        import ida_ida
        return ida_ida.inf_get_start_ea()
    except (ImportError, AttributeError):
        try:
            # Alternative method: idc.get_inf_attr to get
            import idc
            return idc.get_inf_attr(idc.INF_START_EA)
        except (ImportError, AttributeError):
            # Last alternative method: use cvar.inf
            return idaapi.cvar.inf.start_ea


@mcp.tool()
@execute_on_main_thread
def make_function(ea: int) -> None:
    """Make a function at specified address."""
    ida_funcs.add_func(ea)


@mcp.tool()
@execute_on_main_thread
def undefine_function(ea: int) -> None:
    """Undefine a function at specified address."""
    ida_funcs.del_func(ea)


@mcp.tool()
@execute_on_main_thread
def get_dword_at(ea: int) -> int:
    """Get the dword at specified address."""
    return idc.get_dword(ea)


@mcp.tool()
@execute_on_main_thread
def get_word_at(ea: int) -> int:
    """Get the word at specified address."""
    return idc.get_word(ea)


@mcp.tool()
@execute_on_main_thread
def get_byte_at(ea: int) -> int:
    """Get the byte at specified address."""
    return idc.get_byte(ea)


@mcp.tool()
@execute_on_main_thread
def get_qword_at(ea: int) -> int:
    """Get the qword at specified address."""
    return idc.get_qword(ea)


@mcp.tool()
@execute_on_main_thread
def get_float_at(ea: int) -> float:
    """Get the float at specified address."""
    return idc.get_float(ea)


@mcp.tool()
@execute_on_main_thread
def get_double_at(ea: int) -> float:
    """Get the double at specified address."""
    return idc.get_double(ea)


@mcp.tool()
@execute_on_main_thread
def get_string_at(ea: int) -> str:
    """Get the string at specified address."""
    return idc.get_strlit_contents(ea)


@mcp.tool()
@execute_on_main_thread
def get_strings():
    strings = []
    for s in idautils.Strings():
        strings.append({"address": s.ea, "string": str(s)})
    return strings

@mcp.tool()
@execute_on_main_thread
def get_current_file_path():
    return idc.get_input_file_path()

@mcp.tool()
@execute_on_main_thread
def list_files_with_relative_path(relative_path: str = ""):
    base_dir = os.path.dirname(idc.get_input_file_path())
    if  ':' in relative_path or '..' in relative_path or '//' in relative_path:
        return json.dumps({"error": "Invalid relative path"})
    if relative_path is None or relative_path == "":
        return glob.glob(os.path.join(base_dir, "*"))
    else:
        return glob.glob(os.path.join(base_dir, relative_path, "*"))

@mcp.tool()
@execute_on_main_thread
def read_file(relative_path: str):
    base_dir = os.path.dirname(idc.get_input_file_path())
    if  ':' in relative_path or '..' in relative_path or '//' in relative_path:
        return json.dumps({"error": "Invalid relative path"})
    if relative_path is "":
        return json.dumps({"error": "Relative path is required"})
    with open(os.path.join(base_dir, relative_path), "r") as f:
        return f.read()

@mcp.tool()
@execute_on_main_thread
def write_file(relative_path: str, content: str):
    base_dir = os.path.dirname(idc.get_input_file_path())
    if  ':' in relative_path or '..' in relative_path or '//' in relative_path:
        return json.dumps({"error": "Invalid relative path"})
    if relative_path is "":
        return json.dumps({"error": "Relative path is required"})
    with open(os.path.join(base_dir, relative_path), "w") as f:
        f.write(content)

@mcp.tool()
@execute_on_main_thread
def read_binary(relative_path: str):
    base_dir = os.path.dirname(idc.get_input_file_path())
    if  ':' in relative_path or '..' in relative_path or '//' in relative_path:
        return json.dumps({"error": "Invalid relative path"})
    if relative_path is "":
        return json.dumps({"error": "Relative path is required"})
    with open(os.path.join(base_dir, relative_path), "rb") as f:
        return f.read()

@mcp.tool()
@execute_on_main_thread
def write_binary(relative_path: str , content: bytes):
    base_dir = os.path.dirname(idc.get_input_file_path())
    if  ':' in relative_path or '..' in relative_path or '//' in relative_path:
        return json.dumps({"error": "Invalid relative path"})
    if relative_path is "":
        return json.dumps({"error": "Relative path is required"})
    with open(os.path.join(base_dir, relative_path), "wb") as f:
        f.write(content)    

@mcp.tool()
@execute_on_main_thread
def eval_pythoni(script: str):
    return eval(script)

@mcp.tool()
@execute_on_main_thread
def get_instruction_length(address: int) -> int:
    """
    Retrieves the length (in bytes) of the instruction at the specified address.

    Args:
        address: The address of the instruction.

    Returns:
        The length (in bytes) of the instruction.  Returns 0 if the instruction cannot be decoded.
    """
    try:
        # Create an insn_t object to store instruction information.
        insn = ida_ua.insn_t()

        # Decode the instruction.
        length = ida_ua.decode_insn(insn, address)
        if length == 0:
            print(f"Failed to decode instruction at address {hex(address)}")
            return 0

        return length
    except Exception as e:
        print(f"Error getting instruction length: {str(e)}")
        return 0

@mcp.prompt()
def binary_analysis_strategy() -> str:
    """
    Guild for analyzing the binary
    """
    return (
        "IDA Pro MCP Server Tools and Best Practices:\n\n."
        "Tools: \n"
        "- get_bytes: Get bytes at specified address.\n"
        "- get_disasm: Get disassembly at specified address.\n"
        "- get_decompiled_func: Get decompiled pseudocode of function containing address.\n"
        "- get_function_name: Get function name at specified address.\n"
        "- get_segments: Get all segments information.\n"
        "- get_functions: Get all functions in the binary.\n"
        "- get_xrefs_to: Get all cross references to a specified address.\n"
        "- get_imports: Get all imports in the binary.\n"
        "- get_exports: Get all exports in the binary.\n"
        "- get_entry_point: Get the entry point of the binary.\n"
        "- make_function: Make a function at specified address.\n"
        "- undefine_function: Undefine a function at specified address.\n"
        "- get_dword_at: Get the dword at specified address.\n"
        "- get_word_at: Get the word at specified address.\n"
        "- get_byte_at: Get the byte at specified address.\n"
        "- get_qword_at: Get the qword at specified address.\n"
        "- get_float_at: Get the float at specified address.\n"
        "- get_double_at: Get the double at specified address.\n"
        "- get_string_at: Get the string at specified address.\n"
        "- get_strings: Get all strings in the binary.\n"
        "- get_current_file_path: Get the current path of the binary.\n"
        "- list_files_with_relative_path: List all files in the specified relative path in the current directory.\n"
        "- read_file: Read the content of a file.\n"
        "- write_file: Write content to a file.\n"
        "- read_binary: Read the content of a binary file.\n"
        "- write_binary: Write content to a binary file.\n"
        "- eval_python: Evaluate a Python script in IDA Pro.\n"
        "- get_instruction_length: Get the length of the instruction at the specified address.\n"
        "Best Practices: \n"
        "- Initial Analysis Phase\n"
        "   1. Examine the Entry Point\n"
        "       - Use the get_entry_point() tool to locate the program's entry point\n"
        "       - Analyze the code at the entry point to understand the program's startup flow\n"
        "       - Look for unusual instructions or jumps\n"
        "   2. Analyze Import Table\n"
        "       - Use the get_imports() tool to view all imported functions\n"
        "       - Look for suspicious API functions, such as:\n"
        "           - File operations: CreateFile, WriteFile\n"
        "           - Network communication: socket, connect, InternetOpen\n"
        "           - Process manipulation: CreateProcess, VirtualAlloc\n"
        "           - Registry operations: RegOpenKey, RegSetValue\n"
        "           - Cryptography related: CryptEncrypt, CryptDecrypt\n"
        "   3. Review Strings\n"
        "       - Use the get_strings() tool to obtain all strings\n"
        "       - Pay attention to IP addresses, URLs, domain names, file paths\n"
        "       - Look for encrypted or obfuscated string patterns\n"
        "       - Analyze command line parameters and error messages\n"
        "   4. In-Depth Analysis Phase\n"
        "       - Track Key API Calls\n"
        "           - Use get_xrefs_to() to find cross-references to suspicious imported functions\n"
        "           - Use get_decompiled_func() to analyze functions that call these APIs\n"
        "           - Analyze how parameters and return values are handled\n"
        "   5. Identify Main Functional Blocks\n"
        "       - Use get_functions() to get a list of all functions\n"
        "       - Sort functions by size and complexity\n"
        "       - Decompile and analyze large, complex functions\n"
        "       - Look for suspicious function names or unnamed functions\n"
        "   6. Analyze Control Flow\n"
        "       - Observe conditional branches and loop structures\n"
        "       - Analyze function call graphs and execution paths\n"
        "       - Look for anti-debugging and anti-VM detection techniques\n"
        "       - Pay attention to unusual jumps and callback mechanisms\n"
        "   7. Identifying Malicious Behaviors\n"
        "       - Identify Common Malicious Functionality\n"
        "           - Persistence mechanisms: Registry modifications, startup items, service creation\n"
        "           - Data theft: File searching, keylogging, screen capturing\n"
        "           - Communication features: C&C communication, data exfiltration channels\n"
        "           - Evasion techniques: Obfuscation, packing, anti-analysis checks\n"
        "           - Destructive behaviors: File encryption, system damage\n"
        "   8. Analyze Algorithms and Encryption Routines\n"
        "       - Identify encryption and decryption functions\n"
        "       - Look for hardcoded keys and cryptographic constants\n"
        "       - Analyze how data is processed in memory\n"
        "   9. Analyze Network Communication\n"
        "       - Identify network communication functions\n"
        "       - Look for IP addresses, URLs, and domain names\n"
        "       - Analyze how data is sent and received over the network\n"
        "   10. Dump payloads if there is any decryption or encoding\n"
        "       - Generate decryption or decodeing script in python and ida python\n"
        "       - Use eval_python to execute the script in IDA Pro\n"
        "       - Dump payloads in the current directory\n"
        "   11. Document Analysis Results\n"
        "       - Add comments to key functions\n"
        "       - Rename functions to reflect their actual functionality\n"
        "       - Create a logical structure diagram of the code\n"
        "   12. Using Advanced Techniques\n"
        "       - Use IDA Pro's advanced features like IDAPython scripting, IDA SDK, and IDA API\n"
        "       - Implement custom analysis scripts to automate repetitive tasks\n"
        "       - Explore IDA Pro's plugin ecosystem for additional analysis capabilities\n"
    )


def create_starlette_app(mcp_server: Server, *, debug: bool = False) -> Starlette:
    """Create a Starlette application that can serve the provided mcp server with SSE."""

    middleware = [
        Middleware(
            CORSMiddleware,
            allow_origins=["*"],
            allow_methods=["*"],
            allow_headers=["*"],
        )
    ]
    return Starlette(
        debug=debug,
        middleware=middleware,
        routes=[
            Mount("/", app=mcp.sse_app()),
        ],
    )


class ModelContextProtocolPlugin(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_FIX | ida_idaapi.PLUGIN_HIDE
    comment = "IDA Model Context Protocol Server"
    help = "Provides REST API and SSE for IDA Pro analysis"
    wanted_name = "IDA MCP Server"
    wanted_hotkey = ""

    def init(self):
        try:
            print("Initializing IDA Model Context Protocol Server...")
            # app = create_starlette_app(mcp, debug=True)

            def run_server():
                try:
                    # 設置將異常轉換為 JSON 響應
                    mcp.run(transport="sse")
                    # uvicorn.run(app, host="localhost", port=3000, log_level="debug")
                except Exception as e:
                    print(f"Server error: {str(e)}")

            server_thread = threading.Thread(target=run_server)
            server_thread.daemon = True
            server_thread.start()
            print("Server started successfully!")
            return ida_idaapi.PLUGIN_KEEP
        except Exception as e:
            print(f"Failed to start server: {str(e)}")
            return ida_idaapi.PLUGIN_SKIP

    def run(self, arg):
        pass

    def term(self):
        print("Terminating IDA Model Context Protocol Server...")


def PLUGIN_ENTRY():
    return ModelContextProtocolPlugin()


if __name__ == "__main__":
    PLUGIN_ENTRY()


repository: vAlerainTech/vAlerain-Ark
file: cmake-build-debug/IDA_Pro_7.7/python/examples/core/install_user_defined_prefix.py
wanted_name: user defined prefix
comment: This is a user defined prefix sample plugin
created_at: 2024-03-25 14:26:36+00:00
pushed_at: 2024-07-27 05:07:46+00:00
url:
https://github.com/vAlerainTech/vAlerain-Ark/blob/df8bcafeda091d30ceade7bb2c7c73b5ae351a38/cmake-build-debug/IDA_Pro_7.7/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.
"""

from __future__ import print_function

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: BraveCattle/ctf
file: tools/ida_pro/IDA_Pro_7.7/python/examples/core/install_user_defined_prefix.py
wanted_name: user defined prefix
comment: This is a user defined prefix sample plugin
created_at: 2024-01-23 08:17:51+00:00
pushed_at: 2024-10-23 09:31:30+00:00
url:
https://github.com/BraveCattle/ctf/blob/5858d0da5c048eb8468ff8be1d63fbc50d61b601/tools/ida_pro/IDA_Pro_7.7/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.
"""

from __future__ import print_function

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: carsond135/malwareanalysis
file: tools/IDA_Pro_7.7/IDA_Pro_7.7/python/examples/core/install_user_defined_prefix.py
description: NKU恶意代码分析
wanted_name: user defined prefix
comment: This is a user defined prefix sample plugin
created_at: 2023-06-25 19:16:29+00:00
pushed_at: 2022-12-18 09:34:14+00:00
url:
https://github.com/carsond135/malwareanalysis/blob/45c8bf862a54239ff1e7c91313eccf9cf7d4a35b/tools/IDA_Pro_7.7/IDA_Pro_7.7/python/examples/core/install_user_defined_prefix.py


repository: vAlerainTech/vAlerain-Ark
file: cmake-build-debug/IDA_Pro_7.7/python/examples/hexrays/vds21.py
wanted_name: Hex-Rays custom prototype provider (IDAPython)
comment: Sample plugin21 for Hex-Rays decompiler
created_at: 2024-03-25 14:26:36+00:00
pushed_at: 2024-07-27 05:07:46+00:00
url:
https://github.com/vAlerainTech/vAlerain-Ark/blob/df8bcafeda091d30ceade7bb2c7c73b5ae351a38/cmake-build-debug/IDA_Pro_7.7/python/examples/hexrays/vds21.py

content:
"""
summary: dynamically provide a custom call type

description:
  This plugin can greatly improve decompilation of indirect calls:

      call    [eax+4]

  For them, the decompiler has to guess the prototype of the called function.
  This has to be done at a very early phase of decompilation because
  the function prototype influences the data flow analysis. On the other
  hand, we do not have global data flow analysis results yet because
  we haven't analyzed all calls in the function. It is a chicked-and-egg
  problem.

  The decompiler uses various techniques to guess the called function
  prototype. While it works very well, it may fail in some cases.

  To fix, the user can specify the call prototype manually, using
  "Edit, Operand types, Set operand type" at the call instruction.

  This plugin illustrates another approach to the problem:
  if you happen to be able to calculate the call prototypes dynamically,
  this is how to inform the decompiler about them.
"""

import ida_idaapi
import ida_nalt
import ida_kernwin
import ida_typeinf
import ida_hexrays

testing = False # only for testing purposes

class callinfo_provider_t(ida_hexrays.Hexrays_Hooks):

    # this callback will be called for all call instructions
    # our plugin may provide the function prototype or even a complete new callinfo
    # object. The callinfo object may be useful if the prototype is not enough
    # to express all details of the call.
    def build_callinfo(self, blk, type):
        # it is a good idea to skip direct calls.
        # note that some indirect calls may be resolved and become direct calls,
        # and will be filtered out here:
        ida_kernwin.msg("%x: got called for: %s\n" % (blk.tail.ea, blk.tail.dstr()))
        tail = blk.tail
        if tail.opcode == ida_hexrays.m_call:
            return

        # also, if the type was specified by the user, do not interfere
        call_ea = tail.ea
        tif = ida_typeinf.tinfo_t()
        if ida_nalt.get_op_tinfo(tif, call_ea, 0):
            return

        global testing
        if not testing:
            # ok, the decompiler really has to guess the type.
            # just for the sake of an example, return a predefined prototype.
            # in real life you will provide the prototype you discovered yourself,
            # using your magic of yours :)
            my_proto = "int f();"
            ida_kernwin.msg("%x: providing prototype %s\n" % (call_ea, my_proto))
            ida_typeinf.parse_decl(type, None, my_proto, 0)
        else:
            # as an alternative to filling 'type', you can
            # choose to return a mcallinfo_t instance.
            mi = ida_hexrays.mcallinfo_t()
            mi.cc = ida_typeinf.CM_CC_STDCALL | ida_typeinf.CM_N32_F48 # let's use stdcall to differentiate from the tinfo_t-filling version
            mi.return_type = ida_typeinf.tinfo_t(ida_typeinf.BT_INT)
            mi.return_argloc.set_reg1(0) # eax
            return mi

# a plugin interface, boilerplate code
class my_plugin_t(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_HIDE
    wanted_name = "Hex-Rays custom prototype provider (IDAPython)"
    wanted_hotkey = ""
    comment = "Sample plugin21 for Hex-Rays decompiler"
    help = ""
    def init(self):
        if ida_hexrays.init_hexrays_plugin():
            self.hooks = callinfo_provider_t()
            self.hooks.hook()
            ida_kernwin.warning(
                "Installed callinfo provider sample (vds21.py)\n" +\
                "Please note that it is just an example\n" +\
                "and will spoil your decompilations!")
            return ida_idaapi.PLUGIN_KEEP # keep us in the memory
    def term(self):
        self.hooks.unhook()
    def run(self, arg):
        pass

def PLUGIN_ENTRY():
    return my_plugin_t()


repository: BraveCattle/ctf
file: tools/ida_pro/IDA_Pro_7.7/python/examples/hexrays/vds21.py
wanted_name: Hex-Rays custom prototype provider (IDAPython)
comment: Sample plugin21 for Hex-Rays decompiler
created_at: 2024-01-23 08:17:51+00:00
pushed_at: 2024-10-23 09:31:30+00:00
url:
https://github.com/BraveCattle/ctf/blob/5858d0da5c048eb8468ff8be1d63fbc50d61b601/tools/ida_pro/IDA_Pro_7.7/python/examples/hexrays/vds21.py

content:
"""
summary: dynamically provide a custom call type

description:
  This plugin can greatly improve decompilation of indirect calls:

      call    [eax+4]

  For them, the decompiler has to guess the prototype of the called function.
  This has to be done at a very early phase of decompilation because
  the function prototype influences the data flow analysis. On the other
  hand, we do not have global data flow analysis results yet because
  we haven't analyzed all calls in the function. It is a chicked-and-egg
  problem.

  The decompiler uses various techniques to guess the called function
  prototype. While it works very well, it may fail in some cases.

  To fix, the user can specify the call prototype manually, using
  "Edit, Operand types, Set operand type" at the call instruction.

  This plugin illustrates another approach to the problem:
  if you happen to be able to calculate the call prototypes dynamically,
  this is how to inform the decompiler about them.
"""

import ida_idaapi
import ida_nalt
import ida_kernwin
import ida_typeinf
import ida_hexrays

testing = False # only for testing purposes

class callinfo_provider_t(ida_hexrays.Hexrays_Hooks):

    # this callback will be called for all call instructions
    # our plugin may provide the function prototype or even a complete new callinfo
    # object. The callinfo object may be useful if the prototype is not enough
    # to express all details of the call.
    def build_callinfo(self, blk, type):
        # it is a good idea to skip direct calls.
        # note that some indirect calls may be resolved and become direct calls,
        # and will be filtered out here:
        ida_kernwin.msg("%x: got called for: %s\n" % (blk.tail.ea, blk.tail.dstr()))
        tail = blk.tail
        if tail.opcode == ida_hexrays.m_call:
            return

        # also, if the type was specified by the user, do not interfere
        call_ea = tail.ea
        tif = ida_typeinf.tinfo_t()
        if ida_nalt.get_op_tinfo(tif, call_ea, 0):
            return

        global testing
        if not testing:
            # ok, the decompiler really has to guess the type.
            # just for the sake of an example, return a predefined prototype.
            # in real life you will provide the prototype you discovered yourself,
            # using your magic of yours :)
            my_proto = "int f();"
            ida_kernwin.msg("%x: providing prototype %s\n" % (call_ea, my_proto))
            ida_typeinf.parse_decl(type, None, my_proto, 0)
        else:
            # as an alternative to filling 'type', you can
            # choose to return a mcallinfo_t instance.
            mi = ida_hexrays.mcallinfo_t()
            mi.cc = ida_typeinf.CM_CC_STDCALL | ida_typeinf.CM_N32_F48 # let's use stdcall to differentiate from the tinfo_t-filling version
            mi.return_type = ida_typeinf.tinfo_t(ida_typeinf.BT_INT)
            mi.return_argloc.set_reg1(0) # eax
            return mi

# a plugin interface, boilerplate code
class my_plugin_t(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_HIDE
    wanted_name = "Hex-Rays custom prototype provider (IDAPython)"
    wanted_hotkey = ""
    comment = "Sample plugin21 for Hex-Rays decompiler"
    help = ""
    def init(self):
        if ida_hexrays.init_hexrays_plugin():
            self.hooks = callinfo_provider_t()
            self.hooks.hook()
            ida_kernwin.warning(
                "Installed callinfo provider sample (vds21.py)\n" +\
                "Please note that it is just an example\n" +\
                "and will spoil your decompilations!")
            return ida_idaapi.PLUGIN_KEEP # keep us in the memory
    def term(self):
        self.hooks.unhook()
    def run(self, arg):
        pass

def PLUGIN_ENTRY():
    return my_plugin_t()


repository: carsond135/malwareanalysis
file: tools/IDA_Pro_7.7/IDA_Pro_7.7/python/examples/hexrays/vds21.py
description: NKU恶意代码分析
wanted_name: Hex-Rays custom prototype provider (IDAPython)
comment: Sample plugin21 for Hex-Rays decompiler
created_at: 2023-06-25 19:16:29+00:00
pushed_at: 2022-12-18 09:34:14+00:00
url:
https://github.com/carsond135/malwareanalysis/blob/45c8bf862a54239ff1e7c91313eccf9cf7d4a35b/tools/IDA_Pro_7.7/IDA_Pro_7.7/python/examples/hexrays/vds21.py


repository: xx-an/DSV
file: lib/ghidra_9.0.4/Extensions/IDAPro/Python/6xx/loaders/xmlldr.py
wanted_name: XML Importer
comment: Import XML PROGRAM file
created_at: 2021-01-28 12:41:04+00:00
pushed_at: 2021-04-04 23:05:31+00:00
url:
https://github.com/xx-an/DSV/blob/439c3e3835f4d8533d53d8fb040c066e1ff97b78/lib/ghidra_9.0.4/Extensions/IDAPro/Python/6xx/loaders/xmlldr.py

content:
#---------------------------------------------------------------------
# xmlldr.py - IDA XML Import loader and plugin
#---------------------------------------------------------------------
"""
Loader and plugin for IDA to import a XML PROGRAM file to a database.
"""

import idaapi
import time
from xml.etree import cElementTree


XML_IMPORTER_VERSION = "2.1.1"
BASELINE_IDA_VERSION = 620
IDA_SDK_VERSION = idaapi.IDA_SDK_VERSION
BADADDR = idaapi.BADADDR


"""
Loader functions
"""
def accept_file(li, n):
    """
    Check if the file is of supported format

    @param li: a file-like object which can be used to access the input data
    @param n : format number. The function will be called with incrementing
                number until it returns zero
    @return: 0 - no more supported formats
                string "name" - format name to display in the chooser dialog
                dictionary { 'format': "name", 'options': integer }
                options: should be 1, possibly ORed with ACCEPT_FIRST (0x8000)
                to indicate preferred format
    """
    if IDA_SDK_VERSION < BASELINE_IDA_VERSION:
        return 0
    # we support only one format per file
    if n > 0: return 0
    # read 16K bytes to allow for the DTD
    data = li.read(0x4000)
    # look for start of <PROGRAM> element
    start = data.find("<PROGRAM")
    if start >= 0: return "XML PROGRAM file"
    return 0


def load_file(li, neflags, format):
    """
    Load the file into database

    @param li: a file-like object which can be used to access the input data
    @param neflags: options selected by the user, see loader.hpp
    @return: 0-failure, 1-ok
    """
    global event, element
    plugin = False
    status = 0
    st = idaapi.setStat(idaapi.st_Work)
    xml = XmlImporter(False, 0)
    try:
        status = xml.import_xml()
    except Cancelled:
        msg = "XML PROGRAM import cancelled!"
        print "\n" + msg
        idaapi.warning(msg)
    except MultipleAddressSpacesNotSupported:
        msg  = "XML Import cancelled!"
        msg += "\n\nXML Import does not currently support"
        msg += "\nimporting multiple address spaces."
        print "\n" + msg
        idaapi.warning(msg)
    except:
        print "\nHouston, we have a problem!"
        msg = "***** Exception occurred: XML loader failed! *****"
        print "\n" + msg + "\n", sys.exc_type, sys.exc_value
        print event, element.tag, element.attrib
        idaapi.warning(msg)
    finally:
        idaapi.setStat(st)
        xml.cleanup()
        return status


class XmlImporterPlugin(idaapi.plugin_t):
    """
    XML Importer plugin class
    """
    flags = 0
    comment = "Import XML PROGRAM file"
    help = "Import XML <PROGRAM> document to database"
    wanted_name = "XML Importer"
    wanted_hotkey = "Ctrl-Alt-l"

    def init(self):
        """
        init function for XML Importer plugin.
        
        Returns:
            Constant PLUGIN_OK if this IDA version supports the plugin,
            else returns PLUGIN_SKIP if this IDA is older than the supported
            baseline version.
        """
        if IDA_SDK_VERSION < BASELINE_IDA_VERSION:
            idaapi.msg('\nXML Importer plugin (xmlldr.py) not supported ' +
                        'by this version of IDA\n')
            return idaapi.PLUGIN_SKIP
        return idaapi.PLUGIN_OK

    def run(self, arg):
        """
        run function for XML Importer plugin.
        
        Args:
            arg: Integer, a non-zero value enables auto-run feature for
                 IDA batch (no gui) processing mode. Default is 0.
        """
        st = idaapi.setStat(idaapi.st_Work)
        xml = XmlImporter(True, arg)
        try:
            try:
                xml.import_xml()
            except Cancelled:
                msg = "XML Import cancelled!"
                print "\n" + msg
                idaapi.warning(msg)
            except MultipleAddressSpacesNotSupported:
                msg  = "XML Import cancelled!"
                msg += "\n\nXML Import does not currently support"
                msg += "\nimporting multiple address spaces."
                print "\n" + msg
                idaapi.warning(msg)
            except:
                msg = "***** Exception occurred: XML Importer failed! *****"
                print "\n" + msg + "\n", sys.exc_type, sys.exc_value
                idaapi.warning(msg)
        finally:
            xml.cleanup()
            idaapi.setStat(st)


    def term(self):
        pass


def PLUGIN_ENTRY():
    return XmlImporterPlugin()


class Cancelled(Exception):
    pass


class FileError(Exception):
    pass


class MultipleAddressSpacesNotSupported(Exception):
    pass


class XmlImporter:
    """
    XmlImporter class contains methods to import an XML PROGRAM
        document into IDA.
    """
    def __init__(self, plugin, arg):
        """
        Initializes the XmlImporter attributes

        Args:
            arg: Integer, non-zero value enables auto-run feature for
                IDA batch (no gui) processing mode. Default is 0.
        """
        self.plugin = plugin
        self.autorun = False
        self.debug = False # set this to True for debug prints
        self.Elements = {}
        self.Counters = []
        self.Tags = []
        self.Timers = dict()
        self.xmlfile = 0
        self.addr_mode = 1
        self.create = True
        self.dataseg = None
        self.deferred = []
        self.callbacks = {
            'start' : {
                BOOKMARKS           : self.update_import,
                CODE                : self.update_import,
                COMMENTS            : self.update_import,
                COMPILER            : self.import_compiler,
                DATA                : self.update_import,
                DATATYPES           : self.update_import,
                EQUATES             : self.update_import,
                FUNCTIONS           : self.update_import,
                INFO_SOURCE         : self.import_info_source,
                MARKUP              : self.update_import,
                MEMORY_MAP          : self.import_memory_map,
                PROCESSOR           : self.import_processor,
                PROGRAM             : self.import_program,
                PROGRAM_ENTRY_POINTS: self.update_import,
                REGISTER_VALUES     : self.update_import,
                SYMBOL_TABLE        : self.update_import },
            'end' : {
                BOOKMARK            : self.import_bookmark,
                CODE_BLOCK          : self.import_codeblock,
                COMMENT             : self.import_comment,
                DEFINED_DATA        : self.import_defined_data,
                ENUM                : self.import_enum,
                EQUATE_GROUP        : self.import_equate_group,
                EQUATE_REFERENCE    : self.import_equate_reference,
                FUNCTION            : self.import_function,
                FUNCTION_DEF        : self.import_function_def,
                MANUAL_INSTRUCTION  : self.import_manual_instruction,
                MANUAL_OPERAND      : self.import_manual_operand,
                MEMORY_REFERENCE    : self.import_memory_reference,
                MEMORY_SECTION      : self.import_memory_section,
                PROGRAM_ENTRY_POINT : self.import_program_entry_point,
                REGISTER_VALUE_RANGE: self.import_register_value_range,
                STACK_REFERENCE     : self.import_stack_reference,
                STRUCTURE           : self.import_structure,
                SYMBOL              : self.import_symbol,
                TYPE_DEF            : self.import_typedef,
                UNION               : self.import_union,
                # end element for elapse time
                BOOKMARKS           : self.display_timer,
                CODE                : self.display_timer,
                COMMENTS            : self.display_timer,
                DATA                : self.display_timer,
                DATATYPES           : self.process_deferred,
                EQUATES             : self.display_timer,
                FUNCTIONS           : self.display_timer,
                MARKUP              : self.display_timer,
                MEMORY_MAP          : self.display_timer,
                PROGRAM             : self.display_total_time,
                PROGRAM_ENTRY_POINTS: self.display_timer,
                REGISTER_VALUES     : self.display_timer,
                SYMBOL_TABLE        : self.display_timer }
            }


    def import_xml(self):
        """
        Imports the XML PROGRAM file into the database.
        """
        self.display_xml_importer_version()
        displayMenu = self.plugin == True and self.autorun == False 
        #displayMenu = self.autorun == False 
        self.get_options(displayMenu)
        if self.plugin:
            self.filename=idaapi.askfile_c(0, "*.xml",
                                        "Enter name of xml file:")
        else:
            self.filename = idaapi.get_input_file_path()
        if (len(self.filename) == 0):
            return
        idaapi.msg('\nImporting from: ' + self.filename + '\n') 
        idc.Wait()
        if self.plugin == False: idaapi.hide_wait_box()
        idaapi.show_wait_box("Importing XML PROGRAM document....")
        n = 0
        for event,element in cElementTree.iterparse(self.filename,
                                                events=("start","end")):
            if idaapi.wasBreak() == True:
                raise Cancelled
            
            if self.debug == True and event == 'start':
                print element.tag, element.attrib, element.text
            
            if event in self.callbacks:
                if element.tag in self.callbacks[event]:
                    if event == 'start':
                        self.Timers[element.tag] = time.clock()
                    self.callbacks[event][element.tag](element)
                    if event == 'end':
                        element.clear()
            if event == 'end':
                n += 1
        end = time.clock()
        idaapi.hide_wait_box()
        self.display_summary()
        idaapi.msg('\nXML Elements parsed: ' + str(n) + '\n\n')
        return 1
    

    def get_options(self, display):
        """
        Displays the options menu and retrieves the option settings. 
        """
        fmt =  "HELP\n"
        fmt += "XML PROGRAM loader/importer plugin (Python)\n"
        fmt += "IDA SDK: "+ str(IDA_SDK_VERSION) + "\n\n"
        fmt +=  "The XML PROGRAM loader loads elements from a "
        fmt +=  "XML <PROGRAM> document to create an idb database.\n\n"
        fmt +=  "ENDHELP\n"
        fmt +=  "Import from XML PROGRAM document...."
        fmt += "\n <##Options##Code Blocks:{CodeBlocks}>"
        fmt += "\n <Entry Points:{EntryPoints}>"
        fmt += "\n <Segment Register Value Ranges:{RegisterValues}>"
        fmt += "\n <Data Types:{DataTypes}>"
        fmt += "\n <Data Definitions:{DataDefinitions}>"
        fmt += "\n <Symbols:{Symbols}>"
        fmt += "\n <Comments:{Comments}>"
        fmt += "\n <Bookmarks:{Bookmarks}>"
        fmt += "\n <Functions:{Functions}>"
        fmt += "\n <Memory References:{MemoryReferences}>"
        fmt += "\n <Equate/Enum References:{EquateReferences}>"
        fmt += "\n <Manual Instructions/Operands:{Manual}>{cGroup1}>"
        fmt += "\n\n"
    
        Opts = { 'cGroup1': idaapi.Form.ChkGroupControl(
                ( 
                "CodeBlocks", "EntryPoints", "RegisterValues",
                "DataTypes", "DataDefinitions",
                "Symbols", "Comments", "Bookmarks",
                "Functions", "MemoryReferences", "EquateReferences",
                "Manual")) }
        
        self.Options = idaapi.Form(fmt, Opts)
        self.Options.Compile()
    
        self.Options.CodeBlocks.checked       = True
        self.Options.EntryPoints.checked      = True
        self.Options.RegisterValues.checked   = True
        self.Options.DataTypes.checked        = True
        self.Options.DataDefinitions.checked  = True
        self.Options.Symbols.checked          = True
        self.Options.Functions.checked        = True
        self.Options.Comments.checked         = True
        self.Options.Bookmarks.checked        = True
        self.Options.MemoryReferences.checked = True
        self.Options.EquateReferences.checked = True
        self.Options.Manual.checked           = True
        
        if display == True:
            ok = self.Options.Execute()
            if (ok == 0):
                raise Cancelled
            
    def cleanup(self):
        """
        Frees memory and closes message box at termination.
        """
        if self.plugin:
            self.Options.Free()
        idaapi.hide_wait_box()
    

    def display_timer(self, element):
        """
        Displays the elapsed processing time for XML elements.
        
        Args:
            element: XML element object value containing the element tag.
        """
        if element.tag == MEMORY_MAP and self.plugin:
            return
        if element.tag in self.Timers:
            idaapi.msg('elapsed time: %.4f' %
                    (time.clock()-self.Timers[element.tag]))
    

    def display_total_time(self, element):
        """
        Displays the total processing time.
        
        Args:
            element: XML element object value (not used).
        """
        TOTAL = 'Total '
        idaapi.msg('\n%35selapsed time: %.4f' %
                (TOTAL,time.clock()-self.Timers[PROGRAM]))
    

    def display_summary(self):
        """
        Displays summary of the XML PROGRAM import in IDA output window.
        """
        summary = ''
        total = 0
        for tag in self.Tags:
            count = self.Counters[self.Elements[tag]]
            summary += "\n%s: %d" % (tag, count)
            total += count
        summary += '\n--------------------------------------'
        header = '\n--------------------------------------'
        header += ('\nTotal XML Elements Processed: %d') % total
        summary = header + summary
        idaapi.msg(summary)
        if self.plugin and self.autorun == False:
            frmt  = "TITLE XML Import Successful!\n"
            frmt += "ICON INFO\n"
            frmt += "AUTOHIDE NONE\n"
            frmt += "HIDECANCEL\n"
            firstline = "\n XML IMPORT SUCCESSFUL!"
            fileline = '\n\nFile: %s' % self.filename
            details = '\n\nSee output window for details...'
            idaapi.info("%s" % (frmt + firstline + fileline + details))
    

    def display_xml_importer_version(self):
        """
        Displays XML Importer plugin version info in IDA output window.
        """
        if self.plugin:
            f = idaapi.idadir(idaapi.PLG_SUBDIR) + '/xmlldr.py'
        else:
            f = idaapi.idadir(idaapi.LDR_SUBDIR) + '/xmlldr.py'
        plugintime = time.localtime(os.path.getmtime(f))
        ts = time.strftime('%b %d %Y %H:%M:%S', plugintime)
        version = "\nXML Importer Version " + XML_IMPORTER_VERSION
        version += " : SDK " + str(IDA_SDK_VERSION)
        version += " : Python : "+ ts + '\n'
        idaapi.msg(version)
    

    def get_address(self, element, attr):
        """
        Returns the address value for an element.
        
        Args:
            element: XML element object.
            attr: String containing the address attribute name.
            
        Returns:
            Numeric value representing the address.
        """
        addrstr = element.get(attr)
        if '::' in addrstr:
            # overlayed addresses not currently handled
            return BADADDR
        elif ':' in addrstr:
            [segstr, offset_str] = string.split(addrstr,':')
            offset = int(offset_str,16)
            if self.is_int(segstr) == True:
                sgmt = int(segstr,16)
                addr = (sgmt << 4) + offset
            else:
                # multiple address spaces not currently implemented
                addr = BADADDR
            return addr
        else:
            return int(element.get(attr), 16)
    

    def get_attribute(self, element, attr):
        """
        Returns the attribute value string.
        
        Args:
            element: XML element object.
            attr: String containing the attribute name.
            
        Returns:
            String representing the attribute value.
        """
        return element.get(attr)
    

    def get_attribute_value(self, element, attr):
        """
        Returns the numeric attribute value.
        
        Args:
            element: XML element object.
            attr: String containing the attribute name.
            
        Returns:
            Numeric value representing the attribute value.
        """
        val = element.get(attr)
        try:
            if val.upper().startswith('0X') or val.upper().startswith('-0X'):
                return int(val, 16)
            return int(val)
        except:
            idaapi.msg('\nUnable to decode string as value: ' + val)
            return 0
        

    def get_cbsize(self):
        """
        Returns the size of the addressable codebyte for the processor.
        
        Returns:
            Integer representing the number of 8-bit bytes in an
            addressable codebyte.
        """
        return (idaapi.ph_get_cnbits()+7)/8
    

    def get_datatype_flags(self, datatype, size):
        """
        Returns the flags bitmask for the datatype.
        
        Args:
            datatype: String representing the datatype.
            size: Integer representing the datatype size.

        Returns:
            Integer representing the bitmask.
        """
        if datatype.lower().startswith("byte"):     return idaapi.byteflag()
        if datatype.lower().startswith("word"):     return idaapi.wordflag()
        if datatype.lower().startswith("dword"):    return idaapi.dwrdflag()
        if datatype.lower().startswith("qword"):    return idaapi.qwrdflag()
        if datatype.lower().startswith("oword"):    return idaapi.owrdflag()
        if datatype.lower().startswith("tbyte"):    return idaapi.tbytflag()
        if datatype.lower().startswith("float"):    return idaapi.floatflag()
        if datatype.lower().startswith("double"):   return idaapi.doubleflag()
        if datatype.lower().startswith("packed"):   return idaapi.packrealflag()
        if self.is_string_type(datatype):           return idaapi.asciflag()
        if self.is_enumeration(datatype):           return idaapi.enumflag()
        if self.is_structure(datatype):             return idaapi.struflag()
        #if size == 4:                               return idaapi.dwrdflag()
        return 0
    

    def get_string_type(self, datatype):
        if datatype.lower() == 'mbcstring':
            return idaapi.ASCSTR_UNICODE
        if datatype.lower().find('unicode') != -1:
            if datatype.lower().find('pascal') != -1:
                return idaapi.ASCSTR_ULEN2
            return idaapi.ASCSTR_UNICODE
        if datatype.lower().find('pascal') != -1:
            return idaapi.ASCSTR_LEN2
        return idaapi.ASCSTR_TERMCHR
            
    
    def has_attribute(self, element, attr):
        """
        Returns true if the XML element contains the named attribute.
        
        Args:
            element: XML element object
            attr: String containing name of the attribute
        
        Returns:
            True if the element contains the named attribute, otherwise False.
        """
        return attr in element.attrib
    

    def is_enumeration(self, datatype):
        """
        Returns true if datatype is an existing enumeration in the database.
        
        Args:
            datatype: String representing the datatype.
            
        Returns:
            True if the datatype is an enumeration in the database,
            otherwise False.
        """
        if idaapi.get_enum(datatype) == idaapi.BADNODE:  return False
        return True
    
    
    def is_int(self, s):
        try:
            int(s, 16)
            return True
        except:
            return False
        

    def is_pointer_type(self, dtype):
        """
        Returns true if the datatype represents a pointer.
        
        Args:
            dtype: String representing the datatype.
            
        Returns:
            True if the datatype represents a pointer, otherwise False.
        """
        if dtype.lower().startswith("pointer") or dtype.endswith('*'):
            return True
        return False
    

    def is_string_type(self, datatype):
        """
        Returns true if the datatype represents a string type.
        
        Args:
            datatype: String representing the datatype.
            
        Returns:
            True if the datatype represents a string, otherwise False.
        """
        if datatype.lower().startswith("unicode"):  return True
        if datatype.lower().startswith("string"):   return True
        return False
    

    def is_structure(self, datatype):
        """
        Returns true if the datatype represents a structure in the database.
        
        Args:
            dtype: String representing the datatype.
            
        Returns:
            True if the datatype represents an existing structure,
            otherwise False.
        """
        if idaapi.get_struc_id(datatype) == idaapi.BADNODE:  return False
        return True
    

    def import_address_range(self, address_range):
        """
        Processes ADDRESS_RANGE element.
        
        Args:
            address_range: XML element object containing start and end address
                attributes for the address range.
                
        Returns:
            Tuple containing two integers, the start and end address values.
        """
        start = self.get_address(address_range,START)
        end = self.get_address(address_range, END)
        self.update_counter(ADDRESS_RANGE)
        return (start, end)
    

    def import_bit_mask(self, bitmask, eid):
        """
        Processes a BIT_MASK element as an enum bitmask member.
        
        Args:
            bitmask: XML element object representing the IDA enum bitmask.
            eid: Integer representing the IDA enum id
        """
        if IDA_SDK_VERSION < 640:
            return
        name = self.get_attribute(bitmask,NAME)
        value = self.get_attribute_value(bitmask,VALUE)
        idaapi.set_bmask_name(eid, value, name)
        cid = idaapi.get_const_by_name(name)
        self.update_counter(BIT_MASK)
        regcmt = bitmask.find(REGULAR_CMT)
        if regcmt != None:
            idaapi.set_const_cmt(cid, regcmt.text, False);
            self.update_counter(BIT_MASK + ':' + REGULAR_CMT)
        rptcmt = bitmask.find(REPEATABLE_CMT)
        if rptcmt != None:
            idaapi.set_const_cmt(cid, rptcmt.txt, True);
            self.update_counter(BIT_MASK + ':' + REPEATABLE_CMT)
    

    def import_bookmark(self, bookmark):
        """
        Processes a BOOKMARK element.
        
        Args:
            bookmark: XML element object containing bookmark data.
        """
        if self.Options.Bookmarks.checked == False:
            return
        try:
            addr = self.get_address(bookmark, ADDRESS)
            if self.has_attribute(bookmark, TYPE):
                typ = self.get_attribute(bookmark, TYPE)
            category = ''
            if self.has_attribute(bookmark, CATEGORY):
                category = self.get_attribute(bookmark, CATEGORY)
            description = ''
            if self.has_attribute(bookmark, DESCRIPTION):
                description = self.get_attribute(bookmark, DESCRIPTION)
            if idaapi.isEnabled(addr) == False:
                msg = ("import_bookmark: address %X not enabled in database"
                       % addr)
                print msg
                return
            self.update_counter(BOOKMARK)
            for slot in range(1,1025):
                if IDA_SDK_VERSION != 695:
                    ea = idc.GetMarkedPos(slot)
                else:
                    import ida_moves
                    import ida_pro
                    curloc = ida_moves.curloc()
                    intp = ida_pro.int_pointer()
                    intp.assign(slot)
                    ea = curloc.markedpos(intp)
                if ea == BADADDR:
                    curloc = idaapi.curloc()
                    curloc.ea = addr
                    curloc.mark(slot, category, description)
                    break
        except:
            msg = "** Exception occurred in import_bookmark **"
            print "\n" + msg + "\n", sys.exc_type, sys.exc_value
    

    def import_cmts(self, element, sid, typ):
        """
        Processes REGULAR_CMT and REPEATABLE_CMT elements for structures.
        
        Args:
            element: XML element object containing a REGULAR_CMT or
                REPEATABLE_CMT element
            sid: Integer representing the structure id
            typ: String indicating structure type (STRUCTURE or UNION)
        """
        regcmt = element.find(REGULAR_CMT)
        if regcmt != None:
            idaapi.set_struc_cmt(sid, regcmt.text, False)
            self.update_counter(typ + ':' + REGULAR_CMT)
        rptcmt = element.find(REPEATABLE_CMT)
        if rptcmt != None:
            idaapi.set_struc_cmt(sid, rptcmt.text, True)
            self.update_counter(typ + ':' + REPEATABLE_CMT)
            
        
    def import_codeblock(self, code_block):
        """
        Processes a CODE_BLOCK element by disassembling the address range.
        
        Args:
            code_block: XML element containing codeblock start and end
                addresses.
        """
        if self.Options.CodeBlocks.checked == False:
            return
        start = self.get_address(code_block, START)
        end = self.get_address(code_block, END)
        idaapi.do_unknown_range(start, end-start+1, 3)
        addr = start
        while (addr <= end):
            length = idaapi.create_insn(addr)
            addr += idaapi.get_item_size(addr) * self.get_cbsize()
        self.update_counter(CODE_BLOCK)
    

    def import_comment(self, comment):
        """
        Processes a COMMENT element by creating the comment at the address.
        
        Args:
            comment: XML element containing the comment address, type,
                and text.
        """
        if self.Options.Comments.checked == False:
            return
        addr = self.get_address(comment, ADDRESS)
        ctype = self.get_attribute(comment,TYPE)
        text = comment.text
        if ctype == 'pre':
            idaapi.add_long_cmt(addr, True, text)
        elif ctype == 'end-of-line':
            idaapi.set_cmt(addr, text, False)
        elif ctype == 'repeatable':
            idaapi.set_cmt(addr, text, True)
        elif ctype == 'post':
            idaapi.add_long_cmt(addr, False, text)
        self.update_counter(COMMENT+':' + ctype)
    

    def import_compiler(self, compiler):
        """
        Processes the COMPILER element containing the compiler name.
        
        Args:
            compiler: XML element containing the compiler name.
        """
        name = self.get_attribute(compiler, NAME)
        comp = idaapi.COMP_UNK
        if   name == "Visual C++":      comp = idaapi.COMP_MS
        elif name == "Borland C++":     comp = idaapi.COMP_BC
        elif name == "Watcom C++":      comp = idaapi.COMP_WATCOM
        elif name == "GNU C++":         comp = idaapi.COMP_GNU
        elif name == "Visual Age C++":  comp = idaapi.COMP_VISAGE
        elif name == "Delphi":          comp = idaapi.COMP_BP
        idaapi.cvar.inf.cc.id = comp
        self.update_counter(COMPILER)
    

    def import_defined_data(self, defined_data):
        """
        Processes a DEFINED_DATA element by creating a data item at the
            specified address.
        
        Args:
            defined_data: XML element containing the address and
                datatype information for the data item
        """
        if self.Options.DataDefinitions.checked == False:
            return
        addr = self.get_address(defined_data, ADDRESS)
        datatype = self.get_attribute(defined_data, DATATYPE)
        size = self.get_attribute_value(defined_data, SIZE)
        self.update_counter(DEFINED_DATA)        
        ti = idaapi.opinfo_t()
        if self.is_pointer_type(datatype):
            #idaapi.set_refinfo(ti, 0, 0, 0, REF_OFF32)
            flag = idaapi.dwrdflag() | idaapi.FF_0OFF
            #idaapi.set_typeinfo(addr, 0, flag, ti)
        else:
            flag = self.get_datatype_flags(datatype, size)
        if flag == idaapi.asciflag():
            idaapi.make_ascii_string(addr, size,
                                     self.get_string_type(datatype))
        elif flag == idaapi.struflag():
            idaapi.doStruct(addr, size, idaapi.get_struc_id(datatype))
        else:
            idaapi.do_data_ex(addr, flag, size, idaapi.BADNODE)        
        typecmt = defined_data.find(TYPEINFO_CMT)
        if typecmt != None:
            self.update_counter(DEFINED_DATA + ':' + TYPEINFO_CMT)


    def import_enum(self, enum):
        """
        Processes an ENUM element by creating the enumeration.
        
        Args:
            enum: XML element containing the enumeration name and
                member data.
        """
        if self.Options.DataTypes.checked == False:
            return
        name = self.get_attribute(enum, NAME)
        if self.has_attribute(enum,NAMESPACE):
            namespace = self.get_attribute(enum, NAMESPACE)
        if self.has_attribute(enum,SIZE):
            size = self.get_attribute_value(enum, SIZE)
        eid = idaapi.add_enum(BADADDR, name,
                              idaapi.hexflag() | idaapi.dwrdflag())
        self.update_counter(ENUM)
        regcmt = enum.find(REGULAR_CMT)
        if regcmt != None:
            idaapi.set_enum_cmt(eid, regcmt.text, False)
            self.update_counter(ENUM + ':' + REGULAR_CMT)
        rptcmt = enum.find(REPEATABLE_CMT)
        if rptcmt != None:
            idaapi.set_enum_cmt(eid, rptcmt.text, True)
            self.update_counter(ENUM + ':' + REPEATABLE_CMT)
        display_settings = enum.find(DISPLAY_SETTINGS)
        if display_settings != None:
            self.update_counter(ENUM + ':' + DISPLAY_SETTINGS)
        enum_entries = enum.findall(ENUM_ENTRY)
        for enum_entry in enum_entries:
            self.import_enum_entry(enum_entry, eid)
    

    def import_enum_entry(self, enum_entry, eid):
        """
        Processes an ENUM_ENTRY by creating a member in the enumeration.
        
        Args:
            enum_entry: XML element containing the member name and value.
            eid: Integer representing the id of the enumeration.
        """
        name = self.get_attribute(enum_entry, NAME)
        value = self.get_attribute_value(enum_entry, VALUE)
        if IDA_SDK_VERSION > 630:
            idaapi.add_const(eid, name, value)
        else:
            idaapi.add_enum_member(eid, name, value)
        if IDA_SDK_VERSION > 630:
            cid = idaapi.get_const_by_name(name)
        else:
            cid = idaapi.get_enum_member_by_name(name)
        self.update_counter(ENUM_ENTRY)
        regcmt = enum_entry.find(REGULAR_CMT)
        if regcmt != None:
            if IDA_SDK_VERSION > 630:
                idaapi.set_const_cmt(cid, regcmt.text, False);
            else:
                idaapi.set_enum_member_cmt(cid, regcmt.text, False);
            self.update_counter(ENUM_ENTRY + ':' + REGULAR_CMT)
        rptcmt = enum_entry.find(REPEATABLE_CMT)
        if rptcmt != None:
            if IDA_SDK_VERSION > 630:
                idaapi.set_const_cmt(cid, rptcmt.text, True);
            else:
                idaapi.set_enum_member_cmt(cid, rptcmt.text, True);
            self.update_counter(ENUM_ENTRY + ':' + REPEATABLE_CMT)
    

    def import_equate(self, equate, eid):
        """
        Processes EQUATE element as member of an enumeration.
        
        Args:
            enum_entry: XML element containing the equate name and value.
            eid: Integer representing the id for the enumeration.
        """
        name = self.get_attribute(equate,NAME)
        value = self.get_attribute_value(equate,VALUE)
        bm = -1
        if self.has_attribute(equate, BIT_MASK):
            bm = self.get_attribute_value(equate, BIT_MASK)
        if IDA_SDK_VERSION > 630:
            #idaapi.add_const(eid, name, value, bm)
            idaapi.add_const(eid, name, value)
            cid = idaapi.get_const_by_name(name)
        else:
            #idaapi.add_enum_member(eid, name, value, bm)
            idaapi.add_enum_member(eid, name, value)
            cid = idaapi.get_enum_member_by_name(name)
        self.update_counter(EQUATE)
        regcmt = equate.find(REGULAR_CMT) 
        if regcmt != None:
            if IDA_SDK_VERSION > 630:
                idaapi.set_const_cmt(cid, regcmt.text, False);
            else:
                idaapi.set_enum_member_cmt(cid, regcmt.text, False);
            self.update_counter(EQUATE + ':' + REGULAR_CMT)
        rptcmt = equate.find(REPEATABLE_CMT)
        if rptcmt != None:
            if IDA_SDK_VERSION > 630:
                idaapi.set_const_cmt(cid, rptcmt.text, True);
            else:
                idaapi.set_enum_member_cmt(cid, rptcmt.text, True);
            self.update_counter(EQUATE + ':' + REPEATABLE_CMT)
    

    def import_equate_group(self, equate_group):
        """
        Processes EQUATE_GROUP as IDA enumeration type.
        
        Args:
            equate_group: XML element containing the group name and
                equate definitions.
        """
        if self.Options.DataTypes.checked == False:
            return
        msg = EQUATE_GROUP
        name = ''
        if self.has_attribute(equate_group, NAME):
            name = self.get_attribute(equate_group, NAME)
        bf = ''
        if self.has_attribute(equate_group, BIT_FIELD):
            bf = self.get_attribute(equate_group, BIT_FIELD)
        eid = idaapi.add_enum(BADADDR, name, idaapi.hexflag())
        idaapi.set_enum_bf(eid, (bf == 'yes'))
        self.update_counter(EQUATE_GROUP)
        regcmt = equate_group.find(REGULAR_CMT)
        if regcmt != None:
            idaapi.set_enum_cmt(eid, regcmt.text, False)
            self.update_counter(EQUATE_GROUP + ':' + REGULAR_CMT)
        rptcmt = equate_group.find(REPEATABLE_CMT)
        if rptcmt != None:
            idaapi.set_enum_cmt(eid, rptcmt.text, True)
            self.update_counter(EQUATE_GROUP + ':' + REPEATABLE_CMT)
        equates = equate_group.findall(EQUATE)
        for equate in equates:
            self.import_equate(equate,eid)
        bit_masks = equate_group.findall(BIT_MASK)
        for bit_mask in bit_masks:
            self.import_bit_mask(bit_mask, eid)
    

    def import_equate_reference(self, equate_reference):
        if (self.Options.DataTypes.checked == False or
            self.Options.EquateReferences.checked == False):
            return
        self.update_counter(EQUATE_REFERENCE)
        addr = self.get_address(equate_reference, ADDRESS)
        name = ''
        if self.has_attribute(equate_reference, NAME):
            name = self.get_attribute(equate_reference, NAME)
        if name == '':
            return
        opnd = 0
        if self.has_attribute(equate_reference, OPERAND_INDEX):
            opnd = self.get_attribute_value(equate_reference, OPERAND_INDEX)
        value = None
        if self.has_attribute(equate_reference, VALUE):
            value = self.get_attribute_value(equate_reference, VALUE)
        if IDA_SDK_VERSION < 650:
            cid = idaapi.get_enum_member_by_name(name)
        else:
            cid = idaapi.get_const_by_name(name)
        if cid == idaapi.BADNODE:
            return
        if IDA_SDK_VERSION < 650:
            eid = idaapi.get_enum_member_enum(cid)
        else:
            eid = idaapi.get_const_enum(cid)
        if eid == idaapi.BADNODE:
            return
        idaapi.op_enum(addr, opnd, eid, 0)
        

    def import_function(self, function):
        """
        Creates a function using the FUNCTION attributes.
        
        Args:
            function: XML element containing the function address and
                attributes.
        """
        if self.Options.Functions.checked == False:
            return
        try:
            entry_point = self.get_address(function, ENTRY_POINT)
            name = ''
            if self.has_attribute(function, NAME):
                name = self.get_attribute(function, NAME)
            libfunc = 'n'
            if self.has_attribute(function, LIBRARY_FUNCTION):
                libfunc = self.get_attribute(function, LIBRARY_FUNCTION)
            if idaapi.isEnabled(entry_point) == False:
                msg = ("import_function: address %X not enabled in database"
                       % entry_point)
                print msg
                return
            idaapi.add_func(entry_point, BADADDR)
            self.update_counter(FUNCTION)
            func = idaapi.get_func(entry_point)
            if libfunc == 'y':
                func.flags |= idaapi.FUNC_LIB
            ranges = function.findall(ADDRESS_RANGE)
            for addr_range in ranges:
                (start, end) = self.import_address_range(addr_range)
                idaapi.append_func_tail(func, start, end)
                #idaapi.analyze_area(start, end+1)
            idc.Wait()
            regcmt = function.find(REGULAR_CMT)
            if regcmt != None:
                self.update_counter(FUNCTION + ':' + REGULAR_CMT)
                idaapi.set_func_cmt(func, regcmt.text, False)
            rptcmt = function.find(REPEATABLE_CMT)
            if rptcmt != None:
                self.update_counter(FUNCTION + ':' + REPEATABLE_CMT)
                idaapi.set_func_cmt(func, rptcmt.text, True)
            typecmt = function.find(TYPEINFO_CMT)
            if typecmt != None:
                self.update_counter(FUNCTION + ':' + TYPEINFO_CMT)
                idc.SetType(entry_point, typecmt.text + ';')
            sf = function.find(STACK_FRAME)
            if sf != None:
                self.import_stack_frame(sf, func)
            register_vars = function.findall(REGISTER_VAR)
            for register_var in register_vars:
                self.import_register_var(register_var, func)
        except:
            msg = "** Exception occurred in import_function **"
            print "\n" + msg + "\n", sys.exc_type, sys.exc_value


    def import_function_def(self, function_def):
        # import_function_def: NOT IMPLEMENTED
        if self.Options.DataTypes.checked == False:
            return
        self.update_counter(FUNCTION_DEF)
        

    def import_info_source(self, info_source):
        """
        Processes INFO_SOURCE containing information about the
            source of the XML PROGRAM file.
        
        Args:
            info_source: XML element containing attributes that identify
                the source of the PROGRAM data.
        """
        if self.has_attribute(info_source, TOOL):
            tool = self.get_attribute(info_source, TOOL)
        if self.has_attribute(info_source, USER):
            user = self.get_attribute(info_source, USER)
        if self.has_attribute(info_source, FILE):
            f = self.get_attribute(info_source, FILE)
        if self.has_attribute(info_source, TIMESTAMP):
            ts = self.get_attribute(info_source, TIMESTAMP)
        self.update_counter(INFO_SOURCE)
    

    def import_manual_instruction(self, manual_instruction):
        """
        Creates a manual instruction.
        
        Args:
            manual_instruction: XML element containing MANUAL_INSTRUCTION.
        """
        if self.Options.Manual.checked == False:
            return
        addr = self.get_address(manual_instruction, ADDRESS)
        idaapi.set_manual_insn(addr, manual_instruction.text)
        self.update_counter(MANUAL_INSTRUCTION)
    

    def import_manual_operand(self, manual_operand):
        """
        Creates a manual operand at an address.
        
        Args:
            manual_operand: MANUAL_OPERAND XML element.
        """
        if self.Options.Manual.checked == False:
            return
        addr = self.get_address(manual_operand, ADDRESS)
        op = self.get_attribute_value(manual_operand, OPERAND_INDEX)
        if idaapi.isEnabled(addr):
            idaapi.set_forced_operand(addr, op, manual_operand.text)
            self.update_counter(MANUAL_OPERAND)


    def process_deferred(self, element):
        """
        Processes the list of deferred structure members when the
        DATATYPES end element is encountered.
        
        Args:
            element: XML end element for DATATYPES 
        """
        for (member, sptr) in self.deferred:
            self.import_member(member, sptr, False)
        self.display_timer(element)


    def import_member(self, member, sptr, defer=True):
        """
        Creates a member for a structure.
        
        Args:
            member: MEMBER XML element.
            sptr:
            defer:  boolean indicating if processing a member should be
                    deferred when the type is unknown. A member should
                    only be deferred on the first pass, not when processing
                    the deferred list.
        """
        offset = self.get_attribute_value(member, OFFSET)
        datatype = self.get_attribute(member, DATATYPE)
        if self.has_attribute(member, DATATYPE_NAMESPACE):
            dt_namespace = self.get_attribute(member, DATATYPE_NAMESPACE)
        name = ''
        if self.has_attribute(member, NAME):
            name = self.get_attribute(member, NAME)
        size = 0
        if self.has_attribute(member, SIZE):
            size = self.get_attribute_value(member, SIZE)
        ti = idaapi.opinfo_t()
        if self.is_pointer_type(datatype):
            #idaapi.set_refinfo(ti, 0, BADADDR, 0, REF_OFF32)
            flag = idaapi.dwrdflag() | idaapi.FF_0OFF
        else:
            flag = self.get_datatype_flags(datatype, size)
        if flag == 0:
            if defer:
                self.deferred.append([member, sptr])
                return                
        if flag == idaapi.enumflag():
            t = idaapi.get_enum(datatype)
            ti.ec.tid = t
            ti.ec.serial = idaapi.get_enum_idx(t)
        if flag == idaapi.struflag():
            t = idaapi.get_struc_id(datatype)
            ti.tid = t
        error = idaapi.add_struc_member(sptr, name, offset, flag, ti, size)
        mbr = idaapi.get_member(sptr, offset)
        self.import_member_cmts(member, mbr)
        self.update_counter(MEMBER)
        
        
    def import_member_cmts(self, member, mbr):
        """
        Processes REGULAR_CMT and REPEATABLE_CMT elements for members.
        
        Args:
            element: XML element object containing a REGULAR_CMT or
                REPEATABLE_CMT element
            mbr: Integer representing the member id
        """
        regcmt = member.find(REGULAR_CMT)
        if regcmt != None:
            idaapi.set_member_cmt(mbr, regcmt.text, False)
            self.update_counter(MEMBER + ':' + REGULAR_CMT)
        rptcmt = member.find(REPEATABLE_CMT)
        if rptcmt != None:
            idaapi.set_member_cmt(mbr, rptcmt.text, True)
            self.update_counter(MEMBER + ':' + REPEATABLE_CMT)
        

    def import_members(self, element, sptr):
        """
        Add data members to a structure.
        
        Args:
            element: STRUCTURE XML element containing MEMBER sub-elements.
            sptr:
        """
        members = element.findall(MEMBER)
        for member in members:
            self.import_member(member, sptr)
                

    def import_memory_contents(self, memory_contents, start, size):
        """
        Processes MEMORY_CONTENTS to load data for a memory block.
        
        Args:
            memory_contents: MEMORY_CONTENTS XML element.
        """
        if memory_contents.get(START_ADDR) == None:
            saddr = start
        else:
            saddr = self.get_address(memory_contents, START_ADDR)
        fname = self.get_attribute(memory_contents, FILE_NAME)
        offset = self.get_attribute_value(memory_contents, FILE_OFFSET)
        if memory_contents.get(LENGTH) == None:
            length = size
        else:
            length = self.get_attribute_value(memory_contents, LENGTH)
        #(binfilename, ext) = os.path.splitext(self.filename)
        #binfilename += ".bytes"
        (binfilename, fileext) = os.path.split(self.filename)
        binfilename += '/' + fname
        binfile = idaapi.loader_input_t()
        binfile.open(binfilename)
        binfile.file2base(offset,saddr,saddr+length,False)
        binfile.close()
        self.update_counter(MEMORY_CONTENTS)
    

    def import_memory_map(self, memory_map):
        """
        Processes the MEMORY_MAP element.
        
        Args:
            memory_map: MEMORY_MAP XML element.
            
        MEMORY_MAP is only processed by the IDA loader. It is ignored when
            run as an IDA plugin.
        """
        # import memory sections only when run as loader
        if self.plugin:
            return
        self.update_import(memory_map)
    

    def import_memory_reference(self, memory_reference):
        """
        Processes the MEMORY_REFERENCE element.
        
        Args:
            memory_reference: MEMORY_REFERENCE XML element.
        """
        if self.Options.MemoryReferences.checked == False:
            return        
        addr = self.get_address(memory_reference, ADDRESS)
        if self.has_attribute(memory_reference, OPERAND_INDEX):
            op = self.get_attribute_value(memory_reference, OPERAND_INDEX)
        if self.has_attribute(memory_reference, USER_DEFINED):
            user = self.get_attribute(memory_reference, USER_DEFINED)
        to_addr = self.get_address(memory_reference, TO_ADDRESS)
        if self.has_attribute(memory_reference, BASE_ADDRESS):
            base_addr = self.get_address(memory_reference, BASE_ADDRESS)
        if self.has_attribute(memory_reference, PRIMARY):
            primary = self.get_attribute(memory_reference, PRIMARY)
        self.update_counter(MEMORY_REFERENCE)
        # TODO: import_memory_reference - add code to store reference
 
 
    def import_memory_section(self, memory_section):
        """
        Creates a memory segment in the database.
        
        Args:
            memory_section: MEMORY_SECTION XML element.
            
        MEMORY_SECTION is only processed by the IDA loader. It is ignored
            when run as an IDA plugin.
        """
        # TODO: import_memory_section - handle overlays?
        # import memory sections only when run as loader
        if self.plugin:
            return
        name = self.get_attribute(memory_section, NAME)
        length = self.get_attribute_value(memory_section, LENGTH)

        s = idaapi.segment_t()
        addrstr = self.get_attribute(memory_section, START_ADDR)
        seg_str = ''
        if '::' in addrstr:
            # overlay - skip for now
            print '  ** Overlayed memory block %s skipped **  ' % name
            msg  = 'Overlayed memory block %s skipped!' % name
            msg += "\n\nXML Import does not currently support"
            msg += "\noverlayed memory blocks."
            idaapi.warning(msg)
            return
        elif ':' in addrstr:
            [seg_str, offset_str] = string.split(addrstr,':')
            offset = int(offset_str, 16)
            if self.is_int(seg_str):
                base = int(seg_str, 16)
                sel = idaapi.setup_selector(base)
                start = self.get_address(memory_section, START_ADDR)
            else:
                raise MultipleAddressSpacesNotSupported
                return
        else:
            sel = idaapi.allocate_selector(0)
            start = int(addrstr, 16)

        s.sel = sel
        s.startEA = start
        s.endEA = start+length
        s.bitness = self.addr_mode

        perms = ''
        if self.has_attribute(memory_section, PERMISSIONS):
            perms = self.get_attribute(memory_section, PERMISSIONS)
        s.perm = 0
        if 'r' in perms: s.perm |= idaapi.SEGPERM_READ
        if 'w' in perms: s.perm |= idaapi.SEGPERM_WRITE
        if 'x' in perms: s.perm |= idaapi.SEGPERM_EXEC
        idaapi.add_segm_ex(s, name, "", idaapi.ADDSEG_OR_DIE |
                                            idaapi.ADDSEG_QUIET)
        self.update_counter(MEMORY_SECTION)
        for memory_contents in memory_section.findall(MEMORY_CONTENTS):
            self.import_memory_contents(memory_contents, start, length)
    

    def import_processor(self, processor):
        """
        Processes the PROCESSOR element.
        
        Args:
            processor: PROCESSOR XML element.
        """
        name = self.get_attribute(processor, NAME)
        self.update_counter(PROCESSOR)
        if self.plugin:
            return
        address_model = self.get_attribute(processor, ADDRESS_MODEL)
        if address_model != None:
            flag = idaapi.ph_get_flag()
            if str.lower(address_model) == '16-bit':
                self.addr_mode = 0
            elif str.lower(address_model) == '32-bit':
                self.addr_mode = 1
            elif str.lower(address_model) == '64-bit':
                self.addr_mode = 2
        

    def import_program(self, program):
        """
        Processes the PROGRAM element.
        
        Args:
            program: PROGRAM XML element.
        """
        self.update_status(PROGRAM)
        name = self.get_attribute(program, NAME)
        if self.has_attribute(program, EXE_PATH):
            epath = self.get_attribute(program, EXE_PATH)
            idaapi.set_root_filename(epath)
        else:
            idaapi.set_root_filename(name)
        if self.has_attribute(program, EXE_FORMAT):
            eformat = self.get_attribute(program, EXE_FORMAT)
            RootNode = idaapi.netnode('Root Node')
            if IDA_SDK_VERSION < 650:
                RootNode.supset(1, eformat)
            else:
                RootNode.supset(idaapi.RIDX_FILE_FORMAT_NAME, eformat)
        if self.has_attribute(program, IMAGE_BASE):
            base = self.get_attribute_value(program, IMAGE_BASE)
            idaapi.set_imagebase(base)
        if self.has_attribute(program, INPUT_MD5):
            input_md5 = self.get_attribute(program, INPUT_MD5)
            # store original md5 in a special netnode
            md5 = idaapi.netnode(INPUT_MD5, len(INPUT_MD5), True)
            if (IDA_SDK_VERSION < 650):
                md5.supset(1302, input_md5)
            else:
                md5.supset(idaapi.RIDX_MD5, input_md5)
        self.update_counter(PROGRAM)
        """
        # TODO: this needs to be on "end" event for PROGRAM
        description = program.find(DESCRIPTION)
        if description != None:
            pass  # figure out what to do with it
        """

        
    def import_program_entry_point(self, program_entry_point):
        """
        Defines a program entry point.
        
        Args:
            program_entry_point: PROGRAM_ENTRY_POINT XML element.
                Contains the entry point address.
        """
        if self.Options.EntryPoints.checked == False:
            return
        addr = self.get_address(program_entry_point, ADDRESS)
        idaapi.add_entry(addr, addr, "", True)
        self.update_counter(PROGRAM_ENTRY_POINT)
    

    def import_register_value_range(self, register_value_range):
        """
        Defines the address range for a register value.
        
        Args:
            register_value_range: REGISTER_VALUE_RANGE XML element.
                Contains the register, value, start address and range length.
        """
        if self.Options.RegisterValues.checked == False:
            return
        self.update_counter(REGISTER_VALUE_RANGE)
        reg = self.get_attribute(register_value_range, REGISTER)
        if reg == 'cs': return
        value = self.get_attribute_value(register_value_range, VALUE)
        addr = self.get_address(register_value_range, START_ADDRESS)
        length = self.get_attribute_value(register_value_range, LENGTH)
        r = idaapi.str2reg(reg)
        if r >= idaapi.ph.regFirstSreg and r <= idaapi.ph.regLastSreg:
            if IDA_SDK_VERSION < 670:
                idaapi.splitSRarea1(addr, r, value, 2)
            else:
                idaapi.splitSRarea1(addr, r, value, idaapi.SR_user)
    

    def import_register_var(self, register_var, func):
        """
        Defines a register variable for a function.
        
        Args:
            register_var: REGISTER_VAR XML element.
                Contains register, variable name, and datatype.
            func: IDA function object
        """
        name = self.get_attribute(register_var, NAME)
        reg = self.get_attribute(register_var, REGISTER)
        if self.has_attribute(register_var, DATATYPE):
            datatype = self.get_attribute(register_var, DATATYPE)
        if self.has_attribute(register_var, DATATYPE_NAMESPACE):
            namespace = self.get_attribute(register_var, DATATYPE_NAMESPACE)
        idc.MakeLocal(func.startEA, func.endEA, reg, name)
        self.update_counter(REGISTER_VAR)
    

    def import_stack_frame(self, stack_frame, func):
        """
        Defines a stack frame for a function.
        
        Args:
            stack_frame: STACK_FRAME element with STACK_VAR child elements.
        """
        if self.has_attribute(stack_frame, LOCAL_VAR_SIZE):
            lvsize = self.get_attribute_value(stack_frame, LOCAL_VAR_SIZE)
        if self.has_attribute(stack_frame, PARAM_OFFSET):
            param_offset = self.get_attribute_value(stack_frame, PARAM_OFFSET)
        if self.has_attribute(stack_frame, REGISTER_SAVE_SIZE):
            reg_save_size = self.get_attribute_value(stack_frame,
                                                     REGISTER_SAVE_SIZE)
        if self.has_attribute(stack_frame, RETURN_ADDR_SIZE):
            retaddr_size = self.get_attribute_value(stack_frame,
                                                    RETURN_ADDR_SIZE)
        if self.has_attribute(stack_frame, BYTES_PURGED):
            bytes_purged = self.get_attribute_value(stack_frame, BYTES_PURGED)
        self.update_counter(STACK_FRAME)
        for stack_var in stack_frame.findall(STACK_VAR):
            self.import_stack_var(stack_var, func)
    

    def import_stack_reference(self, stack_reference):
        # import_stack_reference: NOT IMPLEMENTED
        #self.update_counter(STACK_REFERENCE)
        pass
                    

    def import_stack_var(self, stack_var, func):
        """
        Processes STACK_VAR element.
        
        Args:
            stack_var: STACK_VAR XML element.
            
        Stack variables are created by IDA's function analysis. 
        Only the STACK_VAR NAME attribute is used to set the name for
        a stack variable at the specified stack/frame offset. 
        """
        spoffset = self.get_attribute_value(stack_var, STACK_PTR_OFFSET)
        datatype = self.get_attribute(stack_var, DATATYPE)
        offset = spoffset + func.frsize + func.frregs
        if self.has_attribute(stack_var, FRAME_PTR_OFFSET):
            fpoffset = self.get_attribute_value(stack_var, FRAME_PTR_OFFSET)
            offset = fpoffset + func.frsize
        name = ''
        if self.has_attribute(stack_var, NAME):
            name = self.get_attribute(stack_var, NAME)
        if self.has_attribute(stack_var, DATATYPE_NAMESPACE):
            namespace = self.get_attribute(stack_var, DATATYPE_NAMESPACE)
        if self.has_attribute(stack_var, SIZE):
            size = self.get_attribute_value(stack_var, SIZE)
        self.update_counter(STACK_VAR)
        sf = idaapi.get_frame(func)
        if name != '':
            idaapi.set_member_name(sf, offset, name)
    

    def import_structure(self, structure):
        """
        Adds a structure.
        
        Args:
            structure: STRUCTURE XML element.
                Contains the STRUCTURE attributes and child elements.
        """
        if self.Options.DataTypes.checked == False:
            return
        name = self.get_attribute(structure, NAME)
        dtyp = idaapi.get_struc_id(name)
        if dtyp != idaapi.BADNODE:
            # duplicate name, try adding name space
            if self.has_attribute(structure, NAMESPACE) == False:
                return
            namespace = self.get_attribute(structure, NAMESPACE)
            name = namspace + '__' + name
            name.replace('/','_')
            name.replace('.','_')
            dtyp = idaapi.get_struc_id(name)
            # skip if still duplicate (could add sequence #)
            if dtyp != idaapi.BADNODE:
                return
        size = 0
        if self.has_attribute(structure, SIZE):
            size = self.get_attribute_value(structure, SIZE)
        if self.has_attribute(structure, VARIABLE_LENGTH):
            vl = self.get_attribute_value(structure, VARIABLE_LENGTH)
            isVariableLength = vl == 'y'
        sid = idaapi.add_struc(BADADDR, name, False)
        sptr = idaapi.get_struc(sid)
        self.update_counter(STRUCTURE)
        self.import_cmts(structure, sid, STRUCTURE)
        self.import_members(structure, sptr)
        t = idaapi.opinfo_t()
        if idaapi.get_struc_size(sptr) < size:
            idaapi.add_struc_member(sptr,"",size-1,idaapi.byteflag(),t,1)
        

    def import_symbol(self, symbol):
        """
        Adds a symbol name at the specified address.
        
        Args:
            symbol: SYMBOL XML element.
                Contains symbol name and address. Optionally includes
                type and mangled symbol.
        """
        if self.Options.Symbols.checked == False:
            return
        addr = self.get_address(symbol, ADDRESS)
        name = self.get_attribute(symbol, NAME)
        if self.has_attribute(symbol, MANGLED):
            name = self.get_attribute(symbol, MANGLED)
        flag = idaapi.SN_NOWARN
        if self.has_attribute(symbol, TYPE):
            typ = self.get_attribute(symbol, TYPE)
            if  typ == 'local': flag |= idaapi.SN_LOCAL
        idaapi.set_name(addr, name, flag)
        self.update_counter(SYMBOL)
    

    def import_typedef(self, type_def):
        # import_typedef: NOT IMPLEMENTED
        if self.Options.DataTypes.checked == False:
            return
        self.update_counter(TYPE_DEF)
        

    def import_union(self, union):
        """
        Adds a union datatype.
        
        Args:
            union: UNION XML element.
                Contains UNION attributes and child elements.
        """
        if self.Options.DataTypes.checked == False:
            return
        name = self.get_attribute(union, NAME)
        dtyp = idaapi.get_struc_id(name)
        if dtyp != idaapi.BADNODE:
            # duplicate name, try adding name space
            if self.has_attribute(union, NAMESPACE) == False:
                return
            namespace = self.get_attribute(union, NAMESPACE)
            name = namspace + '__' + name
            name.replace('/','_')
            name.replace('.','_')
            dtyp = idaapi.get_struc_id(name)
            # skip if still duplicate (could add sequence #)
            if dtyp != idaapi.BADNODE:
                return
        size = 0
        if self.has_attribute(union, SIZE):
            size = self.get_attribute_value(union, SIZE)
        sid = idaapi.add_struc(BADADDR, name, True)
        sptr = idaapi.get_struc(sid)
        self.update_counter(UNION)
        self.import_cmts(union, sid, UNION)
        self.import_members(union, sptr)
        t = idaapi.opinfo_t()
        if idaapi.get_struc_size(sptr) < size:
            idaapi.add_struc_member(sptr,"", size-1, idaapi.byteflag(), None, 1)
            

    def open_file(self, filename, mode):
        """
        Opens filename to specified mode.
        
        Args:
            filename: String representing absolute filepath.
            mode: String representing mode for open.
            
        Returns
            File handle.
        
        Exceptions:
            Displays a warning and raises FileError exception
            if open fails.
        """
        try:
            f = open(filename, mode)
            return f
        except:
            fmt = "TITLE ERROR!\n"
            fmt += "ICON ERROR\n"
            fmt += "AUTOHIDE NONE\n"
            fmt += "HIDECANCEL\n"
            fmt += "Error opening file" + filename + "!\n"
            idaapi.warning(fmt)
            raise FileError
    

    def update_counter(self, tag):
        """
        Updates the counter for the element tag.
        
        Args:
            tag: String representing element tag.
        """
        if tag in self.Elements:
            self.Counters[self.Elements[tag]] += 1
        else:
            self.Elements[tag] = len(self.Elements)
            self.Counters.append(1)
            self.Tags.append(tag)
    

    def update_import(self, element):
        """
        Update the element counter and processing status.
        
        Args:
            element: XML element
            
        This function is used to process certain high-level elements
        (such as COMMENTS, CODE_BLOCKS, SYMBOL_TABLE, FUNCTIONS, etc.)
        that are used to group sub-elements.
        """
        self.update_counter(element.tag)
        self.update_status(element.tag)
    

    def update_status(self, tag):
        """
        Displays the processing status in the IDA window.
        
        Args:
            tag: String representing XML element tag
        """
        status = 'Importing ' + tag
        idaapi.msg('\n%-35s' % status)
        idaapi.hide_wait_box()
        idaapi.show_wait_box(status)
    
        
# Global XML string constants for elements and attributes
ADDRESS = 'ADDRESS'
ADDRESS_MODEL = 'ADDRESS_MODEL'
ADDRESS_RANGE = 'ADDRESS_RANGE'
BASE_ADDRESS = 'BASE_ADDRESS'
BIT_FIELD = 'BIT_FIELD'
BIT_MAPPED = 'BIT_MAPPED'
BIT_MASK = 'BIT_MASK'
BOOKMARK = 'BOOKMARK'
BOOKMARKS = 'BOOKMARKS'
BYTES = 'BYTES'
BYTES_PURGED = 'BYTES_PURGED'
CATEGORY = 'CATEGORY'
CODE = 'CODE'
CODE_BLOCK = 'CODE_BLOCK'
COMMENT = 'COMMENT'
COMMENTS = 'COMMENTS'
COMPILER = 'COMPILER'
DATA = 'DATA'
DATATYPE = 'DATATYPE'
DATATYPES = 'DATATYPES'
DATATYPE_NAMESPACE = 'DATATYPE_NAMESPACE'
DEFINED_DATA = 'DEFINED_DATA'
DESCRIPTION = 'DESCRIPTION'
DISPLAY_SETTINGS = 'DISPLAY_SETTINGS'
END = 'END'
ENDIAN = 'ENDIAN'
ENTRY_POINT = 'ENTRY_POINT'
ENUM = 'ENUM'
ENUM_ENTRY = 'ENUM_ENTRY'
EQUATE = 'EQUATE'
EQUATES = 'EQUATES'
EQUATE_GROUP = 'EQUATE_GROUP'
EQUATE_REFERENCE = 'EQUATE_REFERENCE'
EXE_FORMAT = 'EXE_FORMAT'
EXE_PATH = 'EXE_PATH'
EXT_LIBRARY = 'EXT_LIBRARY'
EXT_LIBRARY_REFERENCE = 'EXT_LIBRARY_REFERENCE'
EXT_LIBRARY_TABLE = 'EXT_LIBRARY_TABLE'
FAMILY = 'FAMILY'
FILE = 'FILE'
FILE_NAME = 'FILE_NAME'
FILE_OFFSET = 'FILE_OFFSET'
FOLDER = 'FOLDER'
FORMAT = 'FORMAT'
FRAGMENT = 'FRAGMENT'
FRAME_PTR_OFFSET = 'FRAME_PTR_OFFSET'
FUNCTION = 'FUNCTION'
FUNCTIONS = 'FUNCTIONS'
FUNCTION_DEF = 'FUNCTION_DEF'
IMAGE_BASE = 'IMAGE_BASE'
INPUT_MD5 = 'INPUT_MD5'
INFO_SOURCE = 'INFO_SOURCE'
LANGUAGE_PROVIDER = 'LANGUAGE_PROVIDER'
LENGTH = 'LENGTH'
LIB_ADDR = 'LIB_ADDR'
LIB_LABEL = 'LIB_LABEL'
LIB_ORDINAL = 'LIB_ORDINAL'
LIB_PROG_NAME = 'LIB_PROG_NAME'
LIBRARY_FUNCTION = 'LIBRARY_FUNCTION'
LOCAL_VAR_SIZE = 'LOCAL_VAR_SIZE'
MANGLED = 'MANGLED'
MANUAL_INSTRUCTION = 'MANUAL_INSTRUCTION'
MANUAL_OPERAND = 'MANUAL_OPERAND'
MARKUP = 'MARKUP'
MEMBER = 'MEMBER'
MEMORY_CONTENTS = 'MEMORY_CONTENTS'
MEMORY_MAP = 'MEMORY_MAP'
MEMORY_REFERENCE = 'MEMORY_REFERENCE'
MEMORY_SECTION = 'MEMORY_SECTION'
NAME = 'NAME'
NAMESPACE = 'NAMESPACE'
OFFSET = 'OFFSET'
OPERAND_INDEX = 'OPERAND_INDEX'
PARAM_OFFSET = 'PARAM_OFFSET'
PATH = 'PATH'
PERMISSIONS = 'PERMISSIONS'
PRIMARY = 'PRIMARY'
PROCESSOR = 'PROCESSOR'
PROGRAM = 'PROGRAM'
PROGRAM_ENTRY_POINT = 'PROGRAM_ENTRY_POINT'
PROGRAM_ENTRY_POINTS = 'PROGRAM_ENTRY_POINTS'
PROGRAM_TREES = 'PROGRAM_TREES'
PROPERTIES = 'PROPERTIES'
PROPERTY = 'PROPERTY'
REGISTER = 'REGISTER'
REGISTER_SAVE_SIZE = 'REGISTER_SAVE_SIZE'
REGISTER_VALUES = 'REGISTER_VALUES'
REGISTER_VALUE_RANGE = 'REGISTER_VALUE_RANGE'
REGISTER_VAR = 'REGISTER_VAR'
REGULAR_CMT = 'REGULAR_CMT'
RELOCATION = 'RELOCATION'
RELOCATION_TABLE = 'RELOCATION_TABLE'
REPEATABLE_CMT = 'REPEATABLE_CMT'
RETURN_ADDR_SIZE = 'RETURN_ADDR_SIZE'
RETURN_TYPE = 'RETURN_TYPE'
SHOW_TERMINATOR = 'SHOW_TERMINATOR'
SIGNED = 'SIGNED'
SIZE = 'SIZE'
SOURCE_ADDRESS = 'SOURCE_ADDRESS'
SOURCE_TYPE = 'SOURCE_TYPE'
STACK_FRAME = 'STACK_FRAME'
STACK_PTR_OFFSET = 'STACK_PTR_OFFSET'
STACK_REFERENCE = 'STACK_REFERENCE'
STACK_VAR = 'STACK_VAR'
START = 'START'
START_ADDR = 'START_ADDR'
START_ADDRESS = 'START_ADDRESS'
STRUCTURE = 'STRUCTURE'
SYMBOL = 'SYMBOL'
SYMBOL_TABLE = 'SYMBOL_TABLE'
TIMESTAMP = 'TIMESTAMP'
TOOL = 'TOOL'
TO_ADDRESS = 'TO_ADDRESS'
TREE = 'TREE'
TYPE = 'TYPE'
TYPEINFO_CMT = 'TYPEINFO_CMT'
TYPE_DEF = 'TYPE_DEF'
UNION = 'UNION'
USER = 'USER'
USER_DEFINED = 'USER_DEFINED'
VALUE = 'VALUE'
VARIABLE_LENGTH = 'VARIABLE_LENGTH'
ZERO_PAD = 'ZERO_PAD'


repository: DamianFekete/NSA-Ghidra
file: Extensions/IDAPro/Python/6xx/loaders/xmlldr.py
wanted_name: XML Importer
comment: Import XML PROGRAM file
created_at: 2019-03-06 09:02:58+00:00
pushed_at: 2019-03-06 09:12:12+00:00
url:
https://github.com/DamianFekete/NSA-Ghidra/blob/f286cb8cddb91ca6bf75a10ce43890fae319232f/Extensions/IDAPro/Python/6xx/loaders/xmlldr.py

content:
#---------------------------------------------------------------------
# xmlldr.py - IDA XML Import loader and plugin
#---------------------------------------------------------------------
"""
Loader and plugin for IDA to import a XML PROGRAM file to a database.
"""

import idaapi
import time
from xml.etree import cElementTree


XML_IMPORTER_VERSION = "2.1.1"
BASELINE_IDA_VERSION = 620
IDA_SDK_VERSION = idaapi.IDA_SDK_VERSION
BADADDR = idaapi.BADADDR


"""
Loader functions
"""
def accept_file(li, n):
    """
    Check if the file is of supported format

    @param li: a file-like object which can be used to access the input data
    @param n : format number. The function will be called with incrementing
                number until it returns zero
    @return: 0 - no more supported formats
                string "name" - format name to display in the chooser dialog
                dictionary { 'format': "name", 'options': integer }
                options: should be 1, possibly ORed with ACCEPT_FIRST (0x8000)
                to indicate preferred format
    """
    if IDA_SDK_VERSION < BASELINE_IDA_VERSION:
        return 0
    # we support only one format per file
    if n > 0: return 0
    # read 16K bytes to allow for the DTD
    data = li.read(0x4000)
    # look for start of <PROGRAM> element
    start = data.find("<PROGRAM")
    if start >= 0: return "XML PROGRAM file"
    return 0


def load_file(li, neflags, format):
    """
    Load the file into database

    @param li: a file-like object which can be used to access the input data
    @param neflags: options selected by the user, see loader.hpp
    @return: 0-failure, 1-ok
    """
    global event, element
    plugin = False
    status = 0
    st = idaapi.setStat(idaapi.st_Work)
    xml = XmlImporter(False, 0)
    try:
        status = xml.import_xml()
    except Cancelled:
        msg = "XML PROGRAM import cancelled!"
        print "\n" + msg
        idaapi.warning(msg)
    except MultipleAddressSpacesNotSupported:
        msg  = "XML Import cancelled!"
        msg += "\n\nXML Import does not currently support"
        msg += "\nimporting multiple address spaces."
        print "\n" + msg
        idaapi.warning(msg)
    except:
        print "\nHouston, we have a problem!"
        msg = "***** Exception occurred: XML loader failed! *****"
        print "\n" + msg + "\n", sys.exc_type, sys.exc_value
        print event, element.tag, element.attrib
        idaapi.warning(msg)
    finally:
        idaapi.setStat(st)
        xml.cleanup()
        return status


class XmlImporterPlugin(idaapi.plugin_t):
    """
    XML Importer plugin class
    """
    flags = 0
    comment = "Import XML PROGRAM file"
    help = "Import XML <PROGRAM> document to database"
    wanted_name = "XML Importer"
    wanted_hotkey = "Ctrl-Alt-l"

    def init(self):
        """
        init function for XML Importer plugin.
        
        Returns:
            Constant PLUGIN_OK if this IDA version supports the plugin,
            else returns PLUGIN_SKIP if this IDA is older than the supported
            baseline version.
        """
        if IDA_SDK_VERSION < BASELINE_IDA_VERSION:
            idaapi.msg('\nXML Importer plugin (xmlldr.py) not supported ' +
                        'by this version of IDA\n')
            return idaapi.PLUGIN_SKIP
        return idaapi.PLUGIN_OK

    def run(self, arg):
        """
        run function for XML Importer plugin.
        
        Args:
            arg: Integer, a non-zero value enables auto-run feature for
                 IDA batch (no gui) processing mode. Default is 0.
        """
        st = idaapi.setStat(idaapi.st_Work)
        xml = XmlImporter(True, arg)
        try:
            try:
                xml.import_xml()
            except Cancelled:
                msg = "XML Import cancelled!"
                print "\n" + msg
                idaapi.warning(msg)
            except MultipleAddressSpacesNotSupported:
                msg  = "XML Import cancelled!"
                msg += "\n\nXML Import does not currently support"
                msg += "\nimporting multiple address spaces."
                print "\n" + msg
                idaapi.warning(msg)
            except:
                msg = "***** Exception occurred: XML Importer failed! *****"
                print "\n" + msg + "\n", sys.exc_type, sys.exc_value
                idaapi.warning(msg)
        finally:
            xml.cleanup()
            idaapi.setStat(st)


    def term(self):
        pass


def PLUGIN_ENTRY():
    return XmlImporterPlugin()


class Cancelled(Exception):
    pass


class FileError(Exception):
    pass


class MultipleAddressSpacesNotSupported(Exception):
    pass


class XmlImporter:
    """
    XmlImporter class contains methods to import an XML PROGRAM
        document into IDA.
    """
    def __init__(self, plugin, arg):
        """
        Initializes the XmlImporter attributes

        Args:
            arg: Integer, non-zero value enables auto-run feature for
                IDA batch (no gui) processing mode. Default is 0.
        """
        self.plugin = plugin
        self.autorun = False
        self.debug = False # set this to True for debug prints
        self.Elements = {}
        self.Counters = []
        self.Tags = []
        self.Timers = dict()
        self.xmlfile = 0
        self.addr_mode = 1
        self.create = True
        self.dataseg = None
        self.deferred = []
        self.callbacks = {
            'start' : {
                BOOKMARKS           : self.update_import,
                CODE                : self.update_import,
                COMMENTS            : self.update_import,
                COMPILER            : self.import_compiler,
                DATA                : self.update_import,
                DATATYPES           : self.update_import,
                EQUATES             : self.update_import,
                FUNCTIONS           : self.update_import,
                INFO_SOURCE         : self.import_info_source,
                MARKUP              : self.update_import,
                MEMORY_MAP          : self.import_memory_map,
                PROCESSOR           : self.import_processor,
                PROGRAM             : self.import_program,
                PROGRAM_ENTRY_POINTS: self.update_import,
                REGISTER_VALUES     : self.update_import,
                SYMBOL_TABLE        : self.update_import },
            'end' : {
                BOOKMARK            : self.import_bookmark,
                CODE_BLOCK          : self.import_codeblock,
                COMMENT             : self.import_comment,
                DEFINED_DATA        : self.import_defined_data,
                ENUM                : self.import_enum,
                EQUATE_GROUP        : self.import_equate_group,
                EQUATE_REFERENCE    : self.import_equate_reference,
                FUNCTION            : self.import_function,
                FUNCTION_DEF        : self.import_function_def,
                MANUAL_INSTRUCTION  : self.import_manual_instruction,
                MANUAL_OPERAND      : self.import_manual_operand,
                MEMORY_REFERENCE    : self.import_memory_reference,
                MEMORY_SECTION      : self.import_memory_section,
                PROGRAM_ENTRY_POINT : self.import_program_entry_point,
                REGISTER_VALUE_RANGE: self.import_register_value_range,
                STACK_REFERENCE     : self.import_stack_reference,
                STRUCTURE           : self.import_structure,
                SYMBOL              : self.import_symbol,
                TYPE_DEF            : self.import_typedef,
                UNION               : self.import_union,
                # end element for elapse time
                BOOKMARKS           : self.display_timer,
                CODE                : self.display_timer,
                COMMENTS            : self.display_timer,
                DATA                : self.display_timer,
                DATATYPES           : self.process_deferred,
                EQUATES             : self.display_timer,
                FUNCTIONS           : self.display_timer,
                MARKUP              : self.display_timer,
                MEMORY_MAP          : self.display_timer,
                PROGRAM             : self.display_total_time,
                PROGRAM_ENTRY_POINTS: self.display_timer,
                REGISTER_VALUES     : self.display_timer,
                SYMBOL_TABLE        : self.display_timer }
            }


    def import_xml(self):
        """
        Imports the XML PROGRAM file into the database.
        """
        self.display_xml_importer_version()
        displayMenu = self.plugin == True and self.autorun == False 
        #displayMenu = self.autorun == False 
        self.get_options(displayMenu)
        if self.plugin:
            self.filename=idaapi.askfile_c(0, "*.xml",
                                        "Enter name of xml file:")
        else:
            self.filename = idaapi.get_input_file_path()
        if (len(self.filename) == 0):
            return
        idaapi.msg('\nImporting from: ' + self.filename + '\n') 
        idc.Wait()
        if self.plugin == False: idaapi.hide_wait_box()
        idaapi.show_wait_box("Importing XML PROGRAM document....")
        n = 0
        for event,element in cElementTree.iterparse(self.filename,
                                                events=("start","end")):
            if idaapi.wasBreak() == True:
                raise Cancelled
            
            if self.debug == True and event == 'start':
                print element.tag, element.attrib, element.text
            
            if event in self.callbacks:
                if element.tag in self.callbacks[event]:
                    if event == 'start':
                        self.Timers[element.tag] = time.clock()
                    self.callbacks[event][element.tag](element)
                    if event == 'end':
                        element.clear()
            if event == 'end':
                n += 1
        end = time.clock()
        idaapi.hide_wait_box()
        self.display_summary()
        idaapi.msg('\nXML Elements parsed: ' + str(n) + '\n\n')
        return 1
    

    def get_options(self, display):
        """
        Displays the options menu and retrieves the option settings. 
        """
        fmt =  "HELP\n"
        fmt += "XML PROGRAM loader/importer plugin (Python)\n"
        fmt += "IDA SDK: "+ str(IDA_SDK_VERSION) + "\n\n"
        fmt +=  "The XML PROGRAM loader loads elements from a "
        fmt +=  "XML <PROGRAM> document to create an idb database.\n\n"
        fmt +=  "ENDHELP\n"
        fmt +=  "Import from XML PROGRAM document...."
        fmt += "\n <##Options##Code Blocks:{CodeBlocks}>"
        fmt += "\n <Entry Points:{EntryPoints}>"
        fmt += "\n <Segment Register Value Ranges:{RegisterValues}>"
        fmt += "\n <Data Types:{DataTypes}>"
        fmt += "\n <Data Definitions:{DataDefinitions}>"
        fmt += "\n <Symbols:{Symbols}>"
        fmt += "\n <Comments:{Comments}>"
        fmt += "\n <Bookmarks:{Bookmarks}>"
        fmt += "\n <Functions:{Functions}>"
        fmt += "\n <Memory References:{MemoryReferences}>"
        fmt += "\n <Equate/Enum References:{EquateReferences}>"
        fmt += "\n <Manual Instructions/Operands:{Manual}>{cGroup1}>"
        fmt += "\n\n"
    
        Opts = { 'cGroup1': idaapi.Form.ChkGroupControl(
                ( 
                "CodeBlocks", "EntryPoints", "RegisterValues",
                "DataTypes", "DataDefinitions",
                "Symbols", "Comments", "Bookmarks",
                "Functions", "MemoryReferences", "EquateReferences",
                "Manual")) }
        
        self.Options = idaapi.Form(fmt, Opts)
        self.Options.Compile()
    
        self.Options.CodeBlocks.checked       = True
        self.Options.EntryPoints.checked      = True
        self.Options.RegisterValues.checked   = True
        self.Options.DataTypes.checked        = True
        self.Options.DataDefinitions.checked  = True
        self.Options.Symbols.checked          = True
        self.Options.Functions.checked        = True
        self.Options.Comments.checked         = True
        self.Options.Bookmarks.checked        = True
        self.Options.MemoryReferences.checked = True
        self.Options.EquateReferences.checked = True
        self.Options.Manual.checked           = True
        
        if display == True:
            ok = self.Options.Execute()
            if (ok == 0):
                raise Cancelled
            
    def cleanup(self):
        """
        Frees memory and closes message box at termination.
        """
        if self.plugin:
            self.Options.Free()
        idaapi.hide_wait_box()
    

    def display_timer(self, element):
        """
        Displays the elapsed processing time for XML elements.
        
        Args:
            element: XML element object value containing the element tag.
        """
        if element.tag == MEMORY_MAP and self.plugin:
            return
        if element.tag in self.Timers:
            idaapi.msg('elapsed time: %.4f' %
                    (time.clock()-self.Timers[element.tag]))
    

    def display_total_time(self, element):
        """
        Displays the total processing time.
        
        Args:
            element: XML element object value (not used).
        """
        TOTAL = 'Total '
        idaapi.msg('\n%35selapsed time: %.4f' %
                (TOTAL,time.clock()-self.Timers[PROGRAM]))
    

    def display_summary(self):
        """
        Displays summary of the XML PROGRAM import in IDA output window.
        """
        summary = ''
        total = 0
        for tag in self.Tags:
            count = self.Counters[self.Elements[tag]]
            summary += "\n%s: %d" % (tag, count)
            total += count
        summary += '\n--------------------------------------'
        header = '\n--------------------------------------'
        header += ('\nTotal XML Elements Processed: %d') % total
        summary = header + summary
        idaapi.msg(summary)
        if self.plugin and self.autorun == False:
            frmt  = "TITLE XML Import Successful!\n"
            frmt += "ICON INFO\n"
            frmt += "AUTOHIDE NONE\n"
            frmt += "HIDECANCEL\n"
            firstline = "\n XML IMPORT SUCCESSFUL!"
            fileline = '\n\nFile: %s' % self.filename
            details = '\n\nSee output window for details...'
            idaapi.info("%s" % (frmt + firstline + fileline + details))
    

    def display_xml_importer_version(self):
        """
        Displays XML Importer plugin version info in IDA output window.
        """
        if self.plugin:
            f = idaapi.idadir(idaapi.PLG_SUBDIR) + '/xmlldr.py'
        else:
            f = idaapi.idadir(idaapi.LDR_SUBDIR) + '/xmlldr.py'
        plugintime = time.localtime(os.path.getmtime(f))
        ts = time.strftime('%b %d %Y %H:%M:%S', plugintime)
        version = "\nXML Importer Version " + XML_IMPORTER_VERSION
        version += " : SDK " + str(IDA_SDK_VERSION)
        version += " : Python : "+ ts + '\n'
        idaapi.msg(version)
    

    def get_address(self, element, attr):
        """
        Returns the address value for an element.
        
        Args:
            element: XML element object.
            attr: String containing the address attribute name.
            
        Returns:
            Numeric value representing the address.
        """
        addrstr = element.get(attr)
        if '::' in addrstr:
            # overlayed addresses not currently handled
            return BADADDR
        elif ':' in addrstr:
            [segstr, offset_str] = string.split(addrstr,':')
            offset = int(offset_str,16)
            if self.is_int(segstr) == True:
                sgmt = int(segstr,16)
                addr = (sgmt << 4) + offset
            else:
                # multiple address spaces not currently implemented
                addr = BADADDR
            return addr
        else:
            return int(element.get(attr), 16)
    

    def get_attribute(self, element, attr):
        """
        Returns the attribute value string.
        
        Args:
            element: XML element object.
            attr: String containing the attribute name.
            
        Returns:
            String representing the attribute value.
        """
        return element.get(attr)
    

    def get_attribute_value(self, element, attr):
        """
        Returns the numeric attribute value.
        
        Args:
            element: XML element object.
            attr: String containing the attribute name.
            
        Returns:
            Numeric value representing the attribute value.
        """
        val = element.get(attr)
        try:
            if val.upper().startswith('0X') or val.upper().startswith('-0X'):
                return int(val, 16)
            return int(val)
        except:
            idaapi.msg('\nUnable to decode string as value: ' + val)
            return 0
        

    def get_cbsize(self):
        """
        Returns the size of the addressable codebyte for the processor.
        
        Returns:
            Integer representing the number of 8-bit bytes in an
            addressable codebyte.
        """
        return (idaapi.ph_get_cnbits()+7)/8
    

    def get_datatype_flags(self, datatype, size):
        """
        Returns the flags bitmask for the datatype.
        
        Args:
            datatype: String representing the datatype.
            size: Integer representing the datatype size.

        Returns:
            Integer representing the bitmask.
        """
        if datatype.lower().startswith("byte"):     return idaapi.byteflag()
        if datatype.lower().startswith("word"):     return idaapi.wordflag()
        if datatype.lower().startswith("dword"):    return idaapi.dwrdflag()
        if datatype.lower().startswith("qword"):    return idaapi.qwrdflag()
        if datatype.lower().startswith("oword"):    return idaapi.owrdflag()
        if datatype.lower().startswith("tbyte"):    return idaapi.tbytflag()
        if datatype.lower().startswith("float"):    return idaapi.floatflag()
        if datatype.lower().startswith("double"):   return idaapi.doubleflag()
        if datatype.lower().startswith("packed"):   return idaapi.packrealflag()
        if self.is_string_type(datatype):           return idaapi.asciflag()
        if self.is_enumeration(datatype):           return idaapi.enumflag()
        if self.is_structure(datatype):             return idaapi.struflag()
        #if size == 4:                               return idaapi.dwrdflag()
        return 0
    

    def get_string_type(self, datatype):
        if datatype.lower() == 'mbcstring':
            return idaapi.ASCSTR_UNICODE
        if datatype.lower().find('unicode') != -1:
            if datatype.lower().find('pascal') != -1:
                return idaapi.ASCSTR_ULEN2
            return idaapi.ASCSTR_UNICODE
        if datatype.lower().find('pascal') != -1:
            return idaapi.ASCSTR_LEN2
        return idaapi.ASCSTR_TERMCHR
            
    
    def has_attribute(self, element, attr):
        """
        Returns true if the XML element contains the named attribute.
        
        Args:
            element: XML element object
            attr: String containing name of the attribute
        
        Returns:
            True if the element contains the named attribute, otherwise False.
        """
        return attr in element.attrib
    

    def is_enumeration(self, datatype):
        """
        Returns true if datatype is an existing enumeration in the database.
        
        Args:
            datatype: String representing the datatype.
            
        Returns:
            True if the datatype is an enumeration in the database,
            otherwise False.
        """
        if idaapi.get_enum(datatype) == idaapi.BADNODE:  return False
        return True
    
    
    def is_int(self, s):
        try:
            int(s, 16)
            return True
        except:
            return False
        

    def is_pointer_type(self, dtype):
        """
        Returns true if the datatype represents a pointer.
        
        Args:
            dtype: String representing the datatype.
            
        Returns:
            True if the datatype represents a pointer, otherwise False.
        """
        if dtype.lower().startswith("pointer") or dtype.endswith('*'):
            return True
        return False
    

    def is_string_type(self, datatype):
        """
        Returns true if the datatype represents a string type.
        
        Args:
            datatype: String representing the datatype.
            
        Returns:
            True if the datatype represents a string, otherwise False.
        """
        if datatype.lower().startswith("unicode"):  return True
        if datatype.lower().startswith("string"):   return True
        return False
    

    def is_structure(self, datatype):
        """
        Returns true if the datatype represents a structure in the database.
        
        Args:
            dtype: String representing the datatype.
            
        Returns:
            True if the datatype represents an existing structure,
            otherwise False.
        """
        if idaapi.get_struc_id(datatype) == idaapi.BADNODE:  return False
        return True
    

    def import_address_range(self, address_range):
        """
        Processes ADDRESS_RANGE element.
        
        Args:
            address_range: XML element object containing start and end address
                attributes for the address range.
                
        Returns:
            Tuple containing two integers, the start and end address values.
        """
        start = self.get_address(address_range,START)
        end = self.get_address(address_range, END)
        self.update_counter(ADDRESS_RANGE)
        return (start, end)
    

    def import_bit_mask(self, bitmask, eid):
        """
        Processes a BIT_MASK element as an enum bitmask member.
        
        Args:
            bitmask: XML element object representing the IDA enum bitmask.
            eid: Integer representing the IDA enum id
        """
        if IDA_SDK_VERSION < 640:
            return
        name = self.get_attribute(bitmask,NAME)
        value = self.get_attribute_value(bitmask,VALUE)
        idaapi.set_bmask_name(eid, value, name)
        cid = idaapi.get_const_by_name(name)
        self.update_counter(BIT_MASK)
        regcmt = bitmask.find(REGULAR_CMT)
        if regcmt != None:
            idaapi.set_const_cmt(cid, regcmt.text, False);
            self.update_counter(BIT_MASK + ':' + REGULAR_CMT)
        rptcmt = bitmask.find(REPEATABLE_CMT)
        if rptcmt != None:
            idaapi.set_const_cmt(cid, rptcmt.txt, True);
            self.update_counter(BIT_MASK + ':' + REPEATABLE_CMT)
    

    def import_bookmark(self, bookmark):
        """
        Processes a BOOKMARK element.
        
        Args:
            bookmark: XML element object containing bookmark data.
        """
        if self.Options.Bookmarks.checked == False:
            return
        try:
            addr = self.get_address(bookmark, ADDRESS)
            if self.has_attribute(bookmark, TYPE):
                typ = self.get_attribute(bookmark, TYPE)
            category = ''
            if self.has_attribute(bookmark, CATEGORY):
                category = self.get_attribute(bookmark, CATEGORY)
            description = ''
            if self.has_attribute(bookmark, DESCRIPTION):
                description = self.get_attribute(bookmark, DESCRIPTION)
            if idaapi.isEnabled(addr) == False:
                msg = ("import_bookmark: address %X not enabled in database"
                       % addr)
                print msg
                return
            self.update_counter(BOOKMARK)
            for slot in range(1,1025):
                if IDA_SDK_VERSION != 695:
                    ea = idc.GetMarkedPos(slot)
                else:
                    import ida_moves
                    import ida_pro
                    curloc = ida_moves.curloc()
                    intp = ida_pro.int_pointer()
                    intp.assign(slot)
                    ea = curloc.markedpos(intp)
                if ea == BADADDR:
                    curloc = idaapi.curloc()
                    curloc.ea = addr
                    curloc.mark(slot, category, description)
                    break
        except:
            msg = "** Exception occurred in import_bookmark **"
            print "\n" + msg + "\n", sys.exc_type, sys.exc_value
    

    def import_cmts(self, element, sid, typ):
        """
        Processes REGULAR_CMT and REPEATABLE_CMT elements for structures.
        
        Args:
            element: XML element object containing a REGULAR_CMT or
                REPEATABLE_CMT element
            sid: Integer representing the structure id
            typ: String indicating structure type (STRUCTURE or UNION)
        """
        regcmt = element.find(REGULAR_CMT)
        if regcmt != None:
            idaapi.set_struc_cmt(sid, regcmt.text, False)
            self.update_counter(typ + ':' + REGULAR_CMT)
        rptcmt = element.find(REPEATABLE_CMT)
        if rptcmt != None:
            idaapi.set_struc_cmt(sid, rptcmt.text, True)
            self.update_counter(typ + ':' + REPEATABLE_CMT)
            
        
    def import_codeblock(self, code_block):
        """
        Processes a CODE_BLOCK element by disassembling the address range.
        
        Args:
            code_block: XML element containing codeblock start and end
                addresses.
        """
        if self.Options.CodeBlocks.checked == False:
            return
        start = self.get_address(code_block, START)
        end = self.get_address(code_block, END)
        idaapi.do_unknown_range(start, end-start+1, 3)
        addr = start
        while (addr <= end):
            length = idaapi.create_insn(addr)
            addr += idaapi.get_item_size(addr) * self.get_cbsize()
        self.update_counter(CODE_BLOCK)
    

    def import_comment(self, comment):
        """
        Processes a COMMENT element by creating the comment at the address.
        
        Args:
            comment: XML element containing the comment address, type,
                and text.
        """
        if self.Options.Comments.checked == False:
            return
        addr = self.get_address(comment, ADDRESS)
        ctype = self.get_attribute(comment,TYPE)
        text = comment.text
        if ctype == 'pre':
            idaapi.add_long_cmt(addr, True, text)
        elif ctype == 'end-of-line':
            idaapi.set_cmt(addr, text, False)
        elif ctype == 'repeatable':
            idaapi.set_cmt(addr, text, True)
        elif ctype == 'post':
            idaapi.add_long_cmt(addr, False, text)
        self.update_counter(COMMENT+':' + ctype)
    

    def import_compiler(self, compiler):
        """
        Processes the COMPILER element containing the compiler name.
        
        Args:
            compiler: XML element containing the compiler name.
        """
        name = self.get_attribute(compiler, NAME)
        comp = idaapi.COMP_UNK
        if   name == "Visual C++":      comp = idaapi.COMP_MS
        elif name == "Borland C++":     comp = idaapi.COMP_BC
        elif name == "Watcom C++":      comp = idaapi.COMP_WATCOM
        elif name == "GNU C++":         comp = idaapi.COMP_GNU
        elif name == "Visual Age C++":  comp = idaapi.COMP_VISAGE
        elif name == "Delphi":          comp = idaapi.COMP_BP
        idaapi.cvar.inf.cc.id = comp
        self.update_counter(COMPILER)
    

    def import_defined_data(self, defined_data):
        """
        Processes a DEFINED_DATA element by creating a data item at the
            specified address.
        
        Args:
            defined_data: XML element containing the address and
                datatype information for the data item
        """
        if self.Options.DataDefinitions.checked == False:
            return
        addr = self.get_address(defined_data, ADDRESS)
        datatype = self.get_attribute(defined_data, DATATYPE)
        size = self.get_attribute_value(defined_data, SIZE)
        self.update_counter(DEFINED_DATA)        
        ti = idaapi.opinfo_t()
        if self.is_pointer_type(datatype):
            #idaapi.set_refinfo(ti, 0, 0, 0, REF_OFF32)
            flag = idaapi.dwrdflag() | idaapi.FF_0OFF
            #idaapi.set_typeinfo(addr, 0, flag, ti)
        else:
            flag = self.get_datatype_flags(datatype, size)
        if flag == idaapi.asciflag():
            idaapi.make_ascii_string(addr, size,
                                     self.get_string_type(datatype))
        elif flag == idaapi.struflag():
            idaapi.doStruct(addr, size, idaapi.get_struc_id(datatype))
        else:
            idaapi.do_data_ex(addr, flag, size, idaapi.BADNODE)        
        typecmt = defined_data.find(TYPEINFO_CMT)
        if typecmt != None:
            self.update_counter(DEFINED_DATA + ':' + TYPEINFO_CMT)


    def import_enum(self, enum):
        """
        Processes an ENUM element by creating the enumeration.
        
        Args:
            enum: XML element containing the enumeration name and
                member data.
        """
        if self.Options.DataTypes.checked == False:
            return
        name = self.get_attribute(enum, NAME)
        if self.has_attribute(enum,NAMESPACE):
            namespace = self.get_attribute(enum, NAMESPACE)
        if self.has_attribute(enum,SIZE):
            size = self.get_attribute_value(enum, SIZE)
        eid = idaapi.add_enum(BADADDR, name,
                              idaapi.hexflag() | idaapi.dwrdflag())
        self.update_counter(ENUM)
        regcmt = enum.find(REGULAR_CMT)
        if regcmt != None:
            idaapi.set_enum_cmt(eid, regcmt.text, False)
            self.update_counter(ENUM + ':' + REGULAR_CMT)
        rptcmt = enum.find(REPEATABLE_CMT)
        if rptcmt != None:
            idaapi.set_enum_cmt(eid, rptcmt.text, True)
            self.update_counter(ENUM + ':' + REPEATABLE_CMT)
        display_settings = enum.find(DISPLAY_SETTINGS)
        if display_settings != None:
            self.update_counter(ENUM + ':' + DISPLAY_SETTINGS)
        enum_entries = enum.findall(ENUM_ENTRY)
        for enum_entry in enum_entries:
            self.import_enum_entry(enum_entry, eid)
    

    def import_enum_entry(self, enum_entry, eid):
        """
        Processes an ENUM_ENTRY by creating a member in the enumeration.
        
        Args:
            enum_entry: XML element containing the member name and value.
            eid: Integer representing the id of the enumeration.
        """
        name = self.get_attribute(enum_entry, NAME)
        value = self.get_attribute_value(enum_entry, VALUE)
        if IDA_SDK_VERSION > 630:
            idaapi.add_const(eid, name, value)
        else:
            idaapi.add_enum_member(eid, name, value)
        if IDA_SDK_VERSION > 630:
            cid = idaapi.get_const_by_name(name)
        else:
            cid = idaapi.get_enum_member_by_name(name)
        self.update_counter(ENUM_ENTRY)
        regcmt = enum_entry.find(REGULAR_CMT)
        if regcmt != None:
            if IDA_SDK_VERSION > 630:
                idaapi.set_const_cmt(cid, regcmt.text, False);
            else:
                idaapi.set_enum_member_cmt(cid, regcmt.text, False);
            self.update_counter(ENUM_ENTRY + ':' + REGULAR_CMT)
        rptcmt = enum_entry.find(REPEATABLE_CMT)
        if rptcmt != None:
            if IDA_SDK_VERSION > 630:
                idaapi.set_const_cmt(cid, rptcmt.text, True);
            else:
                idaapi.set_enum_member_cmt(cid, rptcmt.text, True);
            self.update_counter(ENUM_ENTRY + ':' + REPEATABLE_CMT)
    

    def import_equate(self, equate, eid):
        """
        Processes EQUATE element as member of an enumeration.
        
        Args:
            enum_entry: XML element containing the equate name and value.
            eid: Integer representing the id for the enumeration.
        """
        name = self.get_attribute(equate,NAME)
        value = self.get_attribute_value(equate,VALUE)
        bm = -1
        if self.has_attribute(equate, BIT_MASK):
            bm = self.get_attribute_value(equate, BIT_MASK)
        if IDA_SDK_VERSION > 630:
            #idaapi.add_const(eid, name, value, bm)
            idaapi.add_const(eid, name, value)
            cid = idaapi.get_const_by_name(name)
        else:
            #idaapi.add_enum_member(eid, name, value, bm)
            idaapi.add_enum_member(eid, name, value)
            cid = idaapi.get_enum_member_by_name(name)
        self.update_counter(EQUATE)
        regcmt = equate.find(REGULAR_CMT) 
        if regcmt != None:
            if IDA_SDK_VERSION > 630:
                idaapi.set_const_cmt(cid, regcmt.text, False);
            else:
                idaapi.set_enum_member_cmt(cid, regcmt.text, False);
            self.update_counter(EQUATE + ':' + REGULAR_CMT)
        rptcmt = equate.find(REPEATABLE_CMT)
        if rptcmt != None:
            if IDA_SDK_VERSION > 630:
                idaapi.set_const_cmt(cid, rptcmt.text, True);
            else:
                idaapi.set_enum_member_cmt(cid, rptcmt.text, True);
            self.update_counter(EQUATE + ':' + REPEATABLE_CMT)
    

    def import_equate_group(self, equate_group):
        """
        Processes EQUATE_GROUP as IDA enumeration type.
        
        Args:
            equate_group: XML element containing the group name and
                equate definitions.
        """
        if self.Options.DataTypes.checked == False:
            return
        msg = EQUATE_GROUP
        name = ''
        if self.has_attribute(equate_group, NAME):
            name = self.get_attribute(equate_group, NAME)
        bf = ''
        if self.has_attribute(equate_group, BIT_FIELD):
            bf = self.get_attribute(equate_group, BIT_FIELD)
        eid = idaapi.add_enum(BADADDR, name, idaapi.hexflag())
        idaapi.set_enum_bf(eid, (bf == 'yes'))
        self.update_counter(EQUATE_GROUP)
        regcmt = equate_group.find(REGULAR_CMT)
        if regcmt != None:
            idaapi.set_enum_cmt(eid, regcmt.text, False)
            self.update_counter(EQUATE_GROUP + ':' + REGULAR_CMT)
        rptcmt = equate_group.find(REPEATABLE_CMT)
        if rptcmt != None:
            idaapi.set_enum_cmt(eid, rptcmt.text, True)
            self.update_counter(EQUATE_GROUP + ':' + REPEATABLE_CMT)
        equates = equate_group.findall(EQUATE)
        for equate in equates:
            self.import_equate(equate,eid)
        bit_masks = equate_group.findall(BIT_MASK)
        for bit_mask in bit_masks:
            self.import_bit_mask(bit_mask, eid)
    

    def import_equate_reference(self, equate_reference):
        if (self.Options.DataTypes.checked == False or
            self.Options.EquateReferences.checked == False):
            return
        self.update_counter(EQUATE_REFERENCE)
        addr = self.get_address(equate_reference, ADDRESS)
        name = ''
        if self.has_attribute(equate_reference, NAME):
            name = self.get_attribute(equate_reference, NAME)
        if name == '':
            return
        opnd = 0
        if self.has_attribute(equate_reference, OPERAND_INDEX):
            opnd = self.get_attribute_value(equate_reference, OPERAND_INDEX)
        value = None
        if self.has_attribute(equate_reference, VALUE):
            value = self.get_attribute_value(equate_reference, VALUE)
        if IDA_SDK_VERSION < 650:
            cid = idaapi.get_enum_member_by_name(name)
        else:
            cid = idaapi.get_const_by_name(name)
        if cid == idaapi.BADNODE:
            return
        if IDA_SDK_VERSION < 650:
            eid = idaapi.get_enum_member_enum(cid)
        else:
            eid = idaapi.get_const_enum(cid)
        if eid == idaapi.BADNODE:
            return
        idaapi.op_enum(addr, opnd, eid, 0)
        

    def import_function(self, function):
        """
        Creates a function using the FUNCTION attributes.
        
        Args:
            function: XML element containing the function address and
                attributes.
        """
        if self.Options.Functions.checked == False:
            return
        try:
            entry_point = self.get_address(function, ENTRY_POINT)
            name = ''
            if self.has_attribute(function, NAME):
                name = self.get_attribute(function, NAME)
            libfunc = 'n'
            if self.has_attribute(function, LIBRARY_FUNCTION):
                libfunc = self.get_attribute(function, LIBRARY_FUNCTION)
            if idaapi.isEnabled(entry_point) == False:
                msg = ("import_function: address %X not enabled in database"
                       % entry_point)
                print msg
                return
            idaapi.add_func(entry_point, BADADDR)
            self.update_counter(FUNCTION)
            func = idaapi.get_func(entry_point)
            if libfunc == 'y':
                func.flags |= idaapi.FUNC_LIB
            ranges = function.findall(ADDRESS_RANGE)
            for addr_range in ranges:
                (start, end) = self.import_address_range(addr_range)
                idaapi.append_func_tail(func, start, end)
                #idaapi.analyze_area(start, end+1)
            idc.Wait()
            regcmt = function.find(REGULAR_CMT)
            if regcmt != None:
                self.update_counter(FUNCTION + ':' + REGULAR_CMT)
                idaapi.set_func_cmt(func, regcmt.text, False)
            rptcmt = function.find(REPEATABLE_CMT)
            if rptcmt != None:
                self.update_counter(FUNCTION + ':' + REPEATABLE_CMT)
                idaapi.set_func_cmt(func, rptcmt.text, True)
            typecmt = function.find(TYPEINFO_CMT)
            if typecmt != None:
                self.update_counter(FUNCTION + ':' + TYPEINFO_CMT)
                idc.SetType(entry_point, typecmt.text + ';')
            sf = function.find(STACK_FRAME)
            if sf != None:
                self.import_stack_frame(sf, func)
            register_vars = function.findall(REGISTER_VAR)
            for register_var in register_vars:
                self.import_register_var(register_var, func)
        except:
            msg = "** Exception occurred in import_function **"
            print "\n" + msg + "\n", sys.exc_type, sys.exc_value


    def import_function_def(self, function_def):
        # import_function_def: NOT IMPLEMENTED
        if self.Options.DataTypes.checked == False:
            return
        self.update_counter(FUNCTION_DEF)
        

    def import_info_source(self, info_source):
        """
        Processes INFO_SOURCE containing information about the
            source of the XML PROGRAM file.
        
        Args:
            info_source: XML element containing attributes that identify
                the source of the PROGRAM data.
        """
        if self.has_attribute(info_source, TOOL):
            tool = self.get_attribute(info_source, TOOL)
        if self.has_attribute(info_source, USER):
            user = self.get_attribute(info_source, USER)
        if self.has_attribute(info_source, FILE):
            f = self.get_attribute(info_source, FILE)
        if self.has_attribute(info_source, TIMESTAMP):
            ts = self.get_attribute(info_source, TIMESTAMP)
        self.update_counter(INFO_SOURCE)
    

    def import_manual_instruction(self, manual_instruction):
        """
        Creates a manual instruction.
        
        Args:
            manual_instruction: XML element containing MANUAL_INSTRUCTION.
        """
        if self.Options.Manual.checked == False:
            return
        addr = self.get_address(manual_instruction, ADDRESS)
        idaapi.set_manual_insn(addr, manual_instruction.text)
        self.update_counter(MANUAL_INSTRUCTION)
    

    def import_manual_operand(self, manual_operand):
        """
        Creates a manual operand at an address.
        
        Args:
            manual_operand: MANUAL_OPERAND XML element.
        """
        if self.Options.Manual.checked == False:
            return
        addr = self.get_address(manual_operand, ADDRESS)
        op = self.get_attribute_value(manual_operand, OPERAND_INDEX)
        if idaapi.isEnabled(addr):
            idaapi.set_forced_operand(addr, op, manual_operand.text)
            self.update_counter(MANUAL_OPERAND)


    def process_deferred(self, element):
        """
        Processes the list of deferred structure members when the
        DATATYPES end element is encountered.
        
        Args:
            element: XML end element for DATATYPES 
        """
        for (member, sptr) in self.deferred:
            self.import_member(member, sptr, False)
        self.display_timer(element)


    def import_member(self, member, sptr, defer=True):
        """
        Creates a member for a structure.
        
        Args:
            member: MEMBER XML element.
            sptr:
            defer:  boolean indicating if processing a member should be
                    deferred when the type is unknown. A member should
                    only be deferred on the first pass, not when processing
                    the deferred list.
        """
        offset = self.get_attribute_value(member, OFFSET)
        datatype = self.get_attribute(member, DATATYPE)
        if self.has_attribute(member, DATATYPE_NAMESPACE):
            dt_namespace = self.get_attribute(member, DATATYPE_NAMESPACE)
        name = ''
        if self.has_attribute(member, NAME):
            name = self.get_attribute(member, NAME)
        size = 0
        if self.has_attribute(member, SIZE):
            size = self.get_attribute_value(member, SIZE)
        ti = idaapi.opinfo_t()
        if self.is_pointer_type(datatype):
            #idaapi.set_refinfo(ti, 0, BADADDR, 0, REF_OFF32)
            flag = idaapi.dwrdflag() | idaapi.FF_0OFF
        else:
            flag = self.get_datatype_flags(datatype, size)
        if flag == 0:
            if defer:
                self.deferred.append([member, sptr])
                return                
        if flag == idaapi.enumflag():
            t = idaapi.get_enum(datatype)
            ti.ec.tid = t
            ti.ec.serial = idaapi.get_enum_idx(t)
        if flag == idaapi.struflag():
            t = idaapi.get_struc_id(datatype)
            ti.tid = t
        error = idaapi.add_struc_member(sptr, name, offset, flag, ti, size)
        mbr = idaapi.get_member(sptr, offset)
        self.import_member_cmts(member, mbr)
        self.update_counter(MEMBER)
        
        
    def import_member_cmts(self, member, mbr):
        """
        Processes REGULAR_CMT and REPEATABLE_CMT elements for members.
        
        Args:
            element: XML element object containing a REGULAR_CMT or
                REPEATABLE_CMT element
            mbr: Integer representing the member id
        """
        regcmt = member.find(REGULAR_CMT)
        if regcmt != None:
            idaapi.set_member_cmt(mbr, regcmt.text, False)
            self.update_counter(MEMBER + ':' + REGULAR_CMT)
        rptcmt = member.find(REPEATABLE_CMT)
        if rptcmt != None:
            idaapi.set_member_cmt(mbr, rptcmt.text, True)
            self.update_counter(MEMBER + ':' + REPEATABLE_CMT)
        

    def import_members(self, element, sptr):
        """
        Add data members to a structure.
        
        Args:
            element: STRUCTURE XML element containing MEMBER sub-elements.
            sptr:
        """
        members = element.findall(MEMBER)
        for member in members:
            self.import_member(member, sptr)
                

    def import_memory_contents(self, memory_contents, start, size):
        """
        Processes MEMORY_CONTENTS to load data for a memory block.
        
        Args:
            memory_contents: MEMORY_CONTENTS XML element.
        """
        if memory_contents.get(START_ADDR) == None:
            saddr = start
        else:
            saddr = self.get_address(memory_contents, START_ADDR)
        fname = self.get_attribute(memory_contents, FILE_NAME)
        offset = self.get_attribute_value(memory_contents, FILE_OFFSET)
        if memory_contents.get(LENGTH) == None:
            length = size
        else:
            length = self.get_attribute_value(memory_contents, LENGTH)
        #(binfilename, ext) = os.path.splitext(self.filename)
        #binfilename += ".bytes"
        (binfilename, fileext) = os.path.split(self.filename)
        binfilename += '/' + fname
        binfile = idaapi.loader_input_t()
        binfile.open(binfilename)
        binfile.file2base(offset,saddr,saddr+length,False)
        binfile.close()
        self.update_counter(MEMORY_CONTENTS)
    

    def import_memory_map(self, memory_map):
        """
        Processes the MEMORY_MAP element.
        
        Args:
            memory_map: MEMORY_MAP XML element.
            
        MEMORY_MAP is only processed by the IDA loader. It is ignored when
            run as an IDA plugin.
        """
        # import memory sections only when run as loader
        if self.plugin:
            return
        self.update_import(memory_map)
    

    def import_memory_reference(self, memory_reference):
        """
        Processes the MEMORY_REFERENCE element.
        
        Args:
            memory_reference: MEMORY_REFERENCE XML element.
        """
        if self.Options.MemoryReferences.checked == False:
            return        
        addr = self.get_address(memory_reference, ADDRESS)
        if self.has_attribute(memory_reference, OPERAND_INDEX):
            op = self.get_attribute_value(memory_reference, OPERAND_INDEX)
        if self.has_attribute(memory_reference, USER_DEFINED):
            user = self.get_attribute(memory_reference, USER_DEFINED)
        to_addr = self.get_address(memory_reference, TO_ADDRESS)
        if self.has_attribute(memory_reference, BASE_ADDRESS):
            base_addr = self.get_address(memory_reference, BASE_ADDRESS)
        if self.has_attribute(memory_reference, PRIMARY):
            primary = self.get_attribute(memory_reference, PRIMARY)
        self.update_counter(MEMORY_REFERENCE)
        # TODO: import_memory_reference - add code to store reference
 
 
    def import_memory_section(self, memory_section):
        """
        Creates a memory segment in the database.
        
        Args:
            memory_section: MEMORY_SECTION XML element.
            
        MEMORY_SECTION is only processed by the IDA loader. It is ignored
            when run as an IDA plugin.
        """
        # TODO: import_memory_section - handle overlays?
        # import memory sections only when run as loader
        if self.plugin:
            return
        name = self.get_attribute(memory_section, NAME)
        length = self.get_attribute_value(memory_section, LENGTH)

        s = idaapi.segment_t()
        addrstr = self.get_attribute(memory_section, START_ADDR)
        seg_str = ''
        if '::' in addrstr:
            # overlay - skip for now
            print '  ** Overlayed memory block %s skipped **  ' % name
            msg  = 'Overlayed memory block %s skipped!' % name
            msg += "\n\nXML Import does not currently support"
            msg += "\noverlayed memory blocks."
            idaapi.warning(msg)
            return
        elif ':' in addrstr:
            [seg_str, offset_str] = string.split(addrstr,':')
            offset = int(offset_str, 16)
            if self.is_int(seg_str):
                base = int(seg_str, 16)
                sel = idaapi.setup_selector(base)
                start = self.get_address(memory_section, START_ADDR)
            else:
                raise MultipleAddressSpacesNotSupported
                return
        else:
            sel = idaapi.allocate_selector(0)
            start = int(addrstr, 16)

        s.sel = sel
        s.startEA = start
        s.endEA = start+length
        s.bitness = self.addr_mode

        perms = ''
        if self.has_attribute(memory_section, PERMISSIONS):
            perms = self.get_attribute(memory_section, PERMISSIONS)
        s.perm = 0
        if 'r' in perms: s.perm |= idaapi.SEGPERM_READ
        if 'w' in perms: s.perm |= idaapi.SEGPERM_WRITE
        if 'x' in perms: s.perm |= idaapi.SEGPERM_EXEC
        idaapi.add_segm_ex(s, name, "", idaapi.ADDSEG_OR_DIE |
                                            idaapi.ADDSEG_QUIET)
        self.update_counter(MEMORY_SECTION)
        for memory_contents in memory_section.findall(MEMORY_CONTENTS):
            self.import_memory_contents(memory_contents, start, length)
    

    def import_processor(self, processor):
        """
        Processes the PROCESSOR element.
        
        Args:
            processor: PROCESSOR XML element.
        """
        name = self.get_attribute(processor, NAME)
        self.update_counter(PROCESSOR)
        if self.plugin:
            return
        address_model = self.get_attribute(processor, ADDRESS_MODEL)
        if address_model != None:
            flag = idaapi.ph_get_flag()
            if str.lower(address_model) == '16-bit':
                self.addr_mode = 0
            elif str.lower(address_model) == '32-bit':
                self.addr_mode = 1
            elif str.lower(address_model) == '64-bit':
                self.addr_mode = 2
        

    def import_program(self, program):
        """
        Processes the PROGRAM element.
        
        Args:
            program: PROGRAM XML element.
        """
        self.update_status(PROGRAM)
        name = self.get_attribute(program, NAME)
        if self.has_attribute(program, EXE_PATH):
            epath = self.get_attribute(program, EXE_PATH)
            idaapi.set_root_filename(epath)
        else:
            idaapi.set_root_filename(name)
        if self.has_attribute(program, EXE_FORMAT):
            eformat = self.get_attribute(program, EXE_FORMAT)
            RootNode = idaapi.netnode('Root Node')
            if IDA_SDK_VERSION < 650:
                RootNode.supset(1, eformat)
            else:
                RootNode.supset(idaapi.RIDX_FILE_FORMAT_NAME, eformat)
        if self.has_attribute(program, IMAGE_BASE):
            base = self.get_attribute_value(program, IMAGE_BASE)
            idaapi.set_imagebase(base)
        if self.has_attribute(program, INPUT_MD5):
            input_md5 = self.get_attribute(program, INPUT_MD5)
            # store original md5 in a special netnode
            md5 = idaapi.netnode(INPUT_MD5, len(INPUT_MD5), True)
            if (IDA_SDK_VERSION < 650):
                md5.supset(1302, input_md5)
            else:
                md5.supset(idaapi.RIDX_MD5, input_md5)
        self.update_counter(PROGRAM)
        """
        # TODO: this needs to be on "end" event for PROGRAM
        description = program.find(DESCRIPTION)
        if description != None:
            pass  # figure out what to do with it
        """

        
    def import_program_entry_point(self, program_entry_point):
        """
        Defines a program entry point.
        
        Args:
            program_entry_point: PROGRAM_ENTRY_POINT XML element.
                Contains the entry point address.
        """
        if self.Options.EntryPoints.checked == False:
            return
        addr = self.get_address(program_entry_point, ADDRESS)
        idaapi.add_entry(addr, addr, "", True)
        self.update_counter(PROGRAM_ENTRY_POINT)
    

    def import_register_value_range(self, register_value_range):
        """
        Defines the address range for a register value.
        
        Args:
            register_value_range: REGISTER_VALUE_RANGE XML element.
                Contains the register, value, start address and range length.
        """
        if self.Options.RegisterValues.checked == False:
            return
        self.update_counter(REGISTER_VALUE_RANGE)
        reg = self.get_attribute(register_value_range, REGISTER)
        if reg == 'cs': return
        value = self.get_attribute_value(register_value_range, VALUE)
        addr = self.get_address(register_value_range, START_ADDRESS)
        length = self.get_attribute_value(register_value_range, LENGTH)
        r = idaapi.str2reg(reg)
        if r >= idaapi.ph.regFirstSreg and r <= idaapi.ph.regLastSreg:
            if IDA_SDK_VERSION < 670:
                idaapi.splitSRarea1(addr, r, value, 2)
            else:
                idaapi.splitSRarea1(addr, r, value, idaapi.SR_user)
    

    def import_register_var(self, register_var, func):
        """
        Defines a register variable for a function.
        
        Args:
            register_var: REGISTER_VAR XML element.
                Contains register, variable name, and datatype.
            func: IDA function object
        """
        name = self.get_attribute(register_var, NAME)
        reg = self.get_attribute(register_var, REGISTER)
        if self.has_attribute(register_var, DATATYPE):
            datatype = self.get_attribute(register_var, DATATYPE)
        if self.has_attribute(register_var, DATATYPE_NAMESPACE):
            namespace = self.get_attribute(register_var, DATATYPE_NAMESPACE)
        idc.MakeLocal(func.startEA, func.endEA, reg, name)
        self.update_counter(REGISTER_VAR)
    

    def import_stack_frame(self, stack_frame, func):
        """
        Defines a stack frame for a function.
        
        Args:
            stack_frame: STACK_FRAME element with STACK_VAR child elements.
        """
        if self.has_attribute(stack_frame, LOCAL_VAR_SIZE):
            lvsize = self.get_attribute_value(stack_frame, LOCAL_VAR_SIZE)
        if self.has_attribute(stack_frame, PARAM_OFFSET):
            param_offset = self.get_attribute_value(stack_frame, PARAM_OFFSET)
        if self.has_attribute(stack_frame, REGISTER_SAVE_SIZE):
            reg_save_size = self.get_attribute_value(stack_frame,
                                                     REGISTER_SAVE_SIZE)
        if self.has_attribute(stack_frame, RETURN_ADDR_SIZE):
            retaddr_size = self.get_attribute_value(stack_frame,
                                                    RETURN_ADDR_SIZE)
        if self.has_attribute(stack_frame, BYTES_PURGED):
            bytes_purged = self.get_attribute_value(stack_frame, BYTES_PURGED)
        self.update_counter(STACK_FRAME)
        for stack_var in stack_frame.findall(STACK_VAR):
            self.import_stack_var(stack_var, func)
    

    def import_stack_reference(self, stack_reference):
        # import_stack_reference: NOT IMPLEMENTED
        #self.update_counter(STACK_REFERENCE)
        pass
                    

    def import_stack_var(self, stack_var, func):
        """
        Processes STACK_VAR element.
        
        Args:
            stack_var: STACK_VAR XML element.
            
        Stack variables are created by IDA's function analysis. 
        Only the STACK_VAR NAME attribute is used to set the name for
        a stack variable at the specified stack/frame offset. 
        """
        spoffset = self.get_attribute_value(stack_var, STACK_PTR_OFFSET)
        datatype = self.get_attribute(stack_var, DATATYPE)
        offset = spoffset + func.frsize + func.frregs
        if self.has_attribute(stack_var, FRAME_PTR_OFFSET):
            fpoffset = self.get_attribute_value(stack_var, FRAME_PTR_OFFSET)
            offset = fpoffset + func.frsize
        name = ''
        if self.has_attribute(stack_var, NAME):
            name = self.get_attribute(stack_var, NAME)
        if self.has_attribute(stack_var, DATATYPE_NAMESPACE):
            namespace = self.get_attribute(stack_var, DATATYPE_NAMESPACE)
        if self.has_attribute(stack_var, SIZE):
            size = self.get_attribute_value(stack_var, SIZE)
        self.update_counter(STACK_VAR)
        sf = idaapi.get_frame(func)
        if name != '':
            idaapi.set_member_name(sf, offset, name)
    

    def import_structure(self, structure):
        """
        Adds a structure.
        
        Args:
            structure: STRUCTURE XML element.
                Contains the STRUCTURE attributes and child elements.
        """
        if self.Options.DataTypes.checked == False:
            return
        name = self.get_attribute(structure, NAME)
        dtyp = idaapi.get_struc_id(name)
        if dtyp != idaapi.BADNODE:
            # duplicate name, try adding name space
            if self.has_attribute(structure, NAMESPACE) == False:
                return
            namespace = self.get_attribute(structure, NAMESPACE)
            name = namspace + '__' + name
            name.replace('/','_')
            name.replace('.','_')
            dtyp = idaapi.get_struc_id(name)
            # skip if still duplicate (could add sequence #)
            if dtyp != idaapi.BADNODE:
                return
        size = 0
        if self.has_attribute(structure, SIZE):
            size = self.get_attribute_value(structure, SIZE)
        if self.has_attribute(structure, VARIABLE_LENGTH):
            vl = self.get_attribute_value(structure, VARIABLE_LENGTH)
            isVariableLength = vl == 'y'
        sid = idaapi.add_struc(BADADDR, name, False)
        sptr = idaapi.get_struc(sid)
        self.update_counter(STRUCTURE)
        self.import_cmts(structure, sid, STRUCTURE)
        self.import_members(structure, sptr)
        t = idaapi.opinfo_t()
        if idaapi.get_struc_size(sptr) < size:
            idaapi.add_struc_member(sptr,"",size-1,idaapi.byteflag(),t,1)
        

    def import_symbol(self, symbol):
        """
        Adds a symbol name at the specified address.
        
        Args:
            symbol: SYMBOL XML element.
                Contains symbol name and address. Optionally includes
                type and mangled symbol.
        """
        if self.Options.Symbols.checked == False:
            return
        addr = self.get_address(symbol, ADDRESS)
        name = self.get_attribute(symbol, NAME)
        if self.has_attribute(symbol, MANGLED):
            name = self.get_attribute(symbol, MANGLED)
        flag = idaapi.SN_NOWARN
        if self.has_attribute(symbol, TYPE):
            typ = self.get_attribute(symbol, TYPE)
            if  typ == 'local': flag |= idaapi.SN_LOCAL
        idaapi.set_name(addr, name, flag)
        self.update_counter(SYMBOL)
    

    def import_typedef(self, type_def):
        # import_typedef: NOT IMPLEMENTED
        if self.Options.DataTypes.checked == False:
            return
        self.update_counter(TYPE_DEF)
        

    def import_union(self, union):
        """
        Adds a union datatype.
        
        Args:
            union: UNION XML element.
                Contains UNION attributes and child elements.
        """
        if self.Options.DataTypes.checked == False:
            return
        name = self.get_attribute(union, NAME)
        dtyp = idaapi.get_struc_id(name)
        if dtyp != idaapi.BADNODE:
            # duplicate name, try adding name space
            if self.has_attribute(union, NAMESPACE) == False:
                return
            namespace = self.get_attribute(union, NAMESPACE)
            name = namspace + '__' + name
            name.replace('/','_')
            name.replace('.','_')
            dtyp = idaapi.get_struc_id(name)
            # skip if still duplicate (could add sequence #)
            if dtyp != idaapi.BADNODE:
                return
        size = 0
        if self.has_attribute(union, SIZE):
            size = self.get_attribute_value(union, SIZE)
        sid = idaapi.add_struc(BADADDR, name, True)
        sptr = idaapi.get_struc(sid)
        self.update_counter(UNION)
        self.import_cmts(union, sid, UNION)
        self.import_members(union, sptr)
        t = idaapi.opinfo_t()
        if idaapi.get_struc_size(sptr) < size:
            idaapi.add_struc_member(sptr,"", size-1, idaapi.byteflag(), None, 1)
            

    def open_file(self, filename, mode):
        """
        Opens filename to specified mode.
        
        Args:
            filename: String representing absolute filepath.
            mode: String representing mode for open.
            
        Returns
            File handle.
        
        Exceptions:
            Displays a warning and raises FileError exception
            if open fails.
        """
        try:
            f = open(filename, mode)
            return f
        except:
            fmt = "TITLE ERROR!\n"
            fmt += "ICON ERROR\n"
            fmt += "AUTOHIDE NONE\n"
            fmt += "HIDECANCEL\n"
            fmt += "Error opening file" + filename + "!\n"
            idaapi.warning(fmt)
            raise FileError
    

    def update_counter(self, tag):
        """
        Updates the counter for the element tag.
        
        Args:
            tag: String representing element tag.
        """
        if tag in self.Elements:
            self.Counters[self.Elements[tag]] += 1
        else:
            self.Elements[tag] = len(self.Elements)
            self.Counters.append(1)
            self.Tags.append(tag)
    

    def update_import(self, element):
        """
        Update the element counter and processing status.
        
        Args:
            element: XML element
            
        This function is used to process certain high-level elements
        (such as COMMENTS, CODE_BLOCKS, SYMBOL_TABLE, FUNCTIONS, etc.)
        that are used to group sub-elements.
        """
        self.update_counter(element.tag)
        self.update_status(element.tag)
    

    def update_status(self, tag):
        """
        Displays the processing status in the IDA window.
        
        Args:
            tag: String representing XML element tag
        """
        status = 'Importing ' + tag
        idaapi.msg('\n%-35s' % status)
        idaapi.hide_wait_box()
        idaapi.show_wait_box(status)
    
        
# Global XML string constants for elements and attributes
ADDRESS = 'ADDRESS'
ADDRESS_MODEL = 'ADDRESS_MODEL'
ADDRESS_RANGE = 'ADDRESS_RANGE'
BASE_ADDRESS = 'BASE_ADDRESS'
BIT_FIELD = 'BIT_FIELD'
BIT_MAPPED = 'BIT_MAPPED'
BIT_MASK = 'BIT_MASK'
BOOKMARK = 'BOOKMARK'
BOOKMARKS = 'BOOKMARKS'
BYTES = 'BYTES'
BYTES_PURGED = 'BYTES_PURGED'
CATEGORY = 'CATEGORY'
CODE = 'CODE'
CODE_BLOCK = 'CODE_BLOCK'
COMMENT = 'COMMENT'
COMMENTS = 'COMMENTS'
COMPILER = 'COMPILER'
DATA = 'DATA'
DATATYPE = 'DATATYPE'
DATATYPES = 'DATATYPES'
DATATYPE_NAMESPACE = 'DATATYPE_NAMESPACE'
DEFINED_DATA = 'DEFINED_DATA'
DESCRIPTION = 'DESCRIPTION'
DISPLAY_SETTINGS = 'DISPLAY_SETTINGS'
END = 'END'
ENDIAN = 'ENDIAN'
ENTRY_POINT = 'ENTRY_POINT'
ENUM = 'ENUM'
ENUM_ENTRY = 'ENUM_ENTRY'
EQUATE = 'EQUATE'
EQUATES = 'EQUATES'
EQUATE_GROUP = 'EQUATE_GROUP'
EQUATE_REFERENCE = 'EQUATE_REFERENCE'
EXE_FORMAT = 'EXE_FORMAT'
EXE_PATH = 'EXE_PATH'
EXT_LIBRARY = 'EXT_LIBRARY'
EXT_LIBRARY_REFERENCE = 'EXT_LIBRARY_REFERENCE'
EXT_LIBRARY_TABLE = 'EXT_LIBRARY_TABLE'
FAMILY = 'FAMILY'
FILE = 'FILE'
FILE_NAME = 'FILE_NAME'
FILE_OFFSET = 'FILE_OFFSET'
FOLDER = 'FOLDER'
FORMAT = 'FORMAT'
FRAGMENT = 'FRAGMENT'
FRAME_PTR_OFFSET = 'FRAME_PTR_OFFSET'
FUNCTION = 'FUNCTION'
FUNCTIONS = 'FUNCTIONS'
FUNCTION_DEF = 'FUNCTION_DEF'
IMAGE_BASE = 'IMAGE_BASE'
INPUT_MD5 = 'INPUT_MD5'
INFO_SOURCE = 'INFO_SOURCE'
LANGUAGE_PROVIDER = 'LANGUAGE_PROVIDER'
LENGTH = 'LENGTH'
LIB_ADDR = 'LIB_ADDR'
LIB_LABEL = 'LIB_LABEL'
LIB_ORDINAL = 'LIB_ORDINAL'
LIB_PROG_NAME = 'LIB_PROG_NAME'
LIBRARY_FUNCTION = 'LIBRARY_FUNCTION'
LOCAL_VAR_SIZE = 'LOCAL_VAR_SIZE'
MANGLED = 'MANGLED'
MANUAL_INSTRUCTION = 'MANUAL_INSTRUCTION'
MANUAL_OPERAND = 'MANUAL_OPERAND'
MARKUP = 'MARKUP'
MEMBER = 'MEMBER'
MEMORY_CONTENTS = 'MEMORY_CONTENTS'
MEMORY_MAP = 'MEMORY_MAP'
MEMORY_REFERENCE = 'MEMORY_REFERENCE'
MEMORY_SECTION = 'MEMORY_SECTION'
NAME = 'NAME'
NAMESPACE = 'NAMESPACE'
OFFSET = 'OFFSET'
OPERAND_INDEX = 'OPERAND_INDEX'
PARAM_OFFSET = 'PARAM_OFFSET'
PATH = 'PATH'
PERMISSIONS = 'PERMISSIONS'
PRIMARY = 'PRIMARY'
PROCESSOR = 'PROCESSOR'
PROGRAM = 'PROGRAM'
PROGRAM_ENTRY_POINT = 'PROGRAM_ENTRY_POINT'
PROGRAM_ENTRY_POINTS = 'PROGRAM_ENTRY_POINTS'
PROGRAM_TREES = 'PROGRAM_TREES'
PROPERTIES = 'PROPERTIES'
PROPERTY = 'PROPERTY'
REGISTER = 'REGISTER'
REGISTER_SAVE_SIZE = 'REGISTER_SAVE_SIZE'
REGISTER_VALUES = 'REGISTER_VALUES'
REGISTER_VALUE_RANGE = 'REGISTER_VALUE_RANGE'
REGISTER_VAR = 'REGISTER_VAR'
REGULAR_CMT = 'REGULAR_CMT'
RELOCATION = 'RELOCATION'
RELOCATION_TABLE = 'RELOCATION_TABLE'
REPEATABLE_CMT = 'REPEATABLE_CMT'
RETURN_ADDR_SIZE = 'RETURN_ADDR_SIZE'
RETURN_TYPE = 'RETURN_TYPE'
SHOW_TERMINATOR = 'SHOW_TERMINATOR'
SIGNED = 'SIGNED'
SIZE = 'SIZE'
SOURCE_ADDRESS = 'SOURCE_ADDRESS'
SOURCE_TYPE = 'SOURCE_TYPE'
STACK_FRAME = 'STACK_FRAME'
STACK_PTR_OFFSET = 'STACK_PTR_OFFSET'
STACK_REFERENCE = 'STACK_REFERENCE'
STACK_VAR = 'STACK_VAR'
START = 'START'
START_ADDR = 'START_ADDR'
START_ADDRESS = 'START_ADDRESS'
STRUCTURE = 'STRUCTURE'
SYMBOL = 'SYMBOL'
SYMBOL_TABLE = 'SYMBOL_TABLE'
TIMESTAMP = 'TIMESTAMP'
TOOL = 'TOOL'
TO_ADDRESS = 'TO_ADDRESS'
TREE = 'TREE'
TYPE = 'TYPE'
TYPEINFO_CMT = 'TYPEINFO_CMT'
TYPE_DEF = 'TYPE_DEF'
UNION = 'UNION'
USER = 'USER'
USER_DEFINED = 'USER_DEFINED'
VALUE = 'VALUE'
VARIABLE_LENGTH = 'VARIABLE_LENGTH'
ZERO_PAD = 'ZERO_PAD'


repository: fengjixuchui/decompile
file: bin/ghidra_9.0_PUBLIC_20190228/Extensions/IDAPro/Python/6xx/plugins/xmlldr.py
description: 反编译工具集锦--持续收集
wanted_name: XML Importer
comment: Import XML PROGRAM file
created_at: 2019-10-28 01:45:02+00:00
pushed_at: 2020-05-08 02:21:55+00:00
url:
https://github.com/fengjixuchui/decompile/blob/abb3e5892572ff8657ac5255bbbf2466f1477168/bin/ghidra_9.0_PUBLIC_20190228/Extensions/IDAPro/Python/6xx/plugins/xmlldr.py


repository: fengjixuchui/decompile
file: bin/ghidra_9.0_PUBLIC_20190228/Extensions/IDAPro/Python/6xx/loaders/xmlldr.py
description: 反编译工具集锦--持续收集
wanted_name: XML Importer
comment: Import XML PROGRAM file
created_at: 2019-10-28 01:45:02+00:00
pushed_at: 2020-05-08 02:21:55+00:00
url:
https://github.com/fengjixuchui/decompile/blob/abb3e5892572ff8657ac5255bbbf2466f1477168/bin/ghidra_9.0_PUBLIC_20190228/Extensions/IDAPro/Python/6xx/loaders/xmlldr.py


repository: Ph0en1x-XMU/Ph0en1x-Team
file: Reverse/ida/IDA Pro 7.5.201028 (x86,x64,ARM,ARM64,PPC,PPC64,MIPS)/python/examples/hexrays/vds10.py
wanted_name: Optimize DbgRaiseAssertionFailure (IDAPython)
comment: Sample plugin10 for Hex-Rays decompiler
created_at: 2022-06-02 11:06:16+00:00
pushed_at: 2022-06-02 11:11:40+00:00
url:
https://github.com/Ph0en1x-XMU/Ph0en1x-Team/blob/a8d80aaea8227e418f5826bf05ef69a1bc8d84f1/Reverse/ida/IDA%20Pro%207.5.201028%20(x86,x64,ARM,ARM64,PPC,PPC64,MIPS)/python/examples/hexrays/vds10.py

content:
#
#      Hex-Rays Decompiler project
#      Copyright (c) 2007-2020 by Hex-Rays, support@hex-rays.com
#      ALL RIGHTS RESERVED.
#
#      Sample plugin for Hex-Rays Decompiler.
#      It installs a custom microcode optimization rule:
#        call   !DbgRaiseAssertionFailure <fast:>.0
#      =>
#        call   !DbgRaiseAssertionFailure <fast:"char *" "assertion text">.0
#
#      To see this plugin in action please use arm64_brk.i64, in the hexrays sdk
#
#      This is a rewrite in Python of the vds10 example that comes with hexrays sdk.
#

import ida_bytes
import ida_range
import ida_kernwin
import ida_hexrays
import ida_typeinf
import ida_idaapi

class nt_assert_optimizer_t(ida_hexrays.optinsn_t):
    def func(self, blk, ins, optflags):
        if self.handle_nt_assert(ins):
            return 1
        return 0

    def handle_nt_assert(self, ins):
        # recognize call   !DbgRaiseAssertionFailure <fast:>.0
        if not ins.is_helper("DbgRaiseAssertionFailure"):
            return False

        # did we already add an argument?
        fi = ins.d.f;
        if not fi.args.empty():
            return False

        # use a comment from the disassembly listing as the call argument
        cmt = ida_bytes.get_cmt(ins.ea, False)
        if not cmt:
            return False

        # remove "NT_ASSERT("...")" to make the listing nicer
        if cmt.startswith("NT_ASSERT(\""):
            cmt = cmt[11:]
            if cmt.endswith("\")"):
                cmt = cmt[:-2]

        # all ok, transform the instruction by adding one more call argument
        fa = fi.args.push_back()
        fa.t = ida_hexrays.mop_str;
        fa.cstr = cmt
        fa.type = ida_typeinf.tinfo_t.get_stock(ida_typeinf.STI_PCCHAR) # const char *
        fa.size = fa.type.get_size()
        return True

# --------------------------------------------------------------------------
# a plugin interface, boilerplate code
class my_plugin_t(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_HIDE
    wanted_name = "Optimize DbgRaiseAssertionFailure (IDAPython)"
    wanted_hotkey = ""
    comment = "Sample plugin10 for Hex-Rays decompiler"
    help = ""
    def init(self):
        if ida_hexrays.init_hexrays_plugin():
            self.optimizer = nt_assert_optimizer_t()
            self.optimizer.install()
        return ida_idaapi.PLUGIN_KEEP # keep us in the memory
    def term(self):
        self.optimizer.remove()
    def run(self, arg):
        if arg == 1:
            return self.optimizer.remove()
        elif arg == 2:
            return self.optimizer.install()

def PLUGIN_ENTRY():
    return my_plugin_t()


repository: mandiant/flare-ida
file: plugins/stackstrings_plugin.py
description: IDA Pro utilities from FLARE team
created_at: 2014-07-30 18:42:23+00:00
pushed_at: 2024-10-29 14:38:07+00:00
url:
https://github.com/mandiant/flare-ida/blob/fc667de6db279c6f8ec8bcc2f3eedff348db450f/plugins/stackstrings_plugin.py

content:
#!/usr/bin/env python
# Jay Smith
# jay.smith@fireeye.com
# 
########################################################################
# Copyright 2012 Mandiant
# Copyright 2014 FireEye
# Copyright 2019 FireEye
#
# Mandiant licenses this file to you 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.
########################################################################
#
# IDA Plugin wrapper for stack strings search
#
########################################################################


import sys
import logging

import idc 
import idautils  
import idaapi

is_py2 = (sys.version_info[0] == 2)

if is_py2:
    # currently depending on vivisect, which will never be ported to py3
    idaapi.require('flare')
    idaapi.require('flare.stackstrings')

PLUGIN_COMMENT = "This is a comment"
PLUGIN_HELP = "This is help"
PLUGIN_NAME = "StackStrings"
PLUGIN_WANTED_HOTKEY = "Alt-0"

# get the IDA version number
ida_major, ida_minor = map(int, idaapi.get_kernel_version().split("."))
using_ida7api = (ida_major > 6)
ex_addmenu_item_ctx = None 


class stackstrings_plugin_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_KEEP
    comment = PLUGIN_COMMENT
    help = PLUGIN_HELP
    wanted_name = PLUGIN_NAME
    wanted_hotkey = PLUGIN_WANTED_HOTKEY

    def init(self):
        try:
            idaapi.msg("StackStrings init() called!\n")
            return idaapi.PLUGIN_OK
        except Exception as err:
            idaapi.msg("Exception during init: %s\n" % str(err))
        
        return idaapi.PLUGIN_SKIP

    def run(self, arg):
        try:
            idaapi.msg("StackStrings run() called with %d!\n" % arg)
            if is_py2:
                flare.stackstrings.main()
                idaapi.msg("StackStrings run() done")
            else:
                idaapi.msg("WARNING: stackstrings only works under python2 due to vivisect dependency\n")
        except Exception as err:
            idaapi.msg("Exception during run: %s\n" % str(err))
            raise
            
        idaapi.msg("StackStrings run() complete!\n")

    def term(self):
        idaapi.msg("StackStrings term() called!\n")

def PLUGIN_ENTRY():
    return stackstrings_plugin_t()


repository: joxeankoret/diaphora
file: jkutils/IDAMagicStrings.py
description: Diaphora, the most advanced Free and Open Source program diffing tool.
wanted_name: Magic Strings
comment: IDA Magic Strings plugin
created_at: 2015-03-20 02:12:26+00:00
pushed_at: 2024-11-24 08:49:59+00:00
url:
https://github.com/joxeankoret/diaphora/blob/d8b898e35750d103cb6a15f7a87910ff6b237427/jkutils/IDAMagicStrings.py

content:
#-------------------------------------------------------------------------------
#
# IDAPython script to show many features extracted from debugging strings. It's
# also able to rename functions based on the guessed function name & rename
# functions based on the source code file they belong to.
#
# Copyright (c) 2018-2024, Joxean Koret
# Licensed under the GNU Affero General Public License v3.
#
#-------------------------------------------------------------------------------

import os
import sys

import re

from collections import Counter

import idaapi

from idc import *
from idaapi import *
from idautils import *

try:
  from PyQt5 import QtWidgets
except ImportError as e:
  print(f'{os.path.basename(__file__)} importerror {e}')

sys.path.append("..")
import diaphora_config as config

try:
  import nltk
  from nltk.tokenize import word_tokenize
  from nltk.tag import pos_tag
  has_nltk = True
except ImportError as e:
  if config.SHOW_IMPORT_WARNINGS:
    print("WARNING: NLTK is not installed. It's recommended to install python-nltk. It's optional, but significantly improves the results.")
    print("INFO: Alternatively, you can silence this warning by changing the value of SHOW_IMPORT_WARNINGS in diaphora_config.py.")
  has_nltk = False

#-------------------------------------------------------------------------------
PROGRAM_NAME = "IMS"

#-------------------------------------------------------------------------------
SOURCE_FILES_REGEXP = r"([a-z_\/\\][a-z0-9_/\\:\-\.@]+\.(c|cc|cxx|c\+\+|cpp|h|hpp|m|rs|go|ml))($|:| )"

LANGS = {}
LANGS["C/C++"] = ["c", "cc", "cxx", "cpp", "h", "hpp"]
LANGS["C"] = ["c"]
LANGS["C++"] = ["cc", "cxx", "cpp", "hpp", "c++"]
LANGS["Obj-C"] = ["m"]
LANGS["Rust"] = ["rs"]
LANGS["Golang"] = ["go"]
LANGS["OCaml"] = ["ml"]

#-------------------------------------------------------------------------------
FUNCTION_NAMES_REGEXP = r"([a-z_][a-z0-9_]+((::)+[a-z_][a-z0-9_]+)*)"
CLASS_NAMES_REGEXP    = r"([a-z_][a-z0-9_]+(::(<[a-z0-9_]+>|~{0,1}[a-z0-9_]+))+)\({0,1}"
NOT_FUNCTION_NAMES = ["copyright", "char", "bool", "int", "unsigned", "long",
  "double", "float", "signed", "license", "version", "cannot", "error",
  "invalid", "null", "warning", "general", "argument", "written", "report",
  "failed", "assert", "object", "integer", "unknown", "localhost", "native",
  "memory", "system", "write", "read", "open", "close", "help", "exit", "test",
  "return", "libs", "home", "ambiguous", "internal", "request", "inserting",
  "deleting", "removing", "updating", "adding", "assertion", "flags",
  "overflow", "enabled", "disabled", "enable", "disable", "virtual", "client",
  "server", "switch", "while", "offset", "abort", "panic", "static", "updated",
  "pointer", "reason", "month", "year", "week", "hour", "minute", "second",
  'monday', 'tuesday', 'wednesday', 'thursday', 'friday', 'saturday', 'sunday',
  'january', 'february', 'march', 'april', 'may', 'june', 'july', 'august',
  'september', 'october', 'november', 'december', "arguments", "corrupt",
  "corrupted", "default", "success", "expecting", "missing", "phrase",
  "unrecognized", "undefined",
  ]

#-------------------------------------------------------------------------------
FOUND_TOKENS = {}
TOKEN_TYPES = ["NN", "NNS", "NNP", "JJ", "VB", "VBD", "VBG", "VBN", "VBP", "VBZ"]
def nltk_preprocess(strings):
  if not has_nltk:
    return

  strings = "\n".join(map(str, list(strings)))
  tokens = re.findall(FUNCTION_NAMES_REGEXP, strings)
  l = []
  for token in tokens:
    l.append(token[0])
  word_tags = nltk.pos_tag(l)
  for word, tag in word_tags:
    try:
      FOUND_TOKENS[word.lower()].add(tag)
    except:
      FOUND_TOKENS[word.lower()] = set([tag])

#-------------------------------------------------------------------------------
def get_strings(strtypes = [0, 1]):
  strings = Strings()
  strings.setup(strtypes = strtypes)
  return strings

#-------------------------------------------------------------------------------
def get_lang(full_path):
  """
  Try to guess the programming language from the file extension.
  """
  _, file_ext  = os.path.splitext(full_path.lower())
  file_ext = file_ext.strip(".")
  for key in LANGS:
    if file_ext in LANGS[key]:
      return key
  return None

#-------------------------------------------------------------------------------
def add_source_file_to(d, src_langs, refs, full_path, s):
  """
  Add single source file to the given dictionary and try to guess which language
  was that source file written in.
  """
  if full_path not in d:
    d[full_path] = []

  lang = get_lang(full_path)
  if lang is not None:
    src_langs[lang] += 1

  for ref in refs:
    d[full_path].append([ref, get_func_name(ref), str(s)])

  return d, src_langs

#-------------------------------------------------------------------------------
def get_source_strings(min_len, strtypes):
  """
  Find source code files from strings in the binary
  """
  strings = get_strings(strtypes)

  # Search string references to source files
  src_langs = Counter()
  total_files = 0
  d = {}
  for s in strings:
    if s and s.length > min_len:
      ret = re.findall(SOURCE_FILES_REGEXP, str(s), re.IGNORECASE)
      if ret and len(ret) > 0:
        refs = list(DataRefsTo(s.ea))
        if len(refs) > 0:
          total_files += 1
          full_path    = ret[0][0]
          d, src_langs = add_source_file_to(d, src_langs, refs, full_path, s)

  # Use the loaded debugging information (if any) to find source files
  for f in list(Functions()):
    done = False
    func = idaapi.get_func(f)
    if func is not None:
      cfg = idaapi.FlowChart(func)
      for block in cfg:
        if done:
          break

        head = block.start_ea
        full_path = get_sourcefile(head)
        if full_path is not None:
          total_files += 1
          d, src_langs = add_source_file_to(d, src_langs, [head], full_path, f"Symbol: {full_path}")

  nltk_preprocess(strings)
  if len(d) > 0 and total_files > 0:
    print("Programming languages found:\n")
    for key in src_langs:
      percent = src_langs[key] * 100. / total_files
      print(f"  {key.ljust(10)} {percent}%")
    print()

  return d, strings

#-------------------------------------------------------------------------------
def handler(item, column_no):
  ea = item.ea
  if is_mapped(ea):
    jumpto(ea)

#-------------------------------------------------------------------------------
class CBaseTreeViewer(PluginForm):
  def populate_tree(self, d):
    # Clear previous items
    self.tree.clear()

    # Build the tree
    for key in d:
      src_file_item = QtWidgets.QTreeWidgetItem(self.tree)
      src_file_item.setText(0, key)
      src_file_item.ea = BADADDR

      for ea, name, str_data in d[key]:
        item = QtWidgets.QTreeWidgetItem(src_file_item)
        item.setText(0, "%s [0x%08x] %s" % (name, ea, str_data))
        item.ea = ea

    self.tree.itemDoubleClicked.connect(handler)

  def OnCreate(self, form):
    # Get parent widget
    self.parent = idaapi.PluginForm.FormToPyQtWidget(form)

    # Create tree control
    try:
      self.tree = QtWidgets.QTreeWidget()
    except NameError as e:
      print(f'[idamagic] nameerror {e}')
      return
    self.tree.setHeaderLabels(("Names",))
    self.tree.setColumnWidth(0, 100)

    if self.d is None:
      self.d, self.s = get_source_strings(4, [0, 1])
    d = self.d

    # Create layout
    layout = QtWidgets.QVBoxLayout()
    layout.addWidget(self.tree)
    self.populate_tree(d)

    # Populate PluginForm
    self.parent.setLayout(layout)

  def Show(self, title, d = None):
    self.d = d
    return PluginForm.Show(self, title, options = PluginForm.WOPN_PERSIST)

#-------------------------------------------------------------------------------
def basename(path):
  pos1 = path[::-1].find("\\")
  pos2 = path[::-1].find("/")

  if pos1 == -1: pos1 = len(path)
  if pos2 == -1: pos2 = len(path)
  pos = min(pos1, pos2)

  return path[len(path)-pos:]

#-------------------------------------------------------------------------------
class command_handler_t(ida_kernwin.action_handler_t):
  def __init__(self, obj, cmd_id, num_args = 1):
    self.obj = obj
    self.cmd_id = cmd_id
    self.num_args = num_args
    ida_kernwin.action_handler_t.__init__(self)

  def activate(self, ctx):
    if self.num_args == 1:
      return self.obj.OnCommand(self.cmd_id)
    return self.obj.OnCommand(self.obj, self.cmd_id)

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

#-------------------------------------------------------------------------------
class CIDAMagicStringsChooser(Choose):
  def __init__(self, title, columns, options):
    Choose.__init__(self, title, columns, options)
    self.actions = []

  def AddCommand(self, menu_name, shortcut=None):
    action_name = "IDAMagicStrings:%s" % menu_name.replace(" ", "")
    self.actions.append([len(self.actions), action_name, menu_name, shortcut])
    return len(self.actions)-1

  def OnPopup(self, form, popup_handle):
    for num, action_name, menu_name, shortcut in self.actions:
      handler = command_handler_t(self, num, 2)
      desc = ida_kernwin.action_desc_t(action_name, menu_name, handler, shortcut)
      ida_kernwin.attach_dynamic_action_to_popup(form, popup_handle, desc)

#-------------------------------------------------------------------------------
class CSourceFilesChooser(CIDAMagicStringsChooser):
  def __init__(self, title):
    columns = [ ["Line", 4], ["Full path", 20], ["Filename", 15], ["EA", 16], ["Function Name", 18], ["String data", 40], ]
    CIDAMagicStringsChooser.__init__(self, title, columns, Choose.CH_MULTI)
    self.n = 0
    self.icon = -1
    self.selcount = 0
    self.modal = False
    self.items = []
    self.selected_items = []

    d, s = get_source_strings(4, [0,1])
    keys = list(d.keys())
    keys.sort()

    i = 0
    for key in keys:
      for ea, name, str_data in d[key]:
        line = ["%03d" % i, key, basename(key), "0x%08x" % ea, name, str_data]
        self.items.append(line)
        i += 1

    self.d = d
    self.s = s

  def show(self):
    ret = self.Show(False)
    if ret < 0:
      return False

    self.cmd_all          = self.AddCommand("Rename all to filename_EA")
    self.cmd_all_sub      = self.AddCommand("Rename all sub_* to filename_EA")
    self.cmd_selected     = self.AddCommand("Rename selected to filename_EA")
    self.cmd_selected_sub = self.AddCommand("Rename selected sub_* to filename_EA")
    return self.d

  def OnCommand(self, n, cmd_id):
    # Aditional right-click-menu commands handles
    if cmd_id == self.cmd_all:
      l = list(range(len(self.items)))
    elif cmd_id == self.cmd_all_sub:
      l = []
      for i, item in enumerate(self.items):
        if item[4] is not None and item[4].startswith("sub_"):
          l.append(i)
    elif cmd_id == self.cmd_selected:
      l = list(self.selected_items)
    elif cmd_id == self.cmd_selected_sub:
      l = []
      for i, item in enumerate(self.items):
        if item[4].startswith("sub_"):
          if i in self.selected_items:
            l.append(i)

    self.rename_items(l)

  def rename_items(self, items):
    for i in items:
      item = self.items[i]
      ea = int(item[3], 16)
      candidate, _ = os.path.splitext(item[2])
      candidate = candidate
      name = "%s_%08x" % (candidate, ea)
      func = idaapi.get_func(ea)
      if func is not None:
        ea = func.start_ea
        set_name(ea, name, SN_CHECK)
      else:
        print(f'[idamagic] {self.title} cannot rename {ea} {name} snc:{SN_CHECK}' )

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

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

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

  def OnRefresh(self, n):
    return n

  def OnSelectLine(self, n):
    self.selcount += 1
    row = self.items[n[0]]
    ea = int(row[3], 16)
    if is_mapped(ea):
      jumpto(ea)

  def OnSelectionChange(self, sel_list):
    self.selected_items = sel_list

#-------------------------------------------------------------------------------
class CCandidateFunctionNames(CIDAMagicStringsChooser):
  def __init__(self, title, final_list):
    columns = [ ["Line", 4], ["EA", 16], ["Function Name", 25], ["Candidate", 25], ["FP?", 2], ["Strings", 50], ]
    CIDAMagicStringsChooser.__init__(self, title, columns, Choose.CH_MULTI)
    self.n = 0
    self.icon = -1
    self.selcount = 0
    self.modal = False
    self.items = []
    self.selected_items = []

    idx = 0
    for item in final_list:
      bin_func  = item[1]
      candidate = item[2]
      seems_false = str(int(self.looks_false(bin_func, candidate)))
      line = ["%03d" % idx, "0x%08x" % item[0], item[1], item[2], seems_false, ", ".join(item[3]) ]
      self.items.append(line)
      idx += 1

    self.items = sorted(self.items, key=lambda x: x[4])

  def show(self):
    ret = self.Show(False)
    if ret < 0:
      return False

    self.cmd_rename_all      = self.AddCommand("Rename all functions")
    self.cmd_rename_sub      = self.AddCommand("Rename all sub_* functions")
    self.cmd_rename_selected = self.AddCommand("Rename selected function(s)")
    self.cmd_rename_sub_sel  = self.AddCommand("Rename selected sub_* function(s)")

  def OnCommand(self, n, cmd_id):
    # Aditional right-click-menu commands handles
    if cmd_id == self.cmd_rename_all:
      l = list(range(len(self.items)))
    elif cmd_id == self.cmd_rename_selected:
      l = list(self.selected_items)
    elif cmd_id == self.cmd_rename_sub:
      l = []
      for i, item in enumerate(self.items):
        if item[2].startswith("sub_"):
          l.append(i)
    elif cmd_id == self.cmd_rename_sub_sel:
      l = []
      for i, item in enumerate(self.items):
        if item[2].startswith("sub_"):
          if i in self.selected_items:
            l.append(i)
    else:
      raise Exception("Unknown menu command!")

    self.rename_items(l)

  def rename_items(self, items):
    idx = 0
    for i in items:
      idx +=1
      item = self.items[i]
      ea = int(item[1], 16)
      candidate = item[3].replace(':','')
      candidate = candidate.replace(' ','_')
      set_name(ea, candidate, SN_CHECK)
      print(f'[idamagic] {self.title} i:{idx}/{len(items)} rename {item[2]} addr: {item[1]} to {candidate} snc:{SN_CHECK}')

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

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

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

  def OnRefresh(self, n):
    return n

  def OnSelectLine(self, n):
    self.selcount += 1
    row = self.items[n[0]]
    ea = int(row[1], 16)
    if is_mapped(ea):
      jumpto(ea)

  def OnSelectionChange(self, sel_list):
    self.selected_items = sel_list

  def looks_false(self, bin_func, candidate):
    bin_func  = bin_func.lower()
    candidate = candidate.lower()
    if not bin_func.startswith("sub_"):
      if bin_func.find(candidate) == -1 and candidate.find(bin_func) == -1:
        return True
    return False

  def OnGetLineAttr(self, n):
    item = self.items[n]
    bin_func  = item[2]
    candidate = item[3]
    if self.looks_false(bin_func, candidate):
      return [0x026AFD, 0]
    return [0xFFFFFF, 0]

#-------------------------------------------------------------------------------
class CClassXRefsChooser(idaapi.Choose):
  def __init__(self, title, items):
    idaapi.Choose.__init__(self, title, [ ["Address", 8], ["String", 80] ])
    self.items = items

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

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

#-------------------------------------------------------------------------------
def get_string(ea):
  tmp = idc.get_strlit_contents(ea, strtype=0)
  if tmp is None or len(tmp) == 1:
    unicode_tmp = idc.get_strlit_contents(ea, strtype=1)
    if unicode_tmp is not None and len(unicode_tmp) > len(tmp):
      tmp = unicode_tmp

  if tmp is None:
    tmp = ""
  elif type(tmp) != str:
    tmp = tmp.decode("utf-8")
  return tmp

#-------------------------------------------------------------------------------
def classes_handler(item, column_no):
  if item.childCount() == 0:
    ea = item.ea
    if is_mapped(ea):
      jumpto(ea)

#-------------------------------------------------------------------------------
class CClassesTreeViewer(PluginForm):
  def populate_tree(self):
    # Clear previous items
    self.tree.clear()
    self.nodes = {}

    self.classes = sorted(self.classes, key=lambda x: x[1][0])
    for ea, tokens in self.classes:
      for i, node_name in enumerate(tokens):
        full_name = "::".join(tokens[:tokens.index(node_name)+1])
        if full_name not in self.nodes:
          if full_name.find("::") == -1:
            parent = self.tree
          else:
            parent_name = "::".join(tokens[:tokens.index(node_name)])
            try:
              parent = self.nodes[parent_name]
            except Exception as e:
              print(f"[idamagic] Error {e.with_traceback} ea:{ea} parent_name: {parent_name} sys:{str(sys.exc_info()[1])} ")  #, self.nodes, parent_name, str(sys.exc_info()[1]))
            except KeyError as e:
              print(f"[idamagic] KeyError {e.with_traceback} ea:{ea} parent_name: {parent_name} sys:{str(sys.exc_info()[1])} ")  #, self.nodes, parent_name, str(sys.exc_info()[1]))

          node = QtWidgets.QTreeWidgetItem(parent)
          node.setText(0, full_name)
          node.ea = ea
          self.nodes[full_name] = node

    self.tree.itemDoubleClicked.connect(classes_handler)

  def OnCreate(self, form):
    # Get parent widget
    self.parent = idaapi.PluginForm.FormToPyQtWidget(form)

    # Create tree control
    self.tree = QtWidgets.QTreeWidget()
    self.tree.setHeaderLabels(("Classes",))
    self.tree.setColumnWidth(0, 100)

    # Create layout
    layout = QtWidgets.QVBoxLayout()
    layout.addWidget(self.tree)
    self.populate_tree()

    # Populate PluginForm
    self.parent.setLayout(layout)

  def Show(self, title, classes):
    self.classes = classes
    return PluginForm.Show(self, title, options = PluginForm.WOPN_PERSIST)

#-------------------------------------------------------------------------------
class CClassesGraph(idaapi.GraphViewer):
  def __init__(self, title, classes, final_list):
    idaapi.GraphViewer.__init__(self, title)
    self.selected = None
    self.classes = classes
    self.final_list = final_list
    self.nodes = {}
    self.nodes_ea = {}
    self.graph = {}

    self.last_cmd = 0

    dones = set()
    for ea, tokens in self.classes:
      refs = DataRefsTo(ea)
      refs_funcs = set()
      for ref in refs:
        func = idaapi.get_func(ref)
        if func is not None:
          refs_funcs.add(func.start_ea)

      if len(refs_funcs) == 1:
        func_ea = list(refs_funcs)[0]
        if func_ea in dones:
          continue
        dones.add(func_ea)

        func_name = get_func_name(func_ea)
        tmp = demangle_name(func_name, INF_SHORT_DN)
        if tmp is not None:
          func_name = tmp

        element = [func_ea, func_name, "::".join(tokens), [get_string(ea)]]
        self.final_list.append(element)

  def OnRefresh(self):
    self.Clear()
    self.graph = {}
    for ea, tokens in self.classes:
      for node_name in tokens:
        full_name = "::".join(tokens[:tokens.index(node_name)+1])
        if full_name not in self.nodes:
          node_id = self.AddNode(node_name)
          self.nodes[full_name] = node_id
          self.graph[node_id] = []
        else:
          node_id = self.nodes[full_name]

        try:
          self.nodes_ea[node_id].add(ea)
        except KeyError:
          self.nodes_ea[node_id] = set([ea])

        parent_name = "::".join(tokens[:tokens.index(node_name)])
        if parent_name != "" and parent_name in self.nodes:
          parent_id = self.nodes[parent_name]
          try:
            self.AddEdge(parent_id, node_id)
          except AssertionError as e:
            print(f'[idamagic] onrefresh AssertionError {e} parent:{parent_name} pid:{parent_id} nid:{node_id}')
          except Exception as e:
            print(f'[idamagic] onrefresh ERR {e} parent:{parent_name} pid:{parent_id} nid:{node_id}')
          self.graph[parent_id].append(node_id)

    return True

  def OnGetText(self, node_id):
    ogt = ''
    try:
      ogt = str(self[node_id])
    except KeyError as e:
      print(f'[idamagic] {self} ogt keyerror {e} nodeid={node_id}')
    return ogt

  def OnDblClick(self, node_id):
    eas = self.nodes_ea[node_id]
    if len(eas) == 1:
      jumpto(list(eas)[0])
    else:
      items = []
      for ea in eas:
        func = idaapi.get_func(ea)
        if func is None:
          s = idc.get_strlit_contents(ea)
          s = s.decode("utf-8")
          if s is not None and s.find(str(self[node_id])) == -1:
            s = idc.get_strlit_contents(ea, strtype=1)
          else:
            s = GetDisasm(ea)
        else:
          s = get_func_name(func.start_ea)

        items.append(["0x%08x" % ea, repr(s)])

      chooser = CClassXRefsChooser("XRefs to %s" % str(self[node_id]), items)
      idx = chooser.Show(1)
      if idx > -1:
        jumpto(list(eas)[idx])

  def OnCommand(self, cmd_id):
    if self.cmd_dot == cmd_id:
      fname = ask_file(1, "*.dot", "Dot file name")
      if fname:
        f = open(fname, "w")
        buf = 'digraph G {\n graph [overlap=scale]; node [fontname=Courier]; \n\n'
        for n in self.graph:
          name = str(self[n])
          buf += ' a%s [shape=box, label = "%s", color="blue"]\n' % (n, name)
        buf += '\n'

        dones = set()
        for node_id in self.graph:
          for child_id in self.graph[node_id]:
            s = str([node_id, child_id])
            if s in dones:
              continue
            dones.add(s)
            buf += " a%s -> a%s [style = bold]\n" % (node_id, child_id)

        buf += '\n'
        buf += '}'
        f.write(buf)
        f.close()
    elif self.cmd_gml == cmd_id:
      fname = ask_file(1, "*.gml", "GML file name")
      if fname:
        f = open(fname, "w")
        buf = 'graph [ \n'
        for n in self.graph:
          name = str(self[n])
          buf += 'node [ id %s \n label "%s"\n fill "blue" \n type "oval"\n LabelGraphics [ type "text" ] ] \n' % (n, name)
        buf += '\n'

        dones = set()
        for node_id in self.graph:
          for child_id in self.graph[node_id]:
            s = str([node_id, child_id])
            if s in dones:
              continue
            dones.add(s)
            buf += " edge [ source %s \n target %s ]\n" % (node_id, child_id)

        buf += '\n'
        buf += ']'
        f.write(buf)
        f.close()

  def OnPopup(self, form, popup_handle):
    self.cmd_dot = 0
    cmd_handler = command_handler_t(self, self.cmd_dot)
    desc = ida_kernwin.action_desc_t("IDAMagicStrings:GraphvizExport", "Export to Graphviz",
                                     cmd_handler, "F2")
    ida_kernwin.attach_dynamic_action_to_popup(form, popup_handle, desc)

    self.cmd_gml = 1
    cmd_handler = command_handler_t(self, self.cmd_gml)
    desc = ida_kernwin.action_desc_t("IDAMagicStrings:GmlExport","Export to GML",
                                     cmd_handler, "F3")
    ida_kernwin.attach_dynamic_action_to_popup(form, popup_handle, desc)

  def OnClick(self, item):
    self.selected = item
    return True

  def Show(self):
    if not idaapi.GraphViewer.Show(self):
      return False
    return True

#-------------------------------------------------------------------------------
def show_tree(d = None):
  tree_frm = CBaseTreeViewer()
  tree_frm.Show(PROGRAM_NAME + ": Source code tree", d)

#-------------------------------------------------------------------------------
def seems_function_name(candidate):
  if len(candidate) >= 6 and candidate.lower() not in NOT_FUNCTION_NAMES:
    if candidate.upper() != candidate:
      return True
  return False

#-------------------------------------------------------------------------------
class CFakeString:
  def __init__(self, ea, s):
    self.ea = ea
    self.s = s

  def __str__(self):
    return str(self.s)

  def __repr__(self):
    return self.__str__()

#-------------------------------------------------------------------------------
def find_function_names(strings_list):
  rarity = {}
  func_names = {}
  raw_func_strings = {}
  class_objects = []

  class_tmp_names = []
  for ea, name in Names():
    func = idaapi.get_func(ea)
    if func is None:
      continue

    true_name = name
    if name.find("::") == -1:
      name = demangle_name(name, INF_SHORT_DN)
      if name is not None and name != "" and name.find("::") > -1:
        true_name = name

    if true_name.find("::") > -1:
      s = CFakeString(ea, true_name)
      class_tmp_names.append(s)

  class_tmp_names.extend(strings_list)
  for s in class_tmp_names:
    # Find class members
    class_ret = re.findall(CLASS_NAMES_REGEXP, str(s), re.IGNORECASE)
    if len(class_ret) > 0:
      for element in class_ret:
        candidate = element[0]
        if candidate.find("::") > 0:
          tokens = candidate.split("::")
          if tokens not in class_objects:
            class_objects.append([s.ea, tokens])

    # Find just function names
    ret = re.findall(FUNCTION_NAMES_REGEXP, str(s), re.IGNORECASE)
    if len(ret) > 0:
      candidate = ret[0][0]
      if seems_function_name(candidate):
        ea = s.ea
        refs = DataRefsTo(ea)
        found = False
        for ref in refs:
          func = idaapi.get_func(ref)
          if func is not None:
            found = True
            key = func.start_ea

            if has_nltk:
              if candidate not in FOUND_TOKENS:
                continue

              found = False
              for tkn_type in TOKEN_TYPES:
                if tkn_type in FOUND_TOKENS[candidate]:
                  found = True
                  break

              if not found:
                continue

            try:
              rarity[candidate].add(key)
            except KeyError:
              rarity[candidate] = set([key])

            try:
              func_names[key].add(candidate)
            except KeyError:
              func_names[key] = set([candidate])

            try:
              raw_func_strings[key].add(str(s))
            except:
              raw_func_strings[key] = set([str(s)])

  return func_names, raw_func_strings, rarity, class_objects

#-------------------------------------------------------------------------------
def show_function_names(strings_list):
  found_func_names = find_function_names(strings_list)
  func_names, raw_func_strings, rarity, classes = found_func_names
  final_list = []
  for key in func_names:
    candidates = set()
    for candidate in func_names[key]:
      if len(rarity[candidate]) == 1:
        candidates.add(candidate)

    if len(candidates) == 1:
      raw_strings = list(raw_func_strings[key])
      raw_strings = list(map(repr, raw_strings))

      func_name = get_func_name(key)

      tmp_func_name = demangle_name(func_name, INF_SHORT_DN)
      if tmp_func_name is not None:
        func_name = tmp_func_name

      candidate = candidate.replace('::', '_')
      final_list.append([key, func_name, list(candidates)[0], raw_strings])

  if len(classes) > 0:
    class_graph = CClassesGraph(PROGRAM_NAME + ": Classes Hierarchy", classes, final_list)
    class_graph.Show()

    class_tree = CClassesTreeViewer()
    class_tree.Show(PROGRAM_NAME + ": Classes Tree", classes)

    final_list = class_graph.final_list

  if len(final_list) > 0:
    cfn = CCandidateFunctionNames(PROGRAM_NAME + ": Candidate Function Names", final_list)
    cfn.show()

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

class magicstrings(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = "IDA Magic Strings plugin"

    help = "This is magicstrings help"
    wanted_name = "Magic Strings"
    wanted_hotkey = ""

    def init(self):
        return idaapi.PLUGIN_OK

    def run(self, arg):
        main()

    def term(self):
        pass


def PLUGIN_ENTRY():
    return magicstrings()

def main():
  ch = CSourceFilesChooser(PROGRAM_NAME + ": Source code files")
  if len(ch.items) > 0:
    ch.show()

  d = ch.d
  if len(d) > 0:
    show_tree(d)

  show_function_names(ch.s)

if __name__ == "__main__":
  main()


repository: idapython/src
file: examples/decompiler/vds6.py
description: IDAPython project for Hex-Ray's IDA Pro
wanted_name: Hex-Rays space remover (IDAPython)
comment: Sample plugin6 for Hex-Rays decompiler
created_at: 2015-03-14 20:09:27+00:00
pushed_at: 2025-03-13 15:18:11+00:00
url:
https://github.com/idapython/src/blob/af978e79e68af701c45bab401afce5145aa4b80b/examples/decompiler/vds6.py

content:
"""
summary: superficially modify the decompilation output

description:
  modifies the decompilation output in a superficial manner,
  by removing some white spaces

  Note: this is rather crude, not quite "pythonic" code.

level: intermediate
"""

import idautils
import idc
import ida_idaapi
import ida_hexrays
import ida_lines

do_dbg = False

def dbg(msg):
    if do_dbg:
        print(msg)
    pass

def is_cident_char(c):
    return c.isalnum() or c == '_'

def my_tag_skipcodes(l, storage):
    n = ida_lines.tag_skipcodes(l)
    if n > 0:
        dbg("Skipping %d chars ('%s')" % (n, l[0:n]))
        storage.append(l[0:n])
    return l[n:]

def remove_spaces(sl):

    dbg("*" * 80)
    l = sl.line

    global do_dbg
    # do_dbg = l.find("const char *") > -1

    out = []

    def push(c):
        dbg("Appending '%s'" % c)
        out.append(c)

    # skip initial spaces, do not compress them
    def eat_spaces(l):
        while True:
            l = my_tag_skipcodes(l, out)
            if not l:
                break
            c = l[0]
            if not c.isspace():
                break
            push(c)
            l = l[1:]
        return l
    l = eat_spaces(l)

    # remove all spaces except in string and char constants
    delim = None # if not None, then we are skipping until 'delim'
    last = None # last seen character

    while True:
        dbg("-" * 60)
        dbg("l: '%s'" % l)
        dbg("d: '%s'" % delim)
        dbg("out: '%s'" % out)

        # go until comments
        l = my_tag_skipcodes(l, out)
        if l.startswith("//"):
            push(l)
            break
        nchars = ida_lines.tag_advance(l, 1)
        push(l[0:nchars])
        l = l[nchars:]
        l = my_tag_skipcodes(l, out)
        if not l:
            break
        c = l[0]
        dbg("c: '%s', last: '%s', l: '%s'" % (c, last, l))
        if delim:
            # we're inside a literal.
            if c == delim:
                delim = None # literal ended
        elif c == '"' or c == "'":
            delim = c # string/char literal started
        elif c.isspace():
            end = l.lstrip()
            nptr = my_tag_skipcodes(end, out)
            dbg("end: '%s', nptr: '%s'" % (end, nptr))
            # do not concatenate idents
            if not is_cident_char(last) or not is_cident_char(nptr[0]):
                l = end
                c = l[0] if l else ''
                dbg("new l: '%s'" % l)
        last = l[0] if l else ''

    sl.line = "".join(out)


class vds6_hooks_t(ida_hexrays.Hexrays_Hooks):
    def func_printed(self, cfunc):
        for sl in cfunc.get_pseudocode():
            remove_spaces(sl);
        return 0

# a plugin interface, boilerplate code
class my_plugin_t(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_HIDE
    wanted_name = "Hex-Rays space remover (IDAPython)"
    wanted_hotkey = ""
    comment = "Sample plugin6 for Hex-Rays decompiler"
    help = ""
    def init(self):
        if ida_hexrays.init_hexrays_plugin():
            self.vds6_hooks = vds6_hooks_t()
            self.vds6_hooks.hook()
            return ida_idaapi.PLUGIN_KEEP # keep us in the memory
    def term(self):
        self.vds6_hooks.unhook()
    def run(self, arg):
        pass

def PLUGIN_ENTRY():
    return my_plugin_t()


repository: lj94093/IDAAndroidBreakpoint
file: AndroidBreakpoint.py
description: IDA plugin aid to set android so breakpoint
wanted_name: AndroidBreakpoint
comment: usual bp in android
created_at: 2017-12-08 13:11:06+00:00
pushed_at: 2017-12-09 11:33:21+00:00
url:
https://github.com/lj94093/IDAAndroidBreakpoint/blob/2234343b48b3ff0bfcda96253d419bf4c8064303/AndroidBreakpoint.py

content:
import idaapi
import idc
import struct

def get_ushort(file,offset):
    file.seek(offset)
    return struct.unpack("H",file.read(2))[0]

def get_ulong(file,offset):
    file.seek(offset)
    return struct.unpack("L",file.read(4))[0]


def static_breakpoint(target_type):
    try:
        base = idc.FirstSeg()
        if base == idc.BADADDR:
            print("can't find first seg")
            return
        filepath = idaapi.get_input_file_path()
        print("filename:", filepath)
        idbFile = open(filepath, "rb")
        e_phoff = get_ushort(idbFile, 0x1C)
        e_phentsize = get_ushort(idbFile, 0x2A)
        e_phnum = get_ushort(idbFile, 0x2C)
        print(e_phoff, e_phentsize, e_phnum)
        dynamic = -1
        for i in xrange(e_phnum):
            idbFile.seek(e_phoff + e_phentsize * i)
            p_type = get_ulong(idbFile, e_phoff + e_phentsize * i)
            # define PT_DYNAMIC 2
            if p_type == 2:
                dynamic = get_ulong(idbFile, e_phoff + e_phentsize * i + 4)
                break
        if dynamic == -1:
            print("\tcan't find dynamic")
            return
        dynamic_type = get_ulong(idbFile, dynamic)
        init = -1

        while dynamic_type != 0:

            dynamic += 8
            dynamic_type = get_ulong(idbFile, dynamic)

            if dynamic_type == target_type:
                init = get_ulong(idbFile, dynamic + 4)
                break
        if init == -1:
            print("\tcan't find init")
            return
        idc.AddBpt(base + init)
        print('\tbp the init on the %x' % (base + init))
        print("\nbp init finished")
    except BaseException, e:
        print(e)


#linker call init and init_array functions
# .text:00002924 14 49                       LDR             R1, =(aLinker - 0x2930)
# .text:00002926 04 20                       MOVS            R0, #4
# .text:00002928 23 46                       MOV             R3, R4
# .text:0000292A 14 4A                       LDR             R2, =(aSCallingCons_0 - 0x2932)
# .text:0000292C 79 44                       ADD             R1, PC  ; "linker"
# .text:0000292E 7A 44                       ADD             R2, PC  ; "\"%s\": calling constructors"
# .text:00002930 01 F0 E6 FE                 BL              sub_4700
# .text:00002934
# .text:00002934             loc_2934                                ; CODE XREF: sub_2884+52j
# .text:00002934 12 49                       LDR             R1, =(aDt_init - 0x2940)
# .text:00002936 20 46                       MOV             R0, R4
# .text:00002938 D4 F8 F0 20                 LDR.W           R2, [R4,#0xF0]
# .text:0000293C 79 44                       ADD             R1, PC  ; "DT_INIT"
# .text:0000293E FF F7 DF FE                 BL              sub_2700
# .text:00002942 10 49                       LDR             R1, =(aDt_init_array - 0x2956)
# .text:00002944 00 20                       MOVS            R0, #0
# .text:00002946 00 90                       STR             R0, [SP,#0x28+var_28]
# .text:00002948 20 46                       MOV             R0, R4
# .text:0000294A D4 F8 E0 20                 LDR.W           R2, [R4,#0xE0]
# .text:0000294E D4 F8 E4 30                 LDR.W           R3, [R4,#0xE4]
# .text:00002952 79 44                       ADD             R1, PC  ; "DT_INIT_ARRAY"
# .text:00002954 FF F7 0E FF                 BL              sub_2774
# .text:00002958
# .text:00002958             locret_2958                             ; CODE XREF: sub_2884+Cj
# .text:00002958 BD E8 FE 83                 POP.W           {R1-R9,PC}
def dynamic_breakpoint(targe_type):
    has_linker = False
    module_base = idc.GetFirstModule()
    while module_base != None:
        module_name = idc.GetModuleName(module_base)
        if module_name.find('linker') >= 0:
            has_linker = True
            break

        module_base = idc.GetNextModule(module_base)

    if has_linker == False:
        print '[*]unable to find linker module base'
        return

    module_size = idc.GetModuleSize(module_base)
    print '[*]found linker base=>0x%08X, Size=0x%08X' % (module_base, module_size)

    print("\t[-]begin to search DT_INIT")
    init_func_ea = 0
    init_array_ea = 0
    # bytecode=b'\x53\x1e\x73\xb5\x03\x33\x06\x46\x0d\x46\x14\x46\x24\xd8\x13\x48\x78\x44\x01\x68\x01\x29'
    bytecode =[0x14,0x49,0x04,0x20,0x23,0x46,0x14,0x4A,0x79,0x44,0x7A,0x44]
    findcode=True
    for ea_offset in range(module_base, module_base + module_size):
        findcode = True
        for i in xrange(len(bytecode)):
            if idaapi.get_byte(ea_offset+i)!=bytecode[i]:
                findcode=False
                break
        if(findcode==True):
            init_func_ea=ea_offset+0x1A
            init_array_ea=ea_offset+0x30
            break
    if(findcode==False):
        print("can't find bytecode")
        return
    print "\t[-]found INIT=>0x%08X INIT_ARRAY=>0x%08X" % (init_func_ea, init_array_ea)
    print("\t[-]try set breakpoint there")
    if targe_type==12:
        idc.AddBpt(init_func_ea)
    if targe_type==25:
        idc.AddBpt(init_array_ea)
    print("[*]script finish")

class bpInit(idaapi.action_handler_t):
    name = "bpInit"
    lable = "bpInit"
    hotkey = None

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

    def activate(self, ctx):
        print("\n.init handler activate")
        idainfo = idaapi.get_inf_structure()
        if idainfo.filetype != idaapi.f_ELF:
            print("\tidb isn't elf,skip")
            return
        if idaapi.is_debugger_on():
            print("\tdebugger is on,prepare to bp by linker.")
            dynamic_breakpoint(12)
        else:
            print("\tdebugger isn't on ,prepare to bp by search the elf program header")
            # define DT_INIT 12
            static_breakpoint(12)

    def update(self, ctx):
        print(".init handler update")
        return idaapi.AST_ENABLE_ALWAYS

class bpInitArray(idaapi.action_handler_t):
    name = "bpInitArray"
    lable = "bpInitArray"
    hotkey = None
    def __init__(self):
        idaapi.action_handler_t.__init__(self)

    def activate(self, ctx):
        print("init array handler activate")
        idainfo = idaapi.get_inf_structure()
        if idainfo.filetype != idaapi.f_ELF:
            print("\tidb isn't elf,skip")
            return
        if idaapi.is_debugger_on():
            print("\tdebugger is on,prepare to bp by linker.")
            dynamic_breakpoint(25)
        else:
            print("\tdebugger isn't on ,prepare to bp by search the elf program header")
            # define DT_INIT_ARRAY	25
            static_breakpoint(25)

    def update(self, ctx):
        print("init array handler update")
        return idaapi.AST_ENABLE_ALWAYS

def register(action,*args):
    idaapi.register_action(
        idaapi.action_desc_t(
            action.name,
            action.lable,
            action(*args),
            action.hotkey,
            "set breakpoint on the .init function",
        )
    )

class myIdaPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_KEEP
    wanted_name = "AndroidBreakpoint"
    wanted_hotkey = "Alt-F9"
    comment = "usual bp in android"
    help = "Something helpful"

    def init(self):
        print("init func")
        try:
            # result=idaapi.register_and_attach_to_menu("Options/path","name","label","",0,self.run(0),None)
            register(bpInit)
            register(bpInitArray)
            idaapi.attach_action_to_menu("Edit/AndroidBreakpoint/",bpInit.name,idaapi.SETMENU_APP)
            idaapi.attach_action_to_menu("Edit/AndroidBreakpoint/", bpInitArray.name, idaapi.SETMENU_APP)
        except BaseException, e:
            print(e)

        idainfo=idaapi.get_inf_structure()
        if idainfo.filetype!=idaapi.f_ELF:
            print("idb isn't elf,skip")
            return idaapi.PLUGIN_SKIP
        return idaapi.PLUGIN_KEEP

    def term(self):
        pass

    def run(self, arg):
        try:
            print("AndroidBreakpoint run")
        except BaseException, e:
            print e

def PLUGIN_ENTRY():

    return myIdaPlugin()


repository: binsync/binsync
file: binsync/binsync_plugin.py
description: A reversing plugin for cross-decompiler collaboration, built on git.
created_at: 2019-07-09 15:53:32+00:00
pushed_at: 2025-04-18 05:17:44+00:00
url:
https://github.com/binsync/binsync/blob/38f8d25713063c0d4f9be9e780b448b7d53c5de5/binsync/binsync_plugin.py

content:
# A cross-decompiler collaboration plugin
# @author BinSync Team
# @category Collaboration
# @menupath Tools.BinSync.Start UI...

plugin_command = "binsync -s ghidra"


def create_plugin(*args, **kwargs):
    # REPLACE_ME this import with an import of your plugin's create_plugin function
    from binsync import create_plugin as _create_plugin
    return _create_plugin(*args, **kwargs)

# =============================================================================
# LibBS generic plugin loader (don't touch things below this)
# =============================================================================

import sys
# Python 2 has special requirements for Ghidra, which forces us to use a different entry point
# and scope for defining plugin entry points.
# The Python 3 side has been edited since currently every supported decompiler must import create_plugin
if sys.version[0] == "2":
    # Do Ghidra Py2 entry point
    import subprocess
    from libbs_vendored.ghidra_bridge_server import GhidraBridgeServer

    GhidraBridgeServer.run_server(background=True)
    process = subprocess.Popen(plugin_command.split(" "))
    if process.poll() is not None:
        raise RuntimeError("Failed to run the Python3 backed. It's likely Python3 is not in your Path inside Ghidra.")


def PLUGIN_ENTRY(*args, **kwargs):
    """
    This is the entry point for IDA to load the plugin.
    """
    return create_plugin(*args, **kwargs)


repository: hasherezade/ida_ifl
file: ifl.py
description: IFL - Interactive Functions List (plugin for IDA Pro)
url: https://github.com/hasherezade/ida_ifl/blob/615acaed1213cb96dc8887c54b00e5a511fea9e1/ifl.py
comment: Interactive Functions List
created_at: 2015-09-22 23:45:51+00:00
pushed_at: 2025-02-23 20:49:11+00:00
content:
#!/usr/bin/env python3
#
# IFL - Interactive Functions List
#
# how to install: copy the script into plugins directory, i.e: C:\Program Files\IDA <version>\plugins
# then:
# run from IDA menu: View -> PLUGIN_NAME
# or press: PLUGIN_HOTKEY
#
"""
CC-BY: hasherezade, run via IDA Pro >= 8.0
"""
__VERSION__ = "1.5.2"
__AUTHOR__ = "hasherezade"

PLUGIN_NAME = "IFL - Interactive Functions List"
PLUGIN_HOTKEY = "Ctrl-Alt-F"
import string
from typing import Any, List, Optional, Tuple, Union

import ida_bytes
import ida_kernwin
import idaapi  # type: ignore
import idc  # type: ignore
from idaapi import PluginForm  # type: ignore
from idaapi import (
    BADADDR,
    SN_NOWARN,
    jumpto,
    next_addr,
    o_void,
    prev_addr,
    print_insn_mnem,
    set_cmt,
    set_name,
)
from idautils import Functions, XrefsFrom, XrefsTo  # type: ignore
from idc import (
    CIC_ITEM,
    FUNCATTR_END,
    FUNCATTR_FRAME,
    FUNCATTR_START,
    INF_SHORT_DN,
    GetDisasm,
    demangle_name,
    get_func_attr,
    get_func_name,
    get_inf_attr,
    get_operand_type,
    get_operand_value,
    get_type,
    set_color,
)
from PyQt5 import QtCore, QtGui, QtWidgets  # type: ignore
from PyQt5.QtCore import QObject, pyqtSignal  # type: ignore

VERSION_INFO = (
    f"IFL v{__VERSION__} - check for updates: https://github.com/hasherezade/ida_ifl"
)

transp_l = 230
light_theme = [
    QtGui.QColor(173, 216, 230, transp_l),
    QtGui.QColor(255, 165, 0, transp_l),
    QtGui.QColor(240, 230, 140, transp_l),
]
transp_d = 70
dark_theme = [
    QtGui.QColor(173, 216, 240, transp_d),
    QtGui.QColor(255, 0, 255, transp_d),
    QtGui.QColor(255, 130, 130, transp_d),
]

COLOR_HILIGHT_FUNC = 0xFFDDBB  # BBGGRR
COLOR_HILIGHT_REFTO = 0xBBFFBB
COLOR_HILIGHT_REFFROM = 0xDDBBFF

IS_ALTERNATE_ROW = False

# --------------------------------------------------------------------------
# custom functions:
# --------------------------------------------------------------------------

# Theme


def get_bg_color():
    w = ida_kernwin.get_current_widget()
    if not w:
        return None
    widget = ida_kernwin.PluginForm.FormToPyQtWidget(w)
    if not widget:
        return None
    color = widget.palette().color(QtGui.QPalette.Background)
    return color


def is_darker(color1, color2):
    if QtGui.QColor(color2).lightness() > QtGui.QColor(color1).lightness():
        return False
    return True


def color_to_val(color):
    return (((color.red() << 8) | color.green()) << 8) | color.blue()


def get_theme():
    bgcolor = get_bg_color()
    if bgcolor is None:
        return None
    if is_darker(bgcolor, COLOR_HILIGHT_FUNC):
        return light_theme
    return dark_theme


# Addressing


def rva_to_va(rva: int) -> int:
    base = idaapi.get_imagebase()
    return rva + base


def va_to_rva(va: int) -> int:
    base = idaapi.get_imagebase()
    return va - base


def _is_hex_str(s):
    # Check if the string starts with "0x" or is a valid hex string
    if s.startswith("0x"):
        s = s[2:]
    return all(c in string.hexdigits for c in s)


# Functions and args


def function_at(ea: int) -> Optional[int]:
    return next(Functions(ea), None)


def parse_function_args(ea: int) -> str:
    arguments = []
    frame = idc.get_func_attr(ea, FUNCATTR_FRAME)

    if frame is None:
        return ""
    
    if 760 <= idaapi.IDA_SDK_VERSION < 900:

        class frame_members_iterator:
            def __init__(self, f, max_count=10000):
                self.frame = f
                self.max_count = max_count
                self.start = idc.get_first_member(f)
                self.end = idc.get_last_member(f)
                self.count = 0

            def __iter__(self):
                while self.start <= self.end and self.count <= self.max_count:
                    size = idc.get_member_size(self.frame, self.start)
                    self.count += 1

                    if size is None:
                        self.start += 1
                        continue  # Skip invalid members

                    name = idc.get_member_name(self.frame, self.start)
                    self.start += size

                    yield name
        skip_names = [" r", " s"]
        udt_data_iter = frame_members_iterator(frame)
    else:
        import operator

        import ida_typeinf

        tif = ida_typeinf.tinfo_t()
        tif.get_type_by_tid(ea)
        udt_data = ida_typeinf.udt_type_data_t()
        if not tif.get_udt_details(udt_data):
            return ""
        skip_names = ["__return_address", "__saved_registers"]
        udt_data_iter = map(operator.attrgetter("name"), udt_data)

    arguments = [arg for arg in udt_data_iter if arg not in skip_names]
    if len(arguments) == 0:
        args_str = "void"
    else:
        args_str = ", ".join(arguments)
    return f"({args_str})"


def parse_function_type(ea: int, end: Optional[int] = None) -> str:
    frame = idc.get_func_attr(ea, FUNCATTR_FRAME)
    if frame is None:
        return ""
    if end is None:  # try to find end
        func = function_at(ea)
        if not func:
            return "?"
        end = prev_addr(idc.get_func_attr(func, FUNCATTR_END))
    end_addr = end
    mnem = GetDisasm(end_addr)

    if "ret" not in mnem:
        # it's not a real end, get instruction before...
        end_addr = prev_addr(end)
        if end_addr == BADADDR:
            # cannot get the real end
            return ""
        mnem = GetDisasm(end_addr)

    if "ret" not in mnem:
        # cannot get the real end
        return ""

    op = get_operand_type(end_addr, 0)
    if op == o_void:
        # retn has NO parameters
        return "__cdecl"
    # retn has parameters
    return "__stdcall"


def _getFunctionType(start: int, end: Optional[int] = None) -> str:
    type = get_type(start)
    if type is None:
        return parse_function_type(start, end)
    args_start = type.find("(")
    if args_start is not None:
        type = type[:args_start]
    return type


def _isFunctionMangled(ea: int) -> bool:
    name = get_func_name(ea)
    disable_mask = get_inf_attr(INF_SHORT_DN)
    if demangle_name(name, disable_mask) is None:
        return False
    return True


def _getFunctionNameAt(ea: int) -> str:
    name = get_func_name(ea)
    disable_mask = get_inf_attr(INF_SHORT_DN)
    demangled_name = demangle_name(name, disable_mask)
    if demangled_name is None:
        return name
    args_start = demangled_name.find("(")
    if args_start is None:
        return demangled_name
    return demangled_name[:args_start]


def _getArgsDescription(ea: int) -> str:
    name = demangle_name(
        get_func_name(ea), get_inf_attr(INF_SHORT_DN)
    )  # get from mangled name
    if not name:
        name = get_type(ea)  # get from type
        if not name:
            return parse_function_args(ea)  # cannot get params from the mangled name
    args_start = name.find("(")
    if args_start is not None and args_start != (-1):
        return name[args_start:]
    return ""


def _getArgsNum(ea: int) -> int:
    args = _getArgsDescription(ea)
    if not args or len(args) == 0:
        return 0
    delimiter = ","
    args_list = args.split(delimiter)
    args_num = 0
    for arg in args_list:
        if arg == "()" or arg == "(void)":
            continue
        args_num += 1
    return args_num


# --------------------------------------------------------------------------
# custom data types:
# --------------------------------------------------------------------------


# Global DataManager
class DataManager(QObject):
    """Keeps track on the changes in data and signalizies them."""

    updateSignal = pyqtSignal()

    def __init__(self, parent=None) -> None:
        QtCore.QObject.__init__(self, parent=parent)
        self.currentRva = BADADDR

    def setFunctionName(self, start: int, func_name: str) -> bool:
        flags = idaapi.SN_NOWARN | idaapi.SN_NOCHECK
        if idc.set_name(start, func_name, flags):
            self.updateSignal.emit()
            return True
        return False

    def setCurrentRva(self, rva: Optional[int]) -> None:
        if rva is None:
            rva = BADADDR
        self.currentRva = rva
        self.updateSignal.emit()

    def refreshData(self) -> None:
        self.updateSignal.emit()


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


class FunctionInfo_t:
    """A class representing a single function's record."""

    def __init__(
        self, start: int, end: int, refs_list, called_list, is_import: bool = False
    ) -> None:
        self.start = start
        self.end = end
        self.args_num = _getArgsNum(start)
        self.type = _getFunctionType(start, end)
        self.is_import = is_import
        self.refs_list = refs_list
        self.called_list = called_list

    def contains(self, addr: int) -> bool:
        """Check if the given address lies inside the function."""
        bgn = self.start
        end = self.end
        # swap if order is opposite:
        if self.start > self.end:
            end = self.start
            bgn = self.end
        if addr >= bgn and addr < end:
            return True
        return False


# --------------------------------------------------------------------------
# custom models:
# --------------------------------------------------------------------------


class TableModel_t(QtCore.QAbstractTableModel):
    """The model for the top view: storing all the functions."""

    COL_START = 0
    COL_END = 1
    COL_NAME = 2
    COL_TYPE = 3
    COL_ARGS = 4
    COL_REFS = 5
    COL_CALLED = 6
    COL_IMPORT = 7
    COL_COUNT = 8
    header_names = [
        "Start",
        "End",
        "Name",
        "Type",
        "Args",
        "Is referred by",
        "Refers to",
        "Imported?",
    ]

    # private:

    def _displayHeader(
        self, orientation: QtCore.Qt.Orientation, col: int
    ) -> Optional[str]:
        if orientation == QtCore.Qt.Vertical:
            return None
        if col == self.COL_START:
            return self.header_names[self.COL_START]
        if col == self.COL_END:
            return self.header_names[self.COL_END]
        if col == self.COL_TYPE:
            return self.header_names[self.COL_TYPE]
        if col == self.COL_ARGS:
            return self.header_names[self.COL_ARGS]
        if col == self.COL_NAME:
            return self.header_names[self.COL_NAME]
        if col == self.COL_REFS:
            return self.header_names[self.COL_REFS]
        if col == self.COL_CALLED:
            return self.header_names[self.COL_CALLED]
        if col == self.COL_IMPORT:
            return self.header_names[self.COL_IMPORT]
        return None

    def _displayData(self, row: int, col: int) -> Optional[Union[int, str]]:
        func_info = self.function_info_list[row]
        if col == self.COL_START:
            return "%08x" % func_info.start
        if col == self.COL_END:
            return "%08x" % func_info.end
        if col == self.COL_TYPE:
            return func_info.type
        if col == self.COL_ARGS:
            return _getArgsDescription(func_info.start)
        if col == self.COL_NAME:
            return _getFunctionNameAt(func_info.start)
        if col == self.COL_REFS:
            return len(func_info.refs_list)
        if col == self.COL_CALLED:
            return len(func_info.called_list)
        if col == self.COL_IMPORT:
            if func_info.is_import:
                return "+"
            return "-"
        return None

    def _displayToolTip(self, row: int, col: int) -> str:
        func_info = self.function_info_list[row]
        if col == self.COL_START or col == self.COL_END:
            return "Double Click to follow"
        if col == self.COL_NAME:
            return "Double Click to edit"
        if col == self.COL_REFS:
            return self._listRefs(func_info.refs_list)
        if col == self.COL_CALLED:
            return self._listRefs(func_info.called_list)
        return ""

    def _displayBackground(self, row: int, col: int) -> Any:
        curr_theme = get_theme()
        if curr_theme is not None:
            self.theme = curr_theme

        func_info = self.function_info_list[row]
        if col == self.COL_START or col == self.COL_END:
            return QtGui.QColor(self.theme[0])
        if col == self.COL_NAME:
            if func_info.is_import:
                return QtGui.QColor(self.theme[1])
            return QtGui.QColor(self.theme[2])
        return None

    def _listRefs(self, refs_list: List[Tuple[int, int]]) -> str:
        str_list = []
        for ea, ea_to in refs_list:
            str = "%08x @ %s" % (ea, _getFunctionNameAt(ea_to))
            str_list.append(str)
        return "\n".join(str_list)

    # public:
    def __init__(self, function_info_list, parent=None, *args) -> None:
        super(TableModel_t, self).__init__()
        self.function_info_list = function_info_list
        self.theme = light_theme

    def isFollowable(self, col: int) -> bool:
        if col == self.COL_START:
            return True
        if col == self.COL_END:
            return True
        return False

    # Qt API
    def rowCount(self, parent) -> int:
        return len(self.function_info_list)

    def columnCount(self, parent) -> int:
        return self.COL_COUNT

    def setData(self, index, content, role) -> bool:
        if not index.isValid():
            return False
        func_info = self.function_info_list[index.row()]
        if index.column() == self.COL_NAME:
            set_name(func_info.start, str(content), SN_NOWARN)
            g_DataManager.refreshData()
        return True

    def data(self, index, role) -> Any:
        if not index.isValid():
            return None
        col = index.column()
        row = index.row()

        func_info = self.function_info_list[row]

        if role == QtCore.Qt.UserRole:
            if col == self.COL_END:
                return func_info.end
            return func_info.start
        elif role == QtCore.Qt.DisplayRole or role == QtCore.Qt.EditRole:
            return self._displayData(row, col)
        elif role == QtCore.Qt.ToolTipRole:
            return self._displayToolTip(row, col)
        elif role == QtCore.Qt.BackgroundColorRole:
            return self._displayBackground(row, col)
        else:
            return None

    def flags(self, index) -> Any:
        if not index.isValid():
            return None
        flags = QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsSelectable
        if index.column() == self.COL_NAME:
            return flags | QtCore.Qt.ItemIsEditable
        return flags

    def headerData(self, section, orientation, role=QtCore.Qt.DisplayRole) -> Any:
        if role == QtCore.Qt.DisplayRole:
            return self._displayHeader(orientation, section)
        else:
            return None


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


class RefsTableModel_t(QtCore.QAbstractTableModel):
    """The model for the bottom view: the references to the functions."""

    COL_NAME = 0
    COL_ADDR = 1
    COL_TOADDR = 2
    COL_COUNT = 3

    # private:
    def _displayHeader(self, orientation, col: int) -> Optional[str]:
        """Retrieves a field description to be displayed in the header."""

        if orientation == QtCore.Qt.Vertical:
            return None
        if col == self.COL_ADDR:
            return "From Address"
        if col == self.COL_TOADDR:
            return "To Address"
        if col == self.COL_NAME:
            return "Foreign Val."
        return None

    def _getTargetAddr(self, row: int) -> int:
        """Retrieves the address from which function was referenced, or to which it references."""

        curr_ref_fromaddr = self.refs_list[row][0]  # fromaddr
        curr_ref_addr = self.refs_list[row][1]  # toaddr
        target_addr = BADADDR
        if self.is_refs_to:
            target_addr = curr_ref_fromaddr
        else:
            target_addr = curr_ref_addr
        return target_addr

    def _getForeignFuncName(self, row: int) -> str:
        """Retrieves a name of the foreign function or the details on the referenced address."""

        # curr_ref_fromaddr = self.refs_list[row][0]  # fromaddr
        # curr_ref_addr = self.refs_list[row][1]  # toaddr

        target_addr = self._getTargetAddr(row)
        if print_insn_mnem(target_addr) != "":
            func_name = _getFunctionNameAt(target_addr)
            if func_name:
                return func_name

        addr_str = "[%08lx]" % target_addr
        # target_name = GetDisasm(target_addr)
        return f"{addr_str} : {GetDisasm(target_addr)}"

    def _displayData(self, row: int, col: int) -> Optional[str]:
        """Retrieves the data to be displayed. appropriately to the row and column."""

        if len(self.refs_list) <= row:
            return None
        curr_ref_fromaddr = self.refs_list[row][0]  # fromaddr
        curr_ref_addr = self.refs_list[row][1]  # toaddr
        if col == self.COL_ADDR:
            return "%08x" % curr_ref_fromaddr
        if col == self.COL_TOADDR:
            return "%08x" % curr_ref_addr
        if col == self.COL_NAME:
            return self._getForeignFuncName(row)
        return None

    def _getAddrToFollow(self, row: int, col: int) -> int:
        """Retrieves the address that can be followed on click."""

        if col == self.COL_ADDR:
            return self.refs_list[row][0]
        if col == self.COL_TOADDR:
            return self.refs_list[row][1]
        return BADADDR

    def _displayBackground(self, row: int, col: int) -> Any:
        """Retrieves a background color appropriate for the data."""
        curr_theme = get_theme()
        if curr_theme is not None:
            self.theme = curr_theme

        if self.isFollowable(col):
            return QtGui.QColor(self.theme[0])
        return None

    # public:
    def __init__(self, function_info_list, is_refs_to=True, parent=None, *args) -> None:
        super(RefsTableModel_t, self).__init__()
        self.function_info_list = function_info_list
        self.curr_index = -1
        self.refs_list = []
        self.is_refs_to = is_refs_to
        self.theme = light_theme

    def isFollowable(self, col: int) -> bool:
        """Is the address possible to follow in the disassembly view?"""
        if col == self.COL_ADDR:
            return True
        if col == self.COL_TOADDR:
            return True
        return False

    def findOffsetIndex(self, data: int) -> int:
        """Serches the given address on the list of functions and returns it if found."""

        index = 0
        for func_info in self.function_info_list:
            if data >= func_info.start and data <= func_info.end:
                return index
            index += 1
        return -1

    def setCurrentIndex(self, curr_index: int) -> None:
        self.curr_index = curr_index
        if self.curr_index == (-1) or self.curr_index >= len(self.function_info_list):
            # reset list
            self.refs_list = []
        else:
            if self.is_refs_to:
                self.refs_list = self.function_info_list[self.curr_index].refs_list
            else:
                self.refs_list = self.function_info_list[self.curr_index].called_list
        self.reset()

    def reset(self) -> None:
        self.beginResetModel()
        self.endResetModel()

    # Qt API
    def rowCount(self, parent=None) -> int:
        return len(self.refs_list)

    def columnCount(self, parent) -> int:
        return self.COL_COUNT

    def data(self, index, role) -> Any:
        if not index.isValid():
            return None
        col = index.column()
        row = index.row()

        # curr_ref_addr = self.refs_list[row][0]

        if role == QtCore.Qt.UserRole:
            if self.isFollowable(col):
                return self._getAddrToFollow(row, col)
        elif role == QtCore.Qt.DisplayRole or role == QtCore.Qt.EditRole:
            return self._displayData(row, col)
        elif role == QtCore.Qt.BackgroundColorRole:
            return self._displayBackground(row, col)
        else:
            return None

    def flags(self, index) -> Any:
        if not index.isValid():
            return None
        flags = QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsSelectable
        return flags

    def headerData(self, section, orientation, role=QtCore.Qt.DisplayRole) -> Any:
        if role == QtCore.Qt.DisplayRole:
            return self._displayHeader(orientation, section)
        else:
            return None


# --------------------------------------------------------------------------
# custom views:


class FunctionsView_t(QtWidgets.QTableView):
    """The top view: listing all the functions."""

    # private
    def _get_default_color(self) -> None:
        return idc.DEFCOLOR

    def _set_segment_color(self, ea, color) -> None:
        seg = idaapi.getseg(ea)
        seg.color = color
        seg.update()

    def _set_item_color(self, addr: int, color: int) -> None:
        ea = addr
        # reset to default
        self._set_segment_color(ea, self.color_normal)
        # set desired item color
        set_color(addr, CIC_ITEM, color)

    def _get_themed_color(self, color: int, theme: list) -> int:
        if theme == dark_theme:
            color_hilight = color_to_val(QtGui.QColor(color).darker(200))
        else:
            color_hilight = color
        return color_hilight

    # public
    def __init__(self, dataManager, color_hilight, func_model, parent=None) -> None:
        super(FunctionsView_t, self).__init__(parent=parent)
        self.setSelectionMode(QtWidgets.QAbstractItemView.SingleSelection)
        #
        self.color_normal = self._get_default_color()
        self.prev_addr = BADADDR
        self.color_hilight = color_hilight
        self.func_model = func_model
        self.dataManager = dataManager

        self.setMouseTracking(True)
        self.setAutoFillBackground(True)
        self.theme = light_theme

    # Qt API
    def currentChanged(self, current, previous) -> None:
        index_data = self.get_index_data(current)
        self.dataManager.setCurrentRva(index_data)

    def hilight_addr(self, addr: int) -> None:
        curr_theme = get_theme()
        if curr_theme is not None:
            self.theme = curr_theme

        color_hilight = self._get_themed_color(self.color_hilight, self.theme)

        if self.prev_addr != BADADDR:
            self._set_item_color(self.prev_addr, self.color_normal)
        if addr != BADADDR:
            self._set_item_color(addr, color_hilight)
        self.prev_addr = addr

    def get_index_data(self, index: Any) -> Any:
        if not index.isValid():
            return None
        try:
            data_val = index.data(QtCore.Qt.UserRole)
            if data_val is None:
                return None
            index_data = data_val
        except ValueError:
            return None
        return index_data

    def mousePressEvent(self, event: Any) -> None:
        event.accept()
        # index = self.indexAt(event.pos())
        # data = self.get_index_data(index)
        super(QtWidgets.QTableView, self).mousePressEvent(event)

    def mouseDoubleClickEvent(self, event: Any) -> None:
        event.accept()
        index = self.indexAt(event.pos())
        if not index.isValid():
            return
        data = self.get_index_data(index)
        if not data:
            super(QtWidgets.QTableView, self).mouseDoubleClickEvent(event)
            return
        col = index.column()
        if self.func_model.isFollowable(col):
            self.hilight_addr(data)
            jumpto(data)
        super(QtWidgets.QTableView, self).mouseDoubleClickEvent(event)

    def mouseMoveEvent(self, event: Any) -> None:
        index = self.indexAt(event.pos())
        if not index.isValid():
            return
        col = index.column()
        if self.func_model.isFollowable(col):
            self.setCursor(QtCore.Qt.PointingHandCursor)
        else:
            self.setCursor(QtCore.Qt.ArrowCursor)

    def leaveEvent(self, event: Any) -> None:
        self.setCursor(QtCore.Qt.ArrowCursor)

    def OnDestroy(self) -> None:
        self.hilight_addr(BADADDR)


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


class FunctionsMapper_t(QObject):
    """The class keeping the mapping of all the functions."""

    # private

    def _isImportStart(self, start: int) -> bool:
        """Check if the given function is imported or internal."""

        if start in self._importsSet:
            return True
        if print_insn_mnem(start) == "call":
            return False
        # print(print_insn_mnem(start))
        op = get_operand_value(start, 0)
        if op in self._importsSet:
            return True
        return False

    def imports_names_callback(self, ea: int, name: str, ord: Any) -> bool:
        """A callback adding a particular name and offset to the internal set of the imported functions."""

        self._importsSet.add(ea)
        self._importNamesSet.add(name)
        # True -> Continue enumeration
        return True

    def _loadImports(self) -> None:
        """Enumerates imported functions with the help of IDA API and adds them to the internal sets."""

        self._importsSet = set()
        self._importNamesSet = set()
        nimps = idaapi.get_import_module_qty()
        for i in range(0, nimps):
            try:
                idaapi.enum_import_names(i, self.imports_names_callback)
            except:
                idaapi.msg("Failed to fetch function name")

    def _isImportName(self, name: str) -> bool:
        """Checks if the given name belongs to the imported function with the help of internal set."""

        if name in self._importNamesSet:
            return True
        return False

    def _listRefsTo(self, start: int) -> List[Tuple[int, int]]:
        """Make a list of all the references to the given function.
        Args:
          func : The function references to which we are searching.
          start : The function's start offset.

        Returns:
          list : A list of tuples.
            Each tuple represents: the offsets:
            0 : the offset from where the given function was referenced by the foreign function
            1 : the function's start address
        """

        func_refs_to = XrefsTo(start, 1)
        refs_list = []
        for ref in func_refs_to:
            if idc.print_insn_mnem(ref.frm) == "":
                continue
            refs_list.append((ref.frm, start))
        return refs_list

    def _getCallingOffset(self, func, called_list) -> List[Tuple[int, int]]:
        """Lists the offsets from where the given function references the list of other function."""

        start = get_func_attr(func, FUNCATTR_START)
        end = prev_addr(get_func_attr(func, FUNCATTR_END))
        # func_name = _getFunctionNameAt(start)
        curr = start
        calling_list = []
        while True:
            if curr >= end:
                break
            op = get_operand_value(curr, 0)
            if op in called_list:
                calling_list.append((curr, op))
            curr = next_addr(curr)
        return calling_list

    def _listRefsFrom(self, func, start: int, end: int) -> List[Tuple[int, int]]:
        """Make a list of all the references made from the given function.

        Args:
          func : The function inside of which we are searching.
          start : The function's start offset.
          end : The function's end offset.

        Returns:
          list : A list of tuples. Each tuple represents:
            0 : the offset from where the given function referenced the other entity
            1 : the address that was referenced
        """

        dif = end - start
        called_list = []
        func_name = _getFunctionNameAt(start)

        for indx in range(0, dif):
            addr = start + indx
            func_refs_from = XrefsFrom(addr, 1)
            for ref in func_refs_from:
                if _getFunctionNameAt(ref.to) == func_name:
                    # skip jumps inside self
                    continue
                called_list.append(ref.to)
        calling_list = self._getCallingOffset(func, called_list)
        return calling_list

    def _loadLocals(self) -> None:
        """Enumerates functions using IDA API and loads them into the internal mapping."""
        self._loadImports()
        for func in Functions():
            start = get_func_attr(func, FUNCATTR_START)
            end = prev_addr(get_func_attr(func, FUNCATTR_END))

            is_import = self._isImportStart(start)

            refs_list = self._listRefsTo(start)
            calling_list = self._listRefsFrom(func, start, end)

            func_info = FunctionInfo_t(start, end, refs_list, calling_list, is_import)
            self._functionsMap[va_to_rva(start)] = func_info
            self._functionsMap[va_to_rva(end)] = func_info
            self.funcList.append(func_info)

    # public
    def __init__(self, parent=None) -> None:
        super(FunctionsMapper_t, self).__init__(parent=parent)
        self._functionsMap = dict()
        self.funcList = []  # public
        self._loadLocals()

    def funcAt(self, rva: int) -> FunctionInfo_t:
        func_info = self._functionsMap[rva]
        return func_info


class FunctionsListForm_t(PluginForm):
    """The main form of the IFL plugin."""

    # private
    _LIVE_FILTER = True

    def _listFunctionsAddr(self) -> List[int]:
        """Lists all the starting addresses of the functions using IDA API."""

        fn_list = list()
        for func in Functions():
            start = get_func_attr(func, FUNCATTR_START)
            fn_list.append(start)
        return fn_list

    def _saveFunctionsNames(
        self, file_name: Optional[str], ext: str, skip_unnamed: bool
    ) -> bool:
        """Saves functions names and offsets from the internal mappings into a file.
        Fromats: CSV (default), or TAG (PE-bear, PE-sieve compatibile).
        """

        if file_name is None or len(file_name) == 0:
            return False
        delim = ","
        if ".tag" in ext:  # a TAG format was chosen
            delim = ";"
        fn_list = list()
        for func in Functions():
            start = get_func_attr(func, FUNCATTR_START)
            func_name = _getFunctionNameAt(start)
            if skip_unnamed:
                if func_name.startswith("sub_"):
                    continue
            start_rva = va_to_rva(start)
            line = "%lx%c%s" % (start_rva, delim, func_name)
            fn_list.append(line)
        idaapi.msg(str(file_name))
        with open(file_name, "w") as f:
            for item in fn_list:
                f.write("%s\n" % item)
            return True
        return False

    def _stripImportName(self, func_name) -> str:
        """Keep only ImportName, without the DLL name, and the ordinal."""

        fn1 = func_name.split(".")
        if len(fn1) >= 2:
            func_name = fn1[1].strip()
        fn1 = func_name.split("#")
        if len(fn1) >= 2:
            func_name = fn1[0].strip()
        return func_name

    def _getBitness(self):
        if idaapi.IDA_SDK_VERSION >= 900:
            if idaapi.inf_is_64bit():
                return 64
            elif idaapi.inf_is_32bit_exactly():
                return 32
        else:
            info = idaapi.get_inf_structure()
            if info.is_64bit():
                return 64
            elif info.is_32bit():
                return 32
        return None

    def _defineImportThunk(self, start, thunk_val):
        """If the binary has the Import Thunk filled, define it as a data chunk of appropriate size."""

        bitness = self._getBitness()
        if bitness is None:
            return False
        if bitness == 64:
            curr_val = idc.get_qword(start)
            if curr_val == thunk_val:
                return ida_bytes.create_data(start, idaapi.FF_QWORD, 8, idaapi.BADADDR)
        elif bitness == 32:
            curr_val = ida_bytes.get_dword(start)
            if curr_val == thunk_val:
                return ida_bytes.create_data(start, idaapi.FF_DWORD, 4, idaapi.BADADDR)
        return False

    def _loadFunctionsNames(
        self, file_name: Optional[str], ext: str, loadBase: int
    ) -> Optional[Tuple[int, int]]:
        """Loads functions names from the given file into the internal mappings.
        Fromats: CSV (default), or TAG (PE-bear, PE-sieve compatibile).
        """

        if file_name is None or len(file_name) == 0:
            return None
        curr_functions = self._listFunctionsAddr()
        delim = ","  # new delimiter (for CSV format)
        delim2 = ":"  # old delimiter
        rva_indx = 0
        cmt_indx = 1
        is_imp_list = False
        if ".imports.txt" in ext:
            is_imp_list = True
            cmt_indx = 2
        if ".tag" in ext:  # a TAG format was chosen
            delim2 = ";"
        functions = 0
        comments = 0
        with open(file_name, "r") as f:
            for line in f.readlines():
                line = line.strip()
                fn = line.split(delim)
                if len(fn) < 2:
                    fn = line.split(delim2)  # try old delimiter
                if len(fn) < 2:
                    continue
                start = 0
                addr_chunk = fn[rva_indx].strip()
                if not _is_hex_str(addr_chunk):
                    continue
                try:
                    start = int(addr_chunk, 16)
                except ValueError:
                    # this line doesn't start from an offset, so skip it
                    continue
                func_name = fn[cmt_indx].strip()
                if start < loadBase:  # it is RVA
                    start = start + loadBase  # convert to VA

                if is_imp_list or (start in curr_functions):
                    if is_imp_list:
                        func_name = self._stripImportName(func_name)
                        thunk_val = int(fn[1].strip(), 16)
                        self._defineImportThunk(start, thunk_val)

                    if self.subDataManager.setFunctionName(start, func_name):
                        functions += 1
                        continue

                set_cmt(start, func_name, 1)  # set the name as a repeatable comment
                set_cmt(
                    start, func_name, 0
                )  # make sure to overwrite the default IDA comments
                comments += 1
        return (functions, comments)

    def _setup_sorted_model(self, view, model) -> QtCore.QSortFilterProxyModel:
        """Connects the given sorted data model with the given view."""

        sorted_model = QtCore.QSortFilterProxyModel()
        sorted_model.setDynamicSortFilter(True)
        sorted_model.setSourceModel(model)
        view.setModel(sorted_model)
        view.setSortingEnabled(True)
        #
        sorted_model.setParent(view)
        model.setParent(sorted_model)
        return sorted_model

    def _update_current_offset(self, view, refs_model, offset) -> None:
        """Update the given data model to follow given offset."""

        if offset:
            index = refs_model.findOffsetIndex(offset)
        else:
            index = -1
        refs_model.setCurrentIndex(index)
        refs_model.reset()
        view.reset()
        view.repaint()

    def _update_function_name(self, ea: int) -> None:
        """Sets on the displayed label the name of the function and it's arguments."""

        try:
            func_info = self.funcMapper.funcAt(va_to_rva(ea))
        except KeyError:
            return

        func_type = func_info.type
        func_args = _getArgsDescription(ea)
        func_name = _getFunctionNameAt(ea)
        self.refs_label.setText(f"{func_type} <b>{func_name}</b> {func_args}")

    def _update_ref_tabs(self, ea: int) -> None:
        """Sets on the tabs headers the numbers of references to the selected function."""

        tocount = 0
        fromcount = 0
        try:
            func_info = self.funcMapper.funcAt(va_to_rva(ea))
            tocount = len(func_info.refs_list)
            fromcount = len(func_info.called_list)
        except KeyError:
            pass
        self.refs_tabs.setTabText(0, "Is referred by %d:" % tocount)
        self.refs_tabs.setTabText(1, "Refers to %d:" % fromcount)

    def adjustColumnsToContents(self) -> None:
        """Adjusts columns' sizes to fit the data."""

        self.addr_view.resizeColumnToContents(0)
        self.addr_view.resizeColumnToContents(1)
        self.addr_view.resizeColumnToContents(2)
        #
        self.addr_view.resizeColumnToContents(5)
        self.addr_view.resizeColumnToContents(6)
        self.addr_view.resizeColumnToContents(7)

    # public
    # @pyqtSlot()

    def longoperationcomplete(self) -> None:
        """A callback executed when the current RVA has changed."""
        global g_DataManager

        data = g_DataManager.currentRva
        self.setRefOffset(data)

    def setRefOffset(self, data: Any) -> None:
        """Updates the views to follow to the given RVA."""

        if not data:
            return
        self._update_current_offset(self.refs_view, self.refsto_model, data)
        self._update_current_offset(self.refsfrom_view, self.refsfrom_model, data)
        self._update_ref_tabs(data)
        self._update_function_name(data)

    def filterByColumn(self, col_num, str) -> None:
        """Applies a filter defined by the string on data model."""

        filter_type = QtCore.QRegExp.FixedString
        sensitivity = QtCore.Qt.CaseInsensitive
        if self.criterium_id != 0:
            filter_type = QtCore.QRegExp.RegExp
        self.addr_sorted_model.setFilterRegExp(
            QtCore.QRegExp(str, sensitivity, filter_type)
        )
        self.addr_sorted_model.setFilterKeyColumn(col_num)

    def filterChanged(self) -> None:
        """A wrapper for the function: filterByColumn(self, col_num, str)"""

        self.filterByColumn(self.filter_combo.currentIndex(), self.filter_edit.text())

    def filterKeyEvent(self, event: Any = None) -> None:
        if event is not None:
            QtWidgets.QLineEdit.keyReleaseEvent(self.filter_edit, event)
        if event and (
            not self.is_livefilter
            and event.key() != QtCore.Qt.Key_Enter
            and event.key() != QtCore.Qt.Key_Return
        ):
            return
        self.filterChanged()

    def criteriumChanged(self) -> None:
        """A callback executed when the criterium of sorting has changed and the data has to be sorted again."""

        self.criterium_id = self.criterium_combo.currentIndex()
        if self.criterium_id == 0:
            self.filter_edit.setPlaceholderText("keyword")
        else:
            self.filter_edit.setPlaceholderText("regex")
        self.filterChanged()

    def liveSearchCheckBox(self) -> None:
        self.is_livefilter = self.livefilter_box.isChecked()
        if self.is_livefilter:
            self.filterByColumn(
                self.filter_combo.currentIndex(), self.filter_edit.text()
            )

    def alternateRowColors(self, enable):
        self.refsfrom_view.setAlternatingRowColors(enable)
        self.addr_view.setAlternatingRowColors(enable)
        self.refs_view.setAlternatingRowColors(enable)

    def OnCreate(self, form) -> None:
        """Called when the plugin form is created"""

        QtWidgets.QApplication.setAttribute(QtCore.Qt.AA_EnableHighDpiScaling, True)
        QtWidgets.QApplication.setAttribute(QtCore.Qt.AA_UseHighDpiPixmaps, True)
    
        # init data structures:
        self.funcMapper = FunctionsMapper_t()
        self.criterium_id = 0

        # Get parent widget
        self.parent = self.FormToPyQtWidget(form)

        # Create models
        self.subDataManager = DataManager()

        self.table_model = TableModel_t(self.funcMapper.funcList)

        # init
        self.addr_sorted_model = QtCore.QSortFilterProxyModel()
        self.addr_sorted_model.setDynamicSortFilter(True)
        self.addr_sorted_model.setSourceModel(self.table_model)
        self.addr_view = FunctionsView_t(
            g_DataManager, COLOR_HILIGHT_FUNC, self.table_model
        )
        self.addr_view.setModel(self.addr_sorted_model)
        self.addr_view.setSortingEnabled(True)
        self.addr_view.setWordWrap(False)
        self.addr_view.horizontalHeader().setStretchLastSection(False)
        self.addr_view.verticalHeader().show()

        self.adjustColumnsToContents()
        #
        self.refsto_model = RefsTableModel_t(self.funcMapper.funcList, True)
        self.refs_view = FunctionsView_t(
            self.subDataManager, COLOR_HILIGHT_REFTO, self.refsto_model
        )
        self._setup_sorted_model(self.refs_view, self.refsto_model)
        self.refs_view.setColumnHidden(RefsTableModel_t.COL_TOADDR, True)
        self.refs_view.setWordWrap(False)

        self.refsfrom_model = RefsTableModel_t(self.funcMapper.funcList, False)
        self.refsfrom_view = FunctionsView_t(
            self.subDataManager, COLOR_HILIGHT_REFFROM, self.refsfrom_model
        )
        self._setup_sorted_model(self.refsfrom_view, self.refsfrom_model)
        self.refsfrom_view.setColumnHidden(RefsTableModel_t.COL_TOADDR, True)
        self.refsfrom_view.setWordWrap(False)

        # add a box to enable/disable live filtering
        self.livefilter_box = QtWidgets.QCheckBox("Live filtering")
        self.livefilter_box.setToolTip(
            "If live filtering is enabled, functions are searched as you type in the edit box.\nOtherwise they are searched when you press Enter."
        )
        self.livefilter_box.setChecked(self._LIVE_FILTER)
        self.is_livefilter = self._LIVE_FILTER
        # connect SIGNAL
        self.livefilter_box.stateChanged.connect(self.liveSearchCheckBox)

        # important for proper order of objects destruction:
        self.table_model.setParent(self.addr_sorted_model)
        self.addr_sorted_model.setParent(self.addr_view)

        # connect SIGNAL
        g_DataManager.updateSignal.connect(self.longoperationcomplete)

        # Create a Tab widget for references:
        self.refs_tabs = QtWidgets.QTabWidget()
        self.refs_tabs.insertTab(0, self.refs_view, "Is refered by")
        self.refs_tabs.insertTab(1, self.refsfrom_view, "Refers to")

        # Create filter
        self.filter_edit = QtWidgets.QLineEdit()
        self.filter_edit.setPlaceholderText("keyword")
        self.filter_edit.keyReleaseEvent = self.filterKeyEvent

        self.filter_combo = QtWidgets.QComboBox()
        self.filter_combo.addItems(TableModel_t.header_names)
        self.filter_combo.setCurrentIndex(TableModel_t.COL_NAME)
        # connect SIGNAL
        self.filter_combo.activated.connect(self.filterChanged)

        self.criterium_combo = QtWidgets.QComboBox()
        criteria = ["contains", "matches"]
        self.criterium_combo.addItems(criteria)
        self.criterium_combo.setCurrentIndex(0)
        # connect SIGNAL
        self.criterium_combo.activated.connect(self.criteriumChanged)

        filter_panel = QtWidgets.QFrame()
        filter_layout = QtWidgets.QHBoxLayout()
        filter_layout.addWidget(QtWidgets.QLabel("Where "))
        filter_layout.addWidget(self.filter_combo)
        filter_layout.addWidget(self.criterium_combo)
        filter_layout.addWidget(self.filter_edit)

        filter_panel.setLayout(filter_layout)
        filter_panel.setAutoFillBackground(True)

        self.refs_label = QtWidgets.QLabel("Function")
        self.refs_label.setTextFormat(QtCore.Qt.RichText)
        self.refs_label.setWordWrap(True)

        panel1 = QtWidgets.QFrame()
        layout1 = QtWidgets.QVBoxLayout()
        panel1.setLayout(layout1)

        layout1.addWidget(filter_panel)
        layout1.addWidget(self.livefilter_box)
        layout1.addWidget(self.addr_view)
        layout1.setContentsMargins(0, 0, 0, 0)

        panel2 = QtWidgets.QFrame()
        layout2 = QtWidgets.QVBoxLayout()
        layout2.addWidget(self.refs_label)
        layout2.addWidget(self.refs_tabs)
        layout2.addWidget(self._makeButtonsPanel())
        layout2.setContentsMargins(0, 10, 0, 0)
        panel2.setLayout(layout2)

        self.main_splitter = QtWidgets.QSplitter()
        self.main_splitter.setOrientation(QtCore.Qt.Vertical)
        self.main_splitter.addWidget(panel1)
        self.main_splitter.addWidget(panel2)

        # Populate PluginForm
        layout = QtWidgets.QVBoxLayout()
        layout.addWidget(self.main_splitter)
        layout.setSpacing(0)
        layout.setContentsMargins(0, 0, 0, 0)
        self.parent.setLayout(layout)
        self.alternateRowColors(IS_ALTERNATE_ROW)

        idaapi.set_dock_pos(PLUGIN_NAME, None, idaapi.DP_RIGHT)

    def _makeButtonsPanel(self) -> QtWidgets.QFrame:
        """Creates on the form's bottom the panel with buttons."""

        buttons_panel = QtWidgets.QFrame()
        buttons_layout = QtWidgets.QHBoxLayout()
        buttons_panel.setLayout(buttons_layout)

        importButton = QtWidgets.QPushButton("Load names")
        importButton.clicked.connect(self.importNames)
        buttons_layout.addWidget(importButton)

        exportButton = QtWidgets.QPushButton("Save names")
        exportButton.clicked.connect(self.exportNames)
        buttons_layout.addWidget(exportButton)
        return buttons_panel

    def importNames(self) -> None:
        """Imports functions list from a file."""

        file_name, ext = QtWidgets.QFileDialog.getOpenFileName(
            None,
            "Import functions names",
            QtCore.QDir.homePath(),
            "CSV Files (*.csv);;TAG Files: PE-bear, PE-sieve compatibile (*.tag);;IMPORTS.TXT: generated by PE-sieve (*.imports.txt);;All files (*)",
        )
        if file_name is None or len(file_name) == 0:
            return
        default_base = idaapi.get_imagebase()
        base_str, ok = QtWidgets.QInputDialog.getText(
            None,
            "Rebase tags:",
            "Current module base:",
            QtWidgets.QLineEdit.Normal,
            hex(default_base),
        )
        if not ok:
            return
        if _is_hex_str(base_str):
            load_base = int(base_str, 16)
        else:
            load_base = default_base
        names = self._loadFunctionsNames(file_name, ext, load_base)
        if names is None:
            idaapi.warning(f"Malformed file %s" % file_name)
            return
        (loaded, comments) = names
        if loaded == 0 and comments == 0:
            idaapi.warning("Failed importing functions names! Not matching offsets!")
        else:
            idaapi.info(
                "Imported %d function names and %d comments" % (loaded, comments)
            )

    def exportNames(self) -> None:
        """Exports functions list into a file."""

        file_name, ext = QtWidgets.QFileDialog.getSaveFileName(
            None,
            "Export functions names",
            QtCore.QDir.homePath(),
            "CSV Files (*.csv);;TAG Files (*.tag)",
        )
        skip_unnamed = True
        if file_name is not None and len(file_name) > 0:
            if not self._saveFunctionsNames(file_name, ext, skip_unnamed):
                idaapi.warning("Failed exporting functions names!")
            else:
                idaapi.info("Exported to: " + file_name)

    def OnClose(self, form: Any) -> None:
        """Called when the plugin form is closed"""

        # clear last selection
        self.addr_view.hilight_addr(BADADDR)
        self.refs_view.hilight_addr(BADADDR)
        self.refsfrom_view.hilight_addr(BADADDR)
        del self

    def Show(self) -> PluginForm.Show:
        """Creates the form if not created or sets the focus if the form already exits."""

        title = PLUGIN_NAME + " v" + __VERSION__
        return PluginForm.Show(self, title, options=PluginForm.WOPN_PERSIST)


# --------------------------------------------------------------------------
class IFLMenuHandler(idaapi.action_handler_t):
    """Manages menu items belonging to IFL."""

    def __init__(self) -> None:
        idaapi.action_handler_t.__init__(self)

    def activate(self, ctx: Any) -> int:
        open_form()
        return 1

    def update(self, ctx: Any) -> int:
        return idaapi.AST_ENABLE_ALWAYS


# --------------------------------------------------------------------------
def open_form() -> None:
    global m_functionInfoForm
    global g_DataManager
    # -----
    try:
        g_DataManager
    except Exception:
        g_DataManager = DataManager()
    # -----
    try:
        m_functionInfoForm
    except Exception:
        idaapi.msg("%s\nLoading Interactive Function List...\n" % VERSION_INFO)
        m_functionInfoForm = FunctionsListForm_t()

    m_functionInfoForm.Show()


# --------------------------------------------------------------------------
# IDA api:


class funclister_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = "Interactive Functions List"

    help = (
        "Interactive Function List. Comments? Remarks? Mail to: hasherezade@gmail.com"
    )
    wanted_name = PLUGIN_NAME
    wanted_hotkey = ""
    
    def init(self) -> None:
        idaapi.register_action(
            idaapi.action_desc_t(
                "ifl:open",  # action name
                PLUGIN_NAME,
                IFLMenuHandler(),
                PLUGIN_HOTKEY,
                "Opens Interactive Function List Pane",
            )
        )

        idaapi.attach_action_to_menu("View/", "ifl:open", idaapi.SETMENU_APP)

        return idaapi.PLUGIN_OK

    def run(self, arg: Any) -> None:
        open_form()


def PLUGIN_ENTRY() -> funclister_t:
    return funclister_t()


if __name__ == "__main__":
    PLUGIN_ENTRY().init()


repository: ynvb/DIE
file: DIE.py
description: Dynamic IDA Enrichment
url: https://github.com/ynvb/DIE/blob/453fbdf137bfe0510e760df73441d41d35ea74d6/DIE.py
wanted_name: DIE
comment: Dynamic IDA Enrichment plugin (aka. DIE)
created_at: 2014-12-12 17:52:39+00:00
pushed_at: 2022-03-02 19:44:30+00:00
content:
import networkx as nx
from awesome.context import ignored
import sark
import sark.qt
from sark.qt import QtGui, QtCore
import logging
import logging.handlers as handlers

import os
from time import ctime

from idaapi import plugin_t
import idaapi
import idautils
import idc
import sark.ui


from DIE.Lib.IDAConnector import *
import DIE.Lib.DieConfig
import DIE.Lib.DIEDb
from DIE.Lib import DebugAPI

#import DIE.UI.BPView
import DIE.UI.FunctionViewEx
import DIE.UI.ValueViewEx
import DIE.UI.ParserView
import DIE.UI.BPView
import DIE.UI.SetupView
import DIE.UI.Die_Icons
from DIE.UI.FuncScopeChooser import ScopeChooser
from DIE.UI.AboutScreen import AboutWindow

from DIE.Lib.DIE_Exceptions import DbFileMismatch

class MenuHelperException(Exception):
    pass

  
class DieManager:
    """
    Manage the DIE framework
    """

    def __init__(self, is_dbg_log=True, is_dbg_pause=False, is_dbg_profile=False):
        ### Logging ###

        log_filename = os.path.join(os.getcwd(), "DIE.log")

        self._menus_names = []

        #TODO: Fix logging to include rotating_file_handler \ console_logging
        if is_dbg_log:
            logging.basicConfig(filename=log_filename,
                        level=logging.DEBUG,
                        format='[%(asctime)s] [%(levelname)s] [%(name)s][%(filename)s:%(lineno)s] : %(message)s')
        else:
             logging.basicConfig(filename=log_filename,
                    level=logging.INFO,
                    format='[%(asctime)s] [%(levelname)s] [%(name)s][%(filename)s:%(lineno)s] : %(message)s')

        idaapi.msg("Logfile created at %s\n" % log_filename)
        self.logger = logging.getLogger(__name__)

        ### DIE Configuration ###
        self.config_file_name = os.path.join(os.getcwd(), "DIE.cfg")
        DIE.Lib.DieConfig.initialize()
        config = DIE.Lib.DieConfig.get_config()
        try:
            config.load(self.config_file_name)
        except IOError:
            pass

        except:
            import traceback
            idaapi.msg(traceback.format_exc())

        self.die_config = config

        self.icon_list = {}

        self.debugAPI = DebugAPI.DebugHooker(is_dbg_pause, is_dbg_profile)
        DIE.Lib.DIEDb.initialize_db()
        self.die_db = DIE.Lib.DIEDb.get_db()

        self.is_marked = False

        DIE.UI.FunctionViewEx.initialize()
        DIE.UI.ValueViewEx.initialize()
        DIE.UI.BPView.initialize()
        DIE.UI.ParserView.initialize()
        DIE.UI.Die_Icons.initlialize()
        self.function_view = DIE.UI.FunctionViewEx.get_view()
        self.value_view = DIE.UI.ValueViewEx.get_view()
        self.bp_view = DIE.UI.BPView.get_view()
        self.parser_view = DIE.UI.ParserView.get_view()

        self.load_icons()

        return

    ###########################################################################
    # Icons

    def load_icon(self, icon_filename, icon_key_name):
        """
        Load a single custom icon
        @param icon_filename: Icon file name
        @param icon_key_name: The key value to store the icon with in the icon_list.
        """
        try:
            icons_path = self.die_config.icons_path

            icon_filename = os.path.join(icons_path, icon_filename)
            icon_num = idaapi.load_custom_icon(icon_filename)
            self.icon_list[icon_key_name.lower()] = icon_num
            return True

        except Exception as ex:
            self.logger.error("Failed to load icon %s: %s", icon_filename, ex)
            return False

    def load_icons(self):
        """
        Load custom DIE Icons
        """
        self.load_icon("dbg.png", "debug")
        self.load_icon("dbg_all.png", "debug_all")
        self.load_icon("dbg_custom.png", "debug_scope")
        self.load_icon("die.png", "die")
        self.load_icon("funcview.png", "function_view")
        self.load_icon("valueview.png", "value_view")
        self.load_icon("stop.png", "exception_view")
        self.load_icon("settings.png", "settings")
        self.load_icon("plugins.png", "plugins")
        self.load_icon("save.png", "save")
        self.load_icon("load.png", "load")


    ###########################################################################
    # Menu Items

    # Handlers
    class DIE_gohere_Handler(idaapi.action_handler_t):
        def __init__(self, outer_instance):
             idaapi.action_handler_t.__init__(self)
             self.outer_instance = outer_instance

        def activate(self, ctx):
            self.outer_instance.go_here()
            return True

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

    class DIE_goall_Handler(idaapi.action_handler_t):
        def __init__(self, outer_instance):
            idaapi.action_handler_t.__init__(self)
            self.outer_instance = outer_instance

        def activate(self, ctx):
            self.outer_instance.go_all()
            return True

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

    class DIE_goscope_Handler(idaapi.action_handler_t):
        def __init__(self, outer_instance):
            idaapi.action_handler_t.__init__(self)
            self.outer_instance = outer_instance

        def activate(self, ctx):
            self.outer_instance.show_scope_chooser()
            return True

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

    class DIE_show_funcview_Handler(idaapi.action_handler_t):
        def __init__(self, outer_instance):
            idaapi.action_handler_t.__init__(self)
            self.outer_instance = outer_instance

        def activate(self, ctx):
            self.outer_instance.show_function_view()
            return True

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

    class DIE_show_valview_Handler(idaapi.action_handler_t):
        def __init__(self, outer_instance):
            idaapi.action_handler_t.__init__(self)
            self.outer_instance = outer_instance

        def activate(self, ctx):
            self.outer_instance.show_value_view()
            return True

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

    class DIE_show_parsersview_Handler(idaapi.action_handler_t):
        def __init__(self, outer_instance):
            idaapi.action_handler_t.__init__(self)
            self.outer_instance = outer_instance

        def activate(self, ctx):
            self.outer_instance.show_parser_view()
            return True

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

    class DIE_show_exceptionsview_Handler(idaapi.action_handler_t):
        def __init__(self, outer_instance):
            idaapi.action_handler_t.__init__(self)
            self.outer_instance = outer_instance

        def activate(self, ctx):
            self.outer_instance.show_breakpoint_view()
            return True

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

    class DIE_load_db_Handler(idaapi.action_handler_t):
        def __init__(self, outer_instance):
            idaapi.action_handler_t.__init__(self)
            self.outer_instance = outer_instance

        def activate(self, ctx):
            self.outer_instance.load_db()
            return True

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

    class DIE_save_db_Handler(idaapi.action_handler_t):
        def __init__(self, outer_instance):
            idaapi.action_handler_t.__init__(self)
            self.outer_instance = outer_instance

        def activate(self, ctx):
            self.outer_instance.save_db()
            return True

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

    class DIE_show_cfg_Handler(idaapi.action_handler_t):
        def __init__(self, outer_instance):
            idaapi.action_handler_t.__init__(self)
            self.outer_instance = outer_instance

        def activate(self, ctx):
            self.outer_instance.show_cfg()
            return True

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

    class DIE_mark_execflow_Handler(idaapi.action_handler_t):
        def __init__(self, outer_instance):
            idaapi.action_handler_t.__init__(self)
            self.outer_instance = outer_instance

        def activate(self, ctx):
            self.outer_instance.mark_exec_flow()
            return True

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

    class DIE_show_about_Handler(idaapi.action_handler_t):
        def __init__(self, outer_instance):
            idaapi.action_handler_t.__init__(self)
            self.outer_instance = outer_instance

        def activate(self, ctx):
            self.outer_instance.show_about()
            return True

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

    class DIE_show_settings_Handler(idaapi.action_handler_t):
        def __init__(self, outer_instance):
            idaapi.action_handler_t.__init__(self)
            self.outer_instance = outer_instance

        def activate(self, ctx):
            self.outer_instance.show_settings()
            return True

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


    def add_menu_item_helper(self, name, text, tooltip, handler, icon, shortcut):
        description = idaapi.action_desc_t(name, text, handler, shortcut, tooltip, icon)
        idaapi.register_action(description)
        idaapi.attach_action_to_menu("DIE/" + text, name, idaapi.SETMENU_APP)

        self._menus_names.append(name)


    def add_menu_items(self):
        # Add root level menu
        idaapi.create_menu("DIE", "DIE")

        # Save DieDB
        self.add_menu_item_helper("DIE:savedb", "Save DieDB", 'DIE Save DieDB',
                                  self.DIE_save_db_Handler(self), self.icon_list["save"], '')

        # Load DieDB
        self.add_menu_item_helper("DIE:loaddb", "Load DieDB", 'DIE Load DieDB',
                                  self.DIE_load_db_Handler(self), self.icon_list["load"], '')

        # Show complete execution CFG
        self.add_menu_item_helper("DIE:showcfg", "Show CFG", 'DIE Show CFG',
                                  self.DIE_show_cfg_Handler(self), -1, '')

        # Mark\Unmark Execution Flow
        self.add_menu_item_helper("DIE:markexecutionflow", "Mark\Unmark Execution Flow",
                                  'DIE Mark\Unmark Execution Flow', self.DIE_mark_execflow_Handler(self), -1, '')

        # Parser View
        self.add_menu_item_helper("DIE:parsersview", "Parsers View", 'DIE Parsers View',
                                  self.DIE_show_parsersview_Handler(self), self.icon_list["plugins"], '')

        # Exceptions View
        self.add_menu_item_helper("DIE:exceptionsview", "Exceptions View", 'DIE Exceptions View',
                                  self.DIE_show_exceptionsview_Handler(self), self.icon_list["exception_view"], '')

        # Value View
        self.add_menu_item_helper("DIE:valueview", "Value View", 'DIE Value View',
                                  self.DIE_show_valview_Handler(self), self.icon_list["value_view"], '')

        # Function View
        self.add_menu_item_helper("DIE:functionview", "Function View", 'DIE Function View',
                                  self.DIE_show_funcview_Handler(self), self.icon_list["function_view"], '')

        # Debug a custom scope
        self.add_menu_item_helper("DIE:debugcustomscope", "Debug a custom scope", 'DIE Debug a custom scope',
                                  self.DIE_goscope_Handler(self), self.icon_list["debug_scope"], 'Ctrl+Alt+c')

        # Debug entire code
        self.add_menu_item_helper("DIE:debugentirecode", "Debug entire code", 'DIE Debug entire code',
                                  self.DIE_goall_Handler(self), self.icon_list["debug_scope"], 'Ctrl+Alt+g')

        # Debug from current location
        self.add_menu_item_helper("DIE:gofromcurrentlocation", "Debug from current location", 'DIE from current location',
                                  self.DIE_gohere_Handler(self), self.icon_list["debug"], 'Ctrl+Alt+f')

        # Settings
        self.add_menu_item_helper("DIE:settings", "Settings", 'DIE Settings',
                                  self.DIE_show_settings_Handler(self), self.icon_list["settings"], '')

        # About
        self.add_menu_item_helper("DIE:about", "About", 'DIE About',
                                  self.DIE_show_about_Handler(self), -1, '')


    def del_menu_items(self):
        for menu_name in self._menus_names:
            idaapi.unregister_action(menu_name)
        idaapi.delete_menu("DIE")


    ###########################################################################
    # Debugging
    def go_here(self):
        self.debugAPI.start_debug(idc.here(), None, auto_start=True)

    def go_all(self):
        self.debugAPI.start_debug(None, None, auto_start=True)

    def show_scope_chooser(self):
        global chooser

        functions = get_functions()
        func_list = functions.keys()

        chooser = ScopeChooser(func_list)
        chooser.Compile()

        ok = chooser.Execute()
        if ok == 1:
            start_func = func_list[chooser.cbStartFunction.value]
            start_func_ea = functions[start_func]

            end_func = func_list[chooser.cbEndFunction.value]
            end_func_ea = functions[end_func]

            self.debugAPI.start_debug(start_func_ea, end_func_ea, True)

        chooser.Free()

    ###########################################################################
    # DIE DB
    def save_db(self):
        db_file = idc.AskFile(1, "*.ddb", "Save DIE Db File")
        if db_file is None:
            return

        self.die_db.save_db(db_file)

    def load_db(self):
        try:
            db_file = idc.AskFile(0, "*.ddb", "Load DIE Db File")
            if db_file is not None:
                self.die_db.load_db(db_file)

            if self.die_db is not None:
                self.show_db_details()

        except DbFileMismatch as mismatch:
            idaapi.msg("Error while loading DIE DB: %s\n" % mismatch)

        except Exception as ex:
            logging.exception("Error while loading DB: %s", ex)
            return False


    ###########################################################################
    # Function View
    def show_function_view(self):
        self.function_view.Show()

    ###########################################################################
    # Value View
    def show_value_view(self):
        self.value_view.Show()

    ###########################################################################
    # Parser View
    def show_parser_view(self):
        self.parser_view.Show()

    ###########################################################################
    # Exceptions View
    def show_breakpoint_view(self):
        self.bp_view.Show()

    ###########################################################################
    # About
    def show_about(self):
        AboutWindow().exec_()

    ###########################################################################
    # Settings View
    def show_settings(self):
        DIE.UI.SetupView.Show(self.config_file_name)

    ###########################################################################
    # Show DB Details

    def show_db_details(self):
        """
        Print DB details
        """
        (start_time,
         end_time,
         filename,
         num_of_functions,
         num_of_threads,
         numof_parsed_val) = self.die_db.get_run_info()

        idaapi.msg("Die DB Loaded.\n")
        idaapi.msg("Start Time: %s, End Time %s\n" % (ctime(start_time), ctime(end_time)))
        idaapi.msg("Functions: %d, Threads: %d\n" % (num_of_functions, num_of_threads))
        idaapi.msg("Parsed Values: %d\n" % numof_parsed_val)

    ###########################################################################
    # Mark\Unmark Execution Flow

    def mark_exec_flow(self):
        """
        Mark \ Unmark execution flow
        """
        color = 0x123456
        if self.is_marked:
            color = None

        with sark.ui.Update():
            for func_ea in self.die_db.get_function_counter():
                try:
                    sark.Function(func_ea).color = color
                except sark.exceptions.SarkNoFunction:
                    pass

        # Swap is_marked value
        self.is_marked = not self.is_marked

    ###########################################################################
    # Show CFG

    def show_cfg(self):
        """
        Show execution Call flow graph
        """
        cfg = self.die_db.get_call_graph_complete()
        graph = nx.DiGraph()

        if not cfg:
            idaapi.msg("No CFG to display")
            return

        for ctxt_node in cfg:
            (from_address, to_address) = ctxt_node
            graph.add_edge(from_address, to_address)

        viewer = sark.ui.NXGraph(graph, "Callgraph for {}".format("Exection CFG"), handler=sark.ui.AddressNodeHandler())
        viewer.Show()

    def show_logo(self):
        """
        Show DIE Logo
        """
        idaapi.msg('-----------------------------------------------------\n'
                   '                           _________-----_____       \n'
                   '        _____------           __      ----_          \n'
                   ' ___----             ___------             /\        \n'
                   '    ----________        ----                 \       \n'
                   '                -----__    |             _____)      \n'
                   '                     __-                /    /\      \n'
                   '         _______-----    ___--          \    /)\     \n'
                   '   ------_______      ---____            \__/  /     \n'
                   '                -----__    \ --    _          //\    \n'
                   '                       --__--__     \_____/   \_/\   \n'
                   '                               ----|   /          |  \n'
                   ' Dynamic                           |  |___________|  \n'
                   ' IDA                               |  | ((_(_)| )_)  \n'
                   ' Enrichment                        |  \_((_(_)|/(_)  \n'
                   '                                   \             (   \n'
                   '                                    \_____________)  \n'
                   ' D.I.E v0.1 is now loaded, enjoy.                    \n'
                   '-----------------------------------------------------\n'
                   )


class die_plugin_t(plugin_t):
    flags = idaapi.PLUGIN_PROC
    comment = "Dynamic IDA Enrichment plugin (aka. DIE)"
    help = "Help is a matter of trust."
    wanted_name = "DIE"
    wanted_hotkey = ""

    def init(self):
        try:
            # For Debugging:
            #self.die_manager = DieManager(is_dbg_log=True, is_dbg_pause=False, is_dbg_profile=True)
            self.die_manager = DieManager()
            self.die_manager.add_menu_items()
            self.die_manager.show_logo()
            return idaapi.PLUGIN_KEEP

        except Exception as ex:
            idaapi.msg("Failed to initialize DIE. {}\n".format(ex))
            self.die_manager.del_menu_items()
            del self.die_manager
            idaapi.msg("Errors and fun!\n")
            return idaapi.PLUGIN_SKIP

    def run(self, arg):
        pass

    def term(self):
        with ignored(AttributeError):
            if not self.die_manager.die_db.is_saved:
                response = idc.AskYN(1, "DIE DB was not saved, Would you like to save it now?")
                if response == 1:
                    self.die_manager.save_db()

            self.die_manager.del_menu_items()


def PLUGIN_ENTRY():
    return die_plugin_t()


repository: tmr232/Sark
file: plugins/autostruct.py
description: IDAPython Made Easy
url: https://github.com/tmr232/Sark/blob/950e2c17e520ab637a3cdcd0da312d224b951f13/plugins/autostruct.py
wanted_name: AutoStruct
comment: AutoStruct struct creator
created_at: 2015-01-19 09:13:48+00:00
pushed_at: 2025-02-22 10:17:16+00:00
content:
import idaapi
import idc

import sark


def message(*messages):
    for msg in messages:
        for line in msg.splitlines():
            idaapi.msg("[Autostruct] {}\n".format(line))


class AutoStruct(idaapi.plugin_t):
    flags = idaapi.PLUGIN_PROC
    comment = "AutoStruct struct creator"
    help = "Automagically Create and Apply Structs"
    wanted_name = "AutoStruct"
    wanted_hotkey = "Shift+T"

    def init(self):
        self._prev_struct_name = ""
        return idaapi.PLUGIN_KEEP

    def term(self):
        pass

    def run(self, arg):
        start, end = sark.get_selection()

        if not sark.structure.selection_has_offsets(start, end):
            message('No structure offsets in selection. Operation cancelled.')
            idaapi.warning('No structure offsets in selection. Operation cancelled.')
            return

        struct_name = idaapi.ask_str(self._prev_struct_name, 0, "Struct Name")
        if not struct_name:
            message("No structure name provided. Operation cancelled.")
            return
        self._prev_struct_name = struct_name

        common_reg = sark.structure.get_common_register(start, end)
        reg_name = idaapi.ask_str(common_reg, 0, "Register")
        if not reg_name:
            message("No offsets found. Operation cancelled.")
            return

        try:
            offsets, operands = sark.structure.infer_struct_offsets(start, end, reg_name)
        except sark.exceptions.InvalidStructOffset:
            message("Invalid offset found. Cannot create structure.",
                    "Make sure there are no negative offsets in the selection.")
            return

        except sark.exceptions.SarkInvalidRegisterName:
            message("Invalid register name {!r}. Cannot create structs.".format(reg_name))
            return

        try:
            sark.structure.create_struct_from_offsets(struct_name, offsets)
        except sark.exceptions.SarkStructAlreadyExists:
            yes_no_cancel = idaapi.ask_yn(idaapi.ASKBTN_NO,
                                      "Struct already exists. Modify?\n"
                                      "Cancel to avoid applying the struct.")
            if yes_no_cancel == idaapi.ASKBTN_CANCEL:
                return

            elif yes_no_cancel == idaapi.ASKBTN_YES:
                sid = sark.structure.get_struct(struct_name)
                sark.structure.set_struct_offsets(offsets, sid)

            else:  # yes_no_cancel == idaapi.ASKBTN_NO:
                pass

        sark.structure.apply_struct(start, end, reg_name, struct_name)


def PLUGIN_ENTRY():
    return AutoStruct()


repository: angr/patcherex
file: ida_plugin/patcherex_ida_main.py
description: Shellphish's automated patching engine, originally created for the Cyber Grand Challenge.
wanted_name: Patcherex
comment: Patcherex plugin for IDA
created_at: 2016-08-20 19:15:54+00:00
pushed_at: 2025-03-21 02:40:26+00:00
url:
https://github.com/angr/patcherex/blob/b2a7b49453a24f67b8e0f7293a021e1ccdbd192b/ida_plugin/patcherex_ida_main.py

content:
#pylint: disable=import-error,wildcard-import,unused-wildcard-import,no-self-use,unused-argument
import idaapi
from idaapi import Choose2
import json
import re
import subprocess
import threading
import tempfile
import shlex
import os
from patcherex_ida.forms import *
from collections import namedtuple


PATCH_TYPES = {"InsertCodePatch": {"desc": "Insert Assembly", "handler_form": InsertCodePatchForm},
               "AddRODataPatch": {"desc": "Insert Read Only Data", "handler_form": AddRODataPatchForm},
               "AddLabelPatch": {"desc": "Add a Label", "handler_form": AddLabelPatchForm},
               "RemoveInstructionPatch": {"desc": "Remove an instruction", "handler_form": RemoveInstructionPatchForm},
               "AddRWDataPatch": {"desc": "Add Space for Read/Write Data", "handler_form": AddRWDataPatchForm},
               "AddCodePatch": {"desc": "Add assembly or C code into the binary", "handler_form": AddCodePatchForm}}

class ItemManager(object):
    def __init__(self, get_items, set_items):
        self.get_items = get_items
        self.set_items = set_items
        self.internal_items = None
        self.internal_items = self.load_serialized(self.get_items())

    def pre_sync(self):
        self._load_serialized(self.get_items())

    def post_sync(self):
        s = self._get_serialized()
        self.set_items(s)

    def __getitem__(self, index):
        return self._get_handler_for_item(index).get_gui_format_of(**self.get_item(index))

    def __len__(self):
        return len(self.internal_items)

    def check_name_collisions(self, name, replacement_index):
        for i, item in enumerate(self.internal_items):
            if (replacement_index is None) or (replacement_index != i):
                if item["name"] == name:
                    return item["name"]
        return False

    def gen_valid_name(self, start_name, replacement_index=None):
        alt_name = start_name
        while True:
            conflict = self.check_name_collisions(alt_name, replacement_index=replacement_index)
            if conflict:
                match = re.search(r"_[0-9]+\Z", conflict)
                if match:
                    num = int(conflict[match.start() + 1:match.end()])
                    base_str = conflict[:match.start()]
                else:
                    num = 0
                    base_str = conflict
                alt_name = base_str + "_" + str(num + 1)
            else:
                break
        return alt_name

    def add_item(self, patch_type, address, name, data):
        self.pre_sync()
        self._add_item(patch_type, address, name, data)
        self.post_sync()

    def _add_item(self, patch_type, address, name, data):
        valid_name = self.gen_valid_name(name)
        self.internal_items.append({"patch_type": patch_type,
                                    "address": address,
                                    "name": valid_name,
                                    "data": data})
        new_index = len(self.internal_items) - 1
        self._get_handler_for_item(new_index).on_perform_post_operations(**self._get_item(new_index))
        return len(self.internal_items) - 1

    def get_item(self, index):
        self.pre_sync()
        return self._get_item(index)

    def _get_item(self, index):
        return self.internal_items[index]

    def set_item(self, index, updates):
        self.pre_sync()
        self._set_item(index, updates)
        self.post_sync()

    def _set_item(self, index, updates):
        if "name" in updates:
            updates = dict(updates)
            updates["name"] = self.gen_valid_name(updates["name"], replacement_index=index)
        self._get_handler_for_item(index).on_pre_update(**self._get_item(index))
        self.internal_items[index].update(updates)
        self._get_handler_for_item(index).on_post_update(**self._get_item(index))

    def delete_item(self, index):
        self.pre_sync()
        self._delete_item(index)
        self.post_sync()
        return index

    def _delete_item(self, index):
        self._get_handler_for_item(index).on_delete(**self._get_item(index))
        del self.internal_items[index]

    def _get_handler_for_item(self, n):
        return PATCH_TYPES[self._get_item(n)["patch_type"]]["handler_form"]

    def _get_serialized(self):
        return json.dumps(self.internal_items)

    def _load_serialized(self, contents):
        if self.internal_items is not None:
            self._uninitialize_patches()
        self.internal_items = []
        stuff = json.loads(contents) if contents else None
        if stuff is not None:
            for item in stuff:
                self._add_item(**item)

    def get_serialized(self):
        self.pre_sync()
        return self._get_serialized()

    def load_serialized(self, contents):
        self._load_serialized(contents)
        self.post_sync()

    def uninitialize_patches(self):
        self.pre_sync()
        self._uninitialize_patches()

    def _uninitialize_patches(self):
        while len(self.internal_items) != 0:
            self._delete_item(0)

class PatcherexWindow(Choose2):
    def __init__(self):
        Choose2.__init__(
            self,
            "Patcherex",
            [
                ["Type", 10 | Choose2.CHCOL_PLAIN],
                ["Address", 10 | Choose2.CHCOL_HEX],
                ["Name", 30 | Choose2.CHCOL_PLAIN],
                ["Data", 30 | Choose2.CHCOL_FORMAT]
            ])
        self.node = idaapi.netnode()
        self.node.create("$ patcherex")
        self.items = ItemManager(lambda: self.node.getblob(0, "I"),
                                 lambda x: self.node.setblob(x, 0, "I"))
        self.popup_names = ["Add Patch", "Remove Patch", "Edit Patch", "Refresh"]

    def load_saved_patches(self):
        SaveLoadHook.loading()

    def OnClose(self):
        pass

    def OnEditLine(self, n):
        params = dict(self.items.get_item(n))
        patch_type = params["patch_type"]
        del params["patch_type"]
        form = PATCH_TYPES[patch_type]["handler_form"](**params)
        form.Compile()
        if form.Execute():
            updates = {}
            updates["address"] = form.get_patch_address()
            updates["name"] = form.get_patch_name()
            updates["data"] = form.get_patch_data()
            self.items.set_item(n, updates)
        form.Free()
        self.Refresh()

    def OnInsertLine(self):
        form = PatchTypeForm(PATCH_TYPES)
        form.Compile()
        if form.Execute():
            option = form.get_chosen_patch_type()
            patch_form = PATCH_TYPES[option]["handler_form"]()
            patch_form.Compile()
            if patch_form.Execute():
                params = {}
                params["patch_type"] = option
                params["name"] = patch_form.get_patch_name()
                params["address"] = patch_form.get_patch_address()
                params["data"] = patch_form.get_patch_data()
                self.items.add_item(**params)
        form.Free()
        self.Refresh()

    def OnSelectLine(self, n):
        self.OnEditLine(n)

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

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

    def OnDeleteLine(self, n):
        return self.items.delete_item(n)

    def OnRefresh(self, n):
        return n

    def OnGetIcon(self, n):
        return 5

    def OnGetLineAttr(self, n): # Can be used to set line color, e.g. blue = [0xFF, 0]
        return None

    def get_serialized_items(self):
        return self.items.get_serialized()

    def load_serialized_items(self, contents):
        self.items.load_serialized(contents)
        self.Refresh()

    def uninitialize_patches(self):
        self.items.uninitialize_patches()

    def is_used(self):
        return len(self.items) != 0

class UnsupportedArchitectureException(Exception):
    pass

class SaveHandler(idaapi.action_handler_t):
    def activate(self, ctx):
        form = SaveForm()
        form.Compile()
        if form.Execute():
            file_name = form.get_file_name()
            self.do_save(file_name)
        form.Free()

    @classmethod
    def do_save(cls, file_name):
        contents = patcherex_window.get_serialized_items()
        if file_name is not None:
            try:
                out_file = open(file_name, "wb")
            except IOError as e:
                idaapi.warning("Unable to open %s for saving (%s)" % (file_name, e.strerror))
            else:
                with out_file:
                    out_file.write(contents)
        else:
            patcherex_window.node.setblob(contents, 0, "I")

    def update(self, ctx):
        return idaapi.AST_ENABLE_FOR_FORM if ctx.widget_title == "Patcherex" else idaapi.AST_DISABLE_FOR_FORM

class LoadHandler(idaapi.action_handler_t):
    def activate(self, ctx):
        form = LoadForm()
        form.Compile()
        if form.Execute():
            file_name = form.get_file_name()
            self.do_load(file_name)
        form.Free()

    @classmethod
    def do_load(cls, file_name):
        if file_name is not None:
            try:
                out_file = open(file_name, "rb")
            except IOError as e:
                idaapi.warning("Unable to open %s for loading (%s)" % (file_name, e.strerror))
            else:
                with out_file:
                    contents = out_file.read()
        else:
            contents = patcherex_window.node.getblob(0, "I")
            if contents is None:
                return
        patcherex_window.load_serialized_items(contents)

    def update(self, ctx):
        return idaapi.AST_ENABLE_FOR_FORM if ctx.widget_title == "Patcherex" else idaapi.AST_DISABLE_FOR_FORM

class RunPatcherexHandler(idaapi.action_handler_t):
    config = """{"techniques": {"manualpatcher": {"options": {"patch_file": null}}}, "backend": {"name": "detourbackend", "options": {"base_address": 0}}}"""
    def activate(self, ctx):
        RunPatcherexHandler.menu_activate(None)

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

    @staticmethod
    def patcherex_finish(proc):
        if proc.returncode != 0:
            out_str = "Patcherex failed. See attached terminal."
            idaapi.warning(out_str)
            print out_str
        else:
            out_str = "Patcherex completed successfully."
            idaapi.info(out_str)
            print out_str

    @staticmethod
    def menu_activate(arg):
        form = RunPatcherexForm()
        form.Compile()
        if form.Execute():
            file_name = form.get_file_name()
            patch_output_file = tempfile.NamedTemporaryFile(delete=False)
            patch_output_file.write(patcherex_window.get_serialized_items())
            patch_output_file.close()
            config_file = tempfile.NamedTemporaryFile(delete=False)
            p_config = json.loads(RunPatcherexHandler.config)
            p_config["techniques"]["manualpatcher"]["options"]["patch_file"] = patch_output_file.name
            p_config["backend"]["options"]["base_address"] = idaapi.get_imagebase()
            config_file.write(json.dumps(p_config))
            config_file.close()
            print "Calling patcherex . . ."
            popen_and_call(RunPatcherexHandler.patcherex_finish,
                           {"args": ["patcherex", "-c", config_file.name, "--json", "single", idaapi.get_input_file_path(), file_name]})


class AddPatchHandler(idaapi.action_handler_t):
    def activate(self, ctx):
        AddPatchHandler.menu_activate(None)

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

    @staticmethod
    def menu_activate(arg):
        patcherex_window.OnInsertLine()

class PopHook(idaapi.UI_Hooks):
    def __init__(self, actname_list):
        idaapi.UI_Hooks.__init__(self)
        self.acts = actname_list

    def finish_populating_tform_popup(self, form, popup):
        if idaapi.get_tform_title(form) == "Patcherex":
            for act in self.acts:
                idaapi.attach_action_to_popup(form, popup, "patcherex:" + act, None)

class SaveLoadHook(idaapi.UI_Hooks):
    def saving(self):
        if patcherex_window.is_used():
            SaveHandler.do_save(None)

    # Not a real idaapi hook; called from PatcherexWindow.load_saved_patches
    @classmethod
    def loading(cls):
        LoadHandler.do_load(None)

    @staticmethod
    def get_preferred_path():
        idb_path = os.path.abspath(idc.GetIdbPath())
        no_idb_ext = os.path.splitext(os.path.basename(idb_path))[0]
        save_path = os.path.join(os.path.split(idb_path)[0], no_idb_ext + "_patcherex.json")
        return save_path

# http://stackoverflow.com/questions/2581817/python-subprocess-callback-when-cmd-exits
def popen_and_call(onExit, popenArgs):
    """
    Runs the given args in a subprocess.Popen, and then calls the function
    onExit when the subprocess completes.
    onExit is a callable object, and popenArgs is a list/tuple of args that
    would give to subprocess.Popen.
    """
    def runInThread(onExit, popenArgs):
        proc = subprocess.Popen(**popenArgs)
        proc.wait()
        onExit(proc)
        return
    thread = threading.Thread(target=runInThread, args=(onExit, popenArgs))
    thread.start()
    # returns immediately after the thread starts
    return thread

PatcherexCommand = namedtuple("PatcherexCommand", ["name",
                                                   "description",
                                                   "handler_class",
                                                   "shortcut",
                                                   "menu_path"])
commands = []

commands.append(PatcherexCommand(**{"name": "save",
                                    "description": "Save patches to file . . .",
                                    "handler_class": SaveHandler,
                                    "shortcut": None,
                                    "menu_path": None}))

commands.append(PatcherexCommand(**{"name": "load",
                                    "description": "Load patches from file . . .",
                                    "handler_class": LoadHandler,
                                    "shortcut": None,
                                    "menu_path": None}))

commands.append(PatcherexCommand(**{"name": "run_patcherex",
                                    "description": "Run Patcherex . . .",
                                    "handler_class": RunPatcherexHandler,
                                    "shortcut": "Ctrl-Shift-R",
                                    "menu_path": "File/Produce file/"}))

commands.append(PatcherexCommand(**{"name": "add_patch",
                                    "description": "Add Patcherex patch . . .",
                                    "handler_class": AddPatchHandler,
                                    "shortcut": "Ctrl-Shift-N",
                                    "menu_path": "Edit/Patch program/"}))

if globals().get("patcherex_window", None) is None:
    patcherex_window = None

if globals().get("patcherex_hooks", None) is None:
    patcherex_hooks = None

class PatcherexPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_KEEP
    comment = "Patcherex plugin for IDA"
    help = "Use Ctrl-Shift-P to access Patcherex"
    wanted_name = "Patcherex"
    wanted_hotkey = "Ctrl-Shift-P"

    def __init__(self):
        super(PatcherexPlugin, self).__init__()
        self.patcherex_window = None

    def init(self):
        global patcherex_window
        global patcherex_hooks
        if idaapi.get_inf_structure().procName != "metapc":
            print "Only x86 metapc is supported by Patcherex"
            return idaapi.PLUGIN_SKIP
        for command in commands:
            name = "patcherex:" + command.name
            idaapi.register_action(
                idaapi.action_desc_t(name,
                                     command.description,
                                     command.handler_class(),
                                     command.shortcut))
            if command.menu_path:
                idaapi.attach_action_to_menu(command.menu_path, name, idaapi.SETMENU_APP)
        if patcherex_hooks is None:
            patcherex_hooks = [PopHook([command.name for command in commands]), SaveLoadHook()]
            for hook in patcherex_hooks:
                hook.hook()
        if patcherex_window is None:
            print "Patcherex starting"
            self.patcherex_window = PatcherexWindow()
            patcherex_window = self.patcherex_window
        else:
            self.patcherex_window = patcherex_window
            patcherex_window.load_saved_patches()
        return idaapi.PLUGIN_OK

    def run(self, arg):
        self.patcherex_window.Show()

    def term(self):
        self.patcherex_window.uninitialize_patches()
        self.patcherex_window.Close()

def PLUGIN_ENTRY():
    return PatcherexPlugin()


repository: Chordp/PatternGen
file: PatternGen.py
url: https://github.com/Chordp/PatternGen/blob/0955142c1a6a287345de8579d07b6b5148f6ac32/PatternGen.py
wanted_name: PatternGen
comment: Pattern Generate tool by:chord
created_at: 2020-07-02 10:15:06+00:00
pushed_at: 2022-04-05 07:57:57+00:00
content:
# -*- coding: utf-8 -*-
import math

import idaapi
import idc
#import clipboard

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_FORM
            return idaapi.AST_DISABLE_FOR_FORM


    class Searcher(Kp_Menu_Context):
        def activate(self, ctx):
            self.plugin.search()
            return 1
except:
    pass

class PatternGen_Plugin_t(idaapi.plugin_t):
    comment = "Pattern Generate tool by:chord"
    help = "todo"
    wanted_name = "PatternGen"
    wanted_hotkey = "ALT+Z"
    flags = idaapi.PLUGIN_KEEP

    def init(self):
        try:
            print ("Pattern Generate tool by:chord")
            Searcher.register(self, "PatternGen")
        except:
            pass
        return idaapi.PLUGIN_KEEP

    def term(self):
        pass

    def printAvd(slef):
        print (25 * "===")

    def formatByte(self,ea):
        return " "+"{:02X}".format(idc.get_wide_byte(ea))

    def calcStr(self,ea, endcount):
        hstr = ""
        firstByte = self.formatByte(ea)
        hstr += self.formatByte(ea)
        hstr = hstr + self.formatByte(ea + 1) if (firstByte == "FF" or firstByte == "66" or firstByte == "67") else hstr
        #print(math.ceil(endcount - len(hstr) / 2))
        hstr = hstr + math.ceil(endcount - len(hstr) / 2) * " ??" if endcount >= 2 else hstr
        return hstr

    def extractCode(self):
        self.printAvd()

        start = idc.read_selection_start()
        end = idc.read_selection_end()
        codeSize = end - start
        ea = start
        # print hex(ea)
        result = ""

        for i in range(codeSize):
            op1 = idc.get_operand_type(ea, 0)
            op2 = idc.get_operand_type(ea, 1)
            instructionSize = idc.get_item_size(ea)

            if op1 == idc.o_reg and (op2 == idc.o_reg or op2 == idc.o_void or op2 == idc.o_phrase):
                for b in range(0, instructionSize):
                    result += self.formatByte(ea + b)
            elif (op1 == idc.o_reg and op2 == idc.o_displ) or (op1 == idc.o_displ and op2 == idc.o_reg) or (
                    op1 == idc.o_displ and op2 == idc.o_imm):
                result += self.formatByte(ea) + self.formatByte(ea + 1)
                for b in range(2, instructionSize):
                    result = result + " ??"
            elif op1 == idc.o_phrase and op2 == idc.o_reg:
                for b in range(0, instructionSize):
                    result += self.formatByte(ea + b)
            else:
                result += self.calcStr(ea, instructionSize)

            ea = ea + instructionSize
            if ea >= (start + codeSize):
                break
        # print (idc.get_event_module_base() -  idc.SelStart());
        print ("%s  Address:0x%x Offset:0x%x" % (idc.get_func_off_str(idc.here()),idc.here(), idc.here() - idaapi.get_imagebase()))
        # print result
        return result

    def run(self, arg):
        if (idc.BADADDR != idc.here()):
            copyContent = self.extractCode();
            print(copyContent)
            # clipboard.copy(copyContent)


# register IDA plugin
def PLUGIN_ENTRY():
    return PatternGen_Plugin_t();


repository: zhkl0228/AndroidAttacher
file: src/AndroidAttacher.py
description: IDA debugging plugin for android armv7 so
wanted_name: Android Debug Attach
created_at: 2017-09-15 19:22:57+00:00
pushed_at: 2018-10-04 07:59:31+00:00
url:
https://github.com/zhkl0228/AndroidAttacher/blob/e9e628c252323c506a9ebb82142b73734e2832fb/src/AndroidAttacher.py

content:
#!/usr/bin/env python
# ! -*- coding: utf-8 -*-

import os

ADB_PATH = "~/Library/Android/sdk/platform-tools/adb"


def get_plugin_home():
    import inspect

    plugin_path = os.path.abspath(inspect.getfile(inspect.currentframe()))
    if os.path.islink(plugin_path):
        plugin_path = os.readlink(plugin_path)

    return os.path.dirname(plugin_path)


import sys

sys.path.append(get_plugin_home())

import idaapi
import idc

from aaf import utils


class android_attacher_plugin(idaapi.plugin_t):
    ACTION_NAME = "Attach to Android app"

    flags = idaapi.PLUGIN_KEEP
    comment = ""
    wanted_name = "Android Debug Attach"
    wanted_hotkey = "Alt-F8"
    help = wanted_name + ": Debugger/" + ACTION_NAME

    def __init__(self):
        self.androidAttacher = None
        self.attaching = False
        pass

    def init(self):
        architecture = utils.getIdaArchitecture()
        if architecture != "arm":
            print "%s unsupported architecture: %s" % (self.wanted_name, architecture)
            return idaapi.PLUGIN_SKIP

        idc.Message("Initializing %s\n" % self.wanted_name)

        from aaf import adb
        wrapper = adb.AdbWrapper(ADB_PATH)

        from aaf import AndroidAttacher
        utilsJar = os.path.join(get_plugin_home(), "aaf", "utils.jar")
        config_file = os.path.splitext(idc.GetIdbPath())[0] + ".aaf.conf"
        self.androidAttacher = AndroidAttacher(wrapper, utilsJar, config_file)
        return idaapi.PLUGIN_KEEP

    def term(self):
        pass

    def run(self, arg=0):
        if self.attaching:
            return

        try:
            if self.androidAttacher is not None:
                self.attaching = True
                self.androidAttacher.attach(arg)
        finally:
            self.attaching = False


def PLUGIN_ENTRY():
    return android_attacher_plugin()


repository: EobOrz/FFxivFuckGreenPlayer
file: DevTools/[IDATools]NySigWorker.py
description: Cracked Siren_Lite leaked source code, Consent has been obtained.
wanted_name: NySigFinder
comment: sig finder
created_at: 2025-03-05 12:18:14+00:00
pushed_at: 2025-03-13 05:43:42+00:00
url:
https://github.com/EobOrz/FFxivFuckGreenPlayer/blob/58682e72e2c1ac51544cc94b5bb2a08af9d1b939/DevTools/[IDATools]NySigWorker.py

content:
from idc import *
from idaapi import *
from idautils import *
import tkinter as tk

min_ea = inf_get_min_ea()
max_ea = inf_get_max_ea()
hex_char = [f'{i:02x} ' for i in range(256)] + ['* ', '? ']


# region make sig

def ea_point(ea):
    return next(v for v in (get_operand_value(ea, i) for i in range(2)) if min_ea < v < max_ea)


def line_sig(ea, mask_op=False):
    decode_insn(insn := ida_ua.insn_t(), ea)
    size = insn.size
    line = get_bytes(ea, size)
    offset_op = next((max(op.offb, 1) for op in insn.ops), len(line))
    if mask_op:
        try:
            mask_idx = next(i for i in range(2) if min_ea < get_operand_value(ea, i) < max_ea)
        except StopIteration:
            print(f"[!]cant find operand at {ea:x}, with operand value {get_operand_value(ea, 0):x} {get_operand_value(ea, 1):x}")
            raise
        mask_start = insn.ops[mask_idx].offb
        mask_end = insn.ops[mask_idx + 1].offb or size
    else:
        mask_start = mask_end = size
    res = "".join(hex_char[b] for b in line[:offset_op])
    for i in range(offset_op, size):
        res += '* ' if mask_start <= i < mask_end else '? '
    return res


def sig_search(sig: str, max_search_cnt: int = 10):
    addr = min_ea
    while max_search_cnt > 0:
        addr = find_binary(addr, max_ea, sig, 16, SEARCH_DOWN | SEARCH_NEXT)
        if addr == BADADDR: break
        yield addr
        max_search_cnt -= 1


def find_xrefs(ea: int) -> dict:
    xrefs = {}
    for xref in XrefsTo(ea, 0):
        xrefs.setdefault(xref.type, []).append(xref.frm)
    return xrefs


fl_name = {
    dr_O: 'dr_O',
    dr_I: 'dr_I',
    dr_R: 'dr_R',
    dr_T: 'dr_T',
    dr_W: 'dr_W',
    dr_S: 'dr_S',
    fl_CF: 'fl_CF',
    fl_CN: 'fl_CN',
    fl_JF: 'fl_JF',
    fl_JN: 'fl_JN',
    fl_F: 'fl_F',
}


def is_only_ref_from(from_, to):
    found = False
    for xref in XrefsTo(to, 0):
        if xref.frm == from_:
            found = True
        elif xref.type in (ida_xref.fl_CF, ida_xref.fl_CN, ida_xref.fl_JF, ida_xref.fl_JN, ida_xref.fl_F):
            return False
    return found


def is_only_ref_to(from_, to):
    found = False
    import ida_xref
    for xref in XrefsFrom(from_, 0):
        # print(xref, fl_name.get(xref.type), hex(xref.to),xref.iscode,xref.user)
        if xref.to == to:
            found = True
        elif xref.type in (ida_xref.fl_CF, ida_xref.fl_JF, ida_xref.fl_JN, ida_xref.fl_F):
            return False
    return found


class SigSearch:
    SUCCESS = 0
    FAIL = -1
    FORWARD = 1
    BACKWARD = 2
    PARENT = 3

    def __init__(self, start_ea, mgr: 'SigSearcher'):
        self.mgr = mgr
        self.start_ea = start_ea
        self.current_ea = start_ea
        if not (func := get_func(start_ea)):
            raise ValueError(f"no function at {start_ea:x}")
        self.func_range = range(func.start_ea, func.end_ea)
        self.state = SigSearch.FORWARD
        self.sigs = [line_sig(self.current_ea)]
        self.offset = 0
        self.inst_offset = 0

    def walk(self):
        if self.state == SigSearch.FORWARD:
            self.walk_forward()
        if self.state == SigSearch.BACKWARD:
            self.walk_backward()
        self.test()

    def walk_forward(self):
        next_ea = next_head(self.current_ea, self.func_range.stop)
        if next_ea not in self.func_range or not is_only_ref_from(self.current_ea, next_ea):
            self.state = SigSearch.BACKWARD
            self.current_ea = self.start_ea
            return
        self.sigs.append(line_sig(next_ea))
        self.current_ea = next_ea

    def walk_backward(self):
        next_ea = prev_head(self.current_ea, self.func_range.start)
        if next_ea not in self.func_range or not is_only_ref_to(next_ea, self.current_ea):
            # print(f"[!]fail at {self.current_ea:x} to {next_ea:x} {next_ea not in self.func_range=} {not is_only_ref_to(next_ea, self.current_ea)=}")
            self.state = SigSearch.FAIL
            return
        self.sigs.insert(0, line_sig(next_ea))
        self.offset += self.current_ea - next_ea
        self.inst_offset += 1
        self.current_ea = next_ea

    sig = property(lambda self: "".join(self.sigs))

    def test(self):
        if self.state == SigSearch.FAIL or self.state == SigSearch.SUCCESS: return
        matches = list(sig_search(self.sig, 2))
        if len(matches) == 1:
            self.state = SigSearch.SUCCESS
        elif len(matches) == 0:
            self.state = SigSearch.FAIL  # should not happen

    def count_match(self, max_=10):
        return sum(1 for _ in sig_search(self.sig, max_))

    def make(self, mask_addr=False):
        count_match = self.count_match()
        if count_match < 10:
            while len(self.sigs) - 1 > self.inst_offset:
                if sum(1 for _ in sig_search("".join(self.sigs[:-1]), count_match + 1)) == count_match:
                    self.sigs.pop()
                else:
                    break
        if mask_addr:
            self.sigs[self.inst_offset] = line_sig(self.start_ea, True)
        res = self.sig
        if self.offset:
            res += f' +{self.offset:#x}'
        return res


class SigSearcher:
    def __init__(self, target_ea, max_workers=500):
        self.target_ea = target_ea
        self.max_workers = max_workers
        self.workers = []
        if get_func(target_ea):
            self.workers.append(SigSearch(target_ea, self))
        for xref in XrefsTo(target_ea, 0):
            if xref.type in (fl_CF, fl_CN, fl_JF, fl_JN, dr_R, dr_W, dr_T):
                try:
                    worker = SigSearch(xref.frm, self)
                except ValueError:
                    continue
                if not self.append(worker): break

        self.found = []
        self.failed = []

    def append(self, worker):
        if len(self.workers) < self.max_workers:
            self.workers.append(worker)
            return True
        return False

    def walk(self):
        for worker in self.workers:
            if worker.state > 0:
                worker.walk()
            if worker.state <= 0:
                self.workers.remove(worker)
                if worker.state == SigSearch.SUCCESS:
                    self.found.append(worker)
                else:
                    self.failed.append(worker)

    def auto_walk(self, max_found=10, max_step=50) -> tuple[int, list[SigSearch]]:
        if not self.workers:
            raise ValueError("no workers")
        for _ in range(max_step):
            self.walk()
            if len(self.workers) == 0: break
            if len(self.found) >= max_found: break
        if not self.found:
            sorted_workers = sorted(self.failed + self.workers, key=lambda w: w.count_match())
            min_match = sorted_workers[0].count_match()
            workers = []
            for worker in sorted_workers:
                if worker.count_match() == min_match:
                    workers.append(worker)
                else:
                    break
        else:
            min_match = 1
            workers = self.found
        return min_match, workers


def make_sig(target):
    try:
        searcher = SigSearcher(target)
        print(f"[.]start search {target:x} {get_name(target)} in {len(searcher.workers)} ea")
        min_match, found = searcher.auto_walk()
        print(f"[+]found {len(found)} sig with {min_match} min match")
        for worker in found:
            try:
                print(f"[+] {worker.start_ea:x} : {worker.make(worker.start_ea != target)}")
            except Exception as e:
                print(f"[!]error when make sig {worker.start_ea:x}: {e}")
                print(traceback.format_exc())
    except Exception as e:
        print(f"[!]error when search {target:x} {get_name(target)}: {e}")
        print(traceback.format_exc())


# endregion

# region search sig

def get_exec_data():
    import sys
    import types
    if '__exec_cache__' not in sys.modules:
        sys.modules['__exec_cache__'] = m = types.ModuleType('__exec_cache__')
        m.data = get_bytes(inf_get_min_ea(), inf_get_max_ea() - inf_get_min_ea())
    return sys.modules['__exec_cache__'].data


# xx xx xx xx 空格分割
# [xx:yy] 单字节从 xx 到 yy
# [xx|yy|zz] 单字节 xx 或 yy 或 zz
# ? ? ? ? 视作变量（默认不储存）
# ^ ^ ^ ^ 视作字串（默认储存）
# * * * * 视作跳转（默认储存）
# ?{n} / *{n} 视作匹配n次
# ?{n:m} / *{n:m} 视作匹配n-m次
# (xx xx xx xx) 不存储的分组
# <xx xx xx xx> 储存的分组
# <* * * *: yy yy yy yy> 对分组数据二级匹配
# <* * * *: yy yy yy yy <* * * *:zz zz zz zz>> 对分组数据多级匹配，仅适用于跳转
import io
import re
import typing

fl_is_ref = 1 << 0
fl_is_byes = 1 << 1
fl_store = 1 << 2

hex_chars = set(b'0123456789abcdefABCDEF')
dec_chars = set(b'0123456789')

special_chars_map = {i for i in b'()[]{}?*+-|^$\\.&~# \t\n\r\v\f'}


def take_dec_number(pattern: str, i: int):
    assert i < len(pattern) and ord(pattern[i]) in dec_chars
    j = i + 1
    while j < len(pattern) and ord(pattern[j]) in dec_chars:
        j += 1
    return int(pattern[i:j]), j


def take_cnt(pattern: str, i: int, regex_pattern: bytearray):
    if i < len(pattern) and pattern[i] == '{':
        regex_pattern.append(123)  # {
        n1, i = take_dec_number(pattern, i + 1)
        regex_pattern.extend(str(n1).encode())
        if pattern[i] == ':':
            n2, i = take_dec_number(pattern, i + 1)
            assert n1 <= n2
            regex_pattern.append(44)  # ,
            regex_pattern.extend(str(n2).encode())
        assert pattern[i] == '}'
        regex_pattern.append(125)  # }
        i += 1
    return i


def take_byte(pattern: str, i: int, regex_pattern: bytearray):
    assert i + 2 <= len(pattern)
    next_byte = int(pattern[i:i + 2], 16)
    if next_byte in special_chars_map:
        regex_pattern.append(92)  # \
    regex_pattern.append(next_byte)
    return i + 2


def _take_unk(pattern: str, i: int):
    start_chr = pattern[i]
    assert start_chr in ('?', '*', '^')
    if i + 1 < len(pattern) and pattern[i + 1] == start_chr:
        i += 1
    return start_chr, i + 1


def take_unk(pattern: str, i: int, regex_pattern: bytearray):
    start_unk, i = _take_unk(pattern, i)
    regex_pattern.append(46)
    i = take_cnt(pattern, i, regex_pattern)
    while i < len(pattern):
        match pattern[i]:
            case ' ':
                i += 1
            case c if c == start_unk:
                start_unk, i = _take_unk(pattern, i)
                regex_pattern.append(46)
                i = take_cnt(pattern, i, regex_pattern)
            case _:
                break
    return start_unk, i


def _compile_pattern(pattern: str, i=0, ret_at=None):
    _i = i
    regex_pattern = bytearray()
    sub_matches = []
    group_flags = []
    while i < len(pattern):
        match pattern[i]:
            case ' ':
                i += 1
            case '[':
                regex_pattern.append(91)  # [
                i += 1
                i = take_byte(pattern, i, regex_pattern)
                while True:
                    match pattern[i]:
                        case ' ':
                            i += 1
                        case ']':
                            regex_pattern.append(93)  # ]
                            i += 1
                            break
                        case '|':
                            i = take_byte(pattern, i + 1, regex_pattern)
                        case ':':
                            regex_pattern.append(45)  # -
                            i = take_byte(pattern, i + 1, regex_pattern)
                        case c:
                            raise ValueError(f'Invalid character {c} in pattern {pattern!r} at {i}')

            case '(':
                base_flag = 0  # not fl_store
                regex_pattern.append(40)  # (
                unk_type, i = take_unk(pattern, i + 1, regex_pattern)
                if unk_type == '*':
                    base_flag |= fl_is_ref
                elif unk_type == '^':
                    base_flag |= fl_is_byes
                sub_pattern = None
                while True:
                    match pattern[i]:
                        case ' ':
                            i += 1
                        case ')':
                            regex_pattern.append(41)  # )
                            i += 1
                            break
                        case ':':
                            sub_pattern, i = _compile_pattern(pattern, i + 1, ret_at=')')
                            assert pattern[i] == ')', f'Expected ) get {pattern[i]} at {i} in pattern {pattern!r}'
                            regex_pattern.append(41)
                            i += 1
                            break
                        case c:
                            raise ValueError(f'Invalid character {c} in pattern {pattern!r} at {i}')
                group_flags.append(base_flag)
                sub_matches.append(sub_pattern)
            case '<':
                base_flag = fl_store
                regex_pattern.append(40)
                unk_type, i = take_unk(pattern, i + 1, regex_pattern)
                if unk_type == '*':
                    base_flag |= fl_is_ref
                elif unk_type == '^':
                    base_flag |= fl_is_byes
                sub_pattern = None
                while True:
                    match pattern[i]:
                        case ' ':
                            i += 1
                        case '>':
                            regex_pattern.append(41)
                            i += 1
                            break
                        case ':':
                            sub_pattern, i = _compile_pattern(pattern, i + 1, ret_at='>')
                            assert pattern[i] == '>', f'Expected > get {pattern[i]} at {i} in pattern {pattern!r}'
                            regex_pattern.append(41)
                            i += 1
                            break
                        case c:
                            raise ValueError(f'Invalid character {c} in pattern {pattern!r} at {i}')
                group_flags.append(base_flag)
                sub_matches.append(sub_pattern)
            case '?' | '*' | '^' as c:
                regex_pattern.append(40)
                unk_type, i = take_unk(pattern, i, regex_pattern)
                regex_pattern.append(41)
                if c == '?':
                    group_flags.append(0)
                elif c == '*':
                    group_flags.append(fl_is_ref | fl_store)
                elif c == '^':
                    group_flags.append(fl_is_byes | fl_store)
                else:
                    raise ValueError(f'Invalid character {c} in pattern {pattern!r} at {i}')
                sub_matches.append(None)
            case c if ord(c) in hex_chars:
                i = take_byte(pattern, i, regex_pattern)
                i = take_cnt(pattern, i, regex_pattern)
            case c if c == ret_at:
                break
            case c:
                fmt_pattern = pattern[:i] + '_' + pattern[i] + '_' + pattern[i + 1:]
                raise ValueError(f'Invalid character {c} in pattern {fmt_pattern!r} at {i} (ret_at={ret_at})')
    try:
        regex = re.compile(bytes(regex_pattern), re.DOTALL)
    except re.error as e:
        raise ValueError(f'{e}: ({pattern!r}, {_i}, {ret_at!r}) -> {bytes(regex_pattern)}')
    return Pattern(regex, sub_matches, group_flags, pattern), i


def compile_pattern(pattern: str):
    return _compile_pattern(pattern)[0]


class Pattern:
    def __init__(self, regex: re.Pattern, sub_matches: 'typing.List[None | Pattern]', group_flags, pattern: str):
        self.regex = regex
        self.sub_matches = sub_matches
        self.group_flags = group_flags
        self.pattern = pattern
        self.res_is_ref = []
        for i, (sub, flag) in enumerate(zip(sub_matches, group_flags)):
            if flag & fl_store:
                self.res_is_ref.append(flag & fl_is_ref)
            if sub is not None:
                self.res_is_ref.extend(sub.res_is_ref)

    def finditer(self, _data: bytes | bytearray | memoryview, ref_base=0):
        data = _data if isinstance(_data, memoryview) else memoryview(_data)
        for match in self.regex.finditer(data):
            res = []
            if self._parse_match(data, match, res, ref_base):
                yield match.start(0), res

    def _parse_match(self, data: memoryview, match: re.Match, res: list, ref_base=0):
        for i, (sub_match, flag) in enumerate(zip(self.sub_matches, self.group_flags)):
            if flag & fl_is_byes:
                res.append(match.group(i + 1))
            else:
                val = int.from_bytes(match.group(i + 1), 'little', signed=True)
                if flag & fl_is_ref:
                    val += match.end(i + 1)
                if flag & fl_store:
                    res.append(val)
                if sub_match is not None:
                    start = val if flag & fl_is_ref else val - ref_base
                    if start < 0 or start >= len(data):
                        return False
                    if not sub_match._match(data, start, res, ref_base):
                        return False
        return True

    def _match(self, _data: memoryview, start_at: int, res: list, ref_base=0):
        if not (match := self.regex.match(_data, start_at)): return False
        return self._parse_match(_data, match, res, ref_base)

    def fmt(self, ind: str | int = ' ', _ind=0):
        if isinstance(ind, int): ind = ' ' * ind
        s = io.StringIO()
        s.write(ind * _ind)
        s.write(fmt_bytes_regex_pattern(self.regex.pattern))
        s.write('\n')
        s.write(ind * _ind)
        s.write('res is ref:')
        for flag in self.res_is_ref:
            s.write(' ref' if flag else ' val')
        s.write('\n')
        for i, (sub, flag) in enumerate(zip(self.sub_matches, self.group_flags)):
            s.write(ind * _ind)
            s.write(f'{i}:{"ref" if flag & fl_is_ref else "val"}{" store" if flag & fl_store else ""}\n')
            if sub is not None:
                s.write(sub.fmt(ind, _ind + 1))
                s.write('\n')
        return s.getvalue().rstrip()


def fmt_bytes_regex_pattern(pat: bytes):
    s = ''
    is_escape = False
    is_in_bracket = 0
    for b in pat:
        if is_escape:
            is_escape = False
            s += f'\\x{b:02x}'
        elif b == 92:  # \
            is_escape = True
        elif b in special_chars_map:
            if b == 123:  # {
                is_in_bracket += 1
            elif b == 125:  # }
                is_in_bracket -= 1
            s += chr(b)
        elif is_in_bracket:
            s += chr(b)
        else:
            s += f'\\x{b:02x}'
    return s


class IPatternScanner:
    def search(self, pattern: str | Pattern) -> typing.Generator[tuple[int, list[int]], None, None]:
        raise NotImplementedError

    def search_unique(self, pattern: str | Pattern) -> tuple[int, list[int]]:
        s = self.search(pattern)
        try:
            res = next(s)
        except StopIteration:
            raise KeyError('pattern not found')
        try:
            next(s)
        except StopIteration:
            return res
        raise KeyError('pattern is not unique, at least 2 is found')

    def find_addresses(self, pattern: str | Pattern):
        for address, _ in self.search(pattern):
            yield address

    def find_vals(self, pattern: str | Pattern):
        for address, args in self.search(pattern):
            yield args

    def find_address(self, pattern: str | Pattern):
        return self.search_unique(pattern)[0]

    def find_val(self, pattern: str | Pattern):
        return self.search_unique(pattern)[1]


class IdaScanner(IPatternScanner):
    def __init__(self):
        self.data = get_exec_data()
        self.min_ea = inf_get_min_ea()

    def search(self, pattern: str | Pattern) -> typing.Generator[tuple[int, list[int]], None, None]:
        if isinstance(pattern, str):
            pattern = compile_pattern(pattern)
        for address, args in pattern.finditer(self.data):
            yield self.min_ea + address, [a + self.min_ea if r else a for a, r in zip(args, pattern.res_is_ref)]


def search_sig(sig: str):
    res = list(IdaScanner().search(sig))
    print(f"[+] {len(res)} found for sig {sig}")
    for ea, args in res:
        print(f"[+] {ea:#x}: " + ", ".join(f"{a:#x}" for a in args))


# endregion

def main():
    def call_then_close(func, *args):
        def f():
            try:
                func(*args)
            finally:
                root.destroy()

        return f

    root = tk.Tk()
    root.title("NySigWorker")

    frame_func1 = tk.LabelFrame(root, text="make sig", padx=5, pady=5)
    frame_func2 = tk.LabelFrame(root, text="search sig", padx=5, pady=5)
    frame_buttons = tk.Frame(root, padx=5, pady=5)

    frame_func1.pack(padx=10, pady=5, fill="x")
    frame_func2.pack(padx=10, pady=5, fill="x")
    frame_buttons.pack(padx=10, pady=5, fill="x")

    # Components for func1
    suggestion = []
    if (_here := here()) != BADADDR:
        suggestion.append(_here)
        try:
            suggestion.append(ea_point(_here))
        except StopIteration:
            pass

    entry_manual_input = tk.Entry(frame_func1)
    button_exec_func1 = tk.Button(frame_func1, text="exec", command=call_then_close(lambda: make_sig(eval(entry_manual_input.get()))))
    for i, s in enumerate(suggestion):
        tk.Button(frame_func1, text=get_name(s) or f'ea_{s:x}', command=call_then_close(make_sig, s)).grid(row=0, column=i)
    entry_manual_input.grid(row=1, column=0, sticky="ew", columnspan=max(len(suggestion), 1))
    button_exec_func1.grid(row=1, column=max(len(suggestion), 1), sticky="ew")

    frame_func1.columnconfigure(0, weight=1)

    # Components for func2
    entry_input_func2 = tk.Entry(frame_func2)
    button_exec_func2 = tk.Button(frame_func2, text="exec", command=call_then_close(lambda: search_sig(entry_input_func2.get())))

    entry_input_func2.grid(row=0, column=0, sticky="ew")
    button_exec_func2.grid(row=0, column=1)

    frame_func2.columnconfigure(0, weight=1)

    # Cancel Button
    button_cancel = tk.Button(frame_buttons, text="cancel", command=root.destroy)
    button_cancel.pack(side="right")

    root.mainloop()


class NySigFinder(plugin_t):
    flags = 0  # 插件类别 或者特性
    wanted_name = "NySigFinder"  # 展示名称
    wanted_hotkey = "Ctrl-Alt-S"  # 其快捷键
    comment = "sig finder"  # 插件描述
    help = "ask gpt if you need"  # 帮助信息

    # 初始化时运行的，可以判断是否要启用这个插件，比如你这个插件主要是针对x86的，
    # 那你就要判断当前分析的是不是x86，然后在决定是否要显示或者启用这个插件
    def init(self):
        return PLUGIN_KEEP

    # 插件退出时要做的操作，比如说你打开了某个文件，要到插件结束时才能关闭，
    # 这里就给你这个机会
    def term(self):
        pass

    # 按快捷键等 插件运行时 的地方
    def run(self, arg):
        main()


# 插件入口
def PLUGIN_ENTRY():
    return NySigFinder()


if __name__ == '__main__':
    NySigFinder().run(0)


repository: Antelcat/ida_copilot
file: ida_copilot.py
description: ChatGPT Agent analyses your IDA pseudocode
wanted_name: Copilot
comment: Copilot
created_at: 2023-11-02 14:23:11+00:00
pushed_at: 2023-11-02 17:14:04+00:00
url:
https://github.com/Antelcat/ida_copilot/blob/0217ce294e6293ec7c3122a7514e3108ad0e82ce/ida_copilot.py

content:
import ida_hexrays
import ida_kernwin
import idaapi

from ida_copilot import panel
from ida_copilot.copilot import Copilot


class CopilotPluginActionHandler(idaapi.action_handler_t):
    def __init__(self):
        super(CopilotPluginActionHandler, self).__init__()

    def activate(self, ctx):
        ida_kernwin.show_wait_box('HIDECANCEL\nRunning Copilot...')
        try:
            Copilot().run()
        finally:
            ida_kernwin.hide_wait_box()
            ida_hexrays.get_widget_vdui(ctx.widget).refresh_view(True)
            ida_kernwin.refresh_idaview_anyway()


    def on_task_complete(self, future):
        # 关闭进度条或状态信息
        ida_kernwin.hide_wait_box()

        # 更新UI...
        ida_kernwin.refresh_idaview_anyway()

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


class CopilotPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = "Copilot"
    help = "Copilot"
    wanted_name = "Copilot"
    wanted_hotkey = ""

    def init(self):
        if not ida_hexrays.init_hexrays_plugin():
            print("Hex-Rays decompiler is not available!")
            return

        run_action = idaapi.action_desc_t(
            'copilot:run',
            'Run Copilot',
            CopilotPluginActionHandler(),
            'Ctrl+Shift+P',
            '使用Copilot分析当前函数',
            -1)
        idaapi.register_action(run_action)
        idaapi.attach_action_to_menu(
            'Edit/Copilot',
            'copilot:run',
            idaapi.SETMENU_APP)

        action_desc = idaapi.action_desc_t(
            'copilot:show_panel',
            'Show Copilot',
            panel.ShowCopilotPanel(panel.CopilotPanel()),
            None,
            'Copilot integration',
            0
        )
        idaapi.register_action(action_desc)

        # 添加菜单项
        idaapi.attach_action_to_menu(
            'Windows/Copilot',
            'copilot:show_panel',
            idaapi.SETMENU_APP)

        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        idaapi.require('ida_copilot')
        print('Copilot reloaded')

    def term(self):
        idaapi.detach_action_from_menu(
            'Edit/Copilot',
            'copilot:run')
        idaapi.unregister_action('copilot:run')

        idaapi.detach_action_from_menu(
            'Windows/Copilot',
            'copilot:show_panel')
        idaapi.unregister_action('copilot:show_panel')


def PLUGIN_ENTRY():
    return CopilotPlugin()


repository: grayhatacademy/ida
file: plugins/localxrefs/localxrefs.py
wanted_name: Localized Xrefs
comment: IDA Localized Xrefs
created_at: 2013-10-16 18:52:21+00:00
pushed_at: 2023-10-23 10:38:19+00:00
url:
https://github.com/grayhatacademy/ida/blob/faf13a13c4f5070fe8f0666c1508f928ffe40e23/plugins/localxrefs/localxrefs.py

content:
# IDA Plugin to search for cross references only within the current defined
# function. Useful, for example, to find instructions that use a particular
# register, or that reference a literal value.
#
# Invoke by highlighting the desired text in IDA, then going to
# Jump->List local xrefs. Highlighting is also supported; once xrefs are found,
# Type the following in the Python command window:
#
#    Python> localxrefs.highlight()       <-- Highlight all xrefs
#    Python> localxrefs.highlight(False)  <-- Un-highlight all xrefs
#
# Craig Heffner
# Tactical Network Solutions

import sys
import idc
import idaapi
import idautils

from shims import ida_shims

localxrefs = None


def add_to_namespace(namespace, name, variable):
    '''
    Add a variable to a different namespace, likely __main__.
    '''
    import importlib

    importer_module = sys.modules[namespace]
    if name in sys.modules.keys():
        if not (sys.version_info.major == 3 and sys.version_info.minor >= 4):
            import imp
            imp.reload(sys.modules[name])
        else:
            importlib.reload(sys.modules[name])
    else:
        m = importlib.import_module(name, None)
        sys.modules[name] = m

    setattr(importer_module, name, variable)


class LocalXrefs(object):
    UP = 'Up  '
    DOWN = 'Down'
    THIS = '-   '

    READ = 'r'
    WRITE = 'w'

    OPND_WRITE_FLAGS = {
            0: idaapi.CF_CHG1,
            1: idaapi.CF_CHG2,
            2: idaapi.CF_CHG3,
            3: idaapi.CF_CHG4,
            4: idaapi.CF_CHG5,
            5: idaapi.CF_CHG6}

    def __init__(self):
        self.xrefs = {}
        self.function = ''
        self._profile_function()

    def _profile_function(self):
        current_ea = ida_shims.get_screen_ea()
        current_function = ida_shims.get_func_name(current_ea)
        current_function_ea = ida_shims.get_name_ea_simple(current_function)

        if current_function:
            self.function = current_function

        ea = ida_shims.get_func_attr(current_function_ea,  idc.FUNCATTR_START)
        end_ea = ida_shims.get_func_attr(current_function_ea, idc.FUNCATTR_END)

        try:
            self.highlighted = ida_shims.get_highlighted_identifier()
        except TypeError:
            raise Exception("No text was highlighted. Nothing to show.") 

        while ea < end_ea and ea != idc.BADADDR and self.highlighted:
            i = 0
            match = False
            optype = self.READ

            insn = ida_shims.decode_insn(ea)

            mnem = ida_shims.print_insn_mnem(ea)

            if self.highlighted in mnem:
                match = True
            elif idaapi.is_call_insn(ea):
                for xref in idautils.XrefsFrom(ea):
                    if xref.type != 21:
                        name = ida_shims.get_name(xref.to)
                        if name and self.highlighted in name:
                            match = True
                            break
            else:
                while True:
                    opnd = ida_shims.print_operand(ea, i)
                    if opnd:
                        if self.highlighted in opnd:
                            canon_feature = ida_shims.get_canon_feature(insn)
                            match = True
                            if canon_feature & self.OPND_WRITE_FLAGS[i]:
                                optype = self.WRITE
                        i += 1
                    else:
                        break

            if not match:
                comment = idc.GetCommentEx(ea, 0)
                if comment and self.highlighted in comment:
                    match = True
                else:
                    comment = idc.GetCommentEx(ea, 1)
                    if comment and self.highlighted in comment:
                        match = True

            if match:
                if ea > current_ea:
                    direction = self.DOWN
                elif ea < current_ea:
                    direction = self.UP
                else:
                    direction = self.THIS

                self.xrefs[ea] = {
                    'offset': ida_shims.get_func_off_str(ea),
                    'mnem': mnem,
                    'type': optype,
                    'direction': direction,
                    'text': idc.GetDisasm(ea),
                }

            ea = ida_shims.next_head(ea)

    def highlight(self, highlight=True, mnem=None, optype=None, direction=None,
                  text=None):
        for (ea, info) in self.xrefs.items():
            if mnem and info['mnem'] != mnem:
                highlight = False
            elif optype and info['optype'] != optype:
                highlight = False
            elif direction and info['direction'] != direction:
                highlight = False
            elif text and info['text'] != text:
                highlight = False

            if highlight:
                color = 0x00ff00
            else:
                color = idc.DEFCOLOR

            ida_shims.set_color(ea, idc.CIC_ITEM, color)

    def unhighlight(self):
        self.highlight(False)


def show_local_xrefs(arg=None):
    delim = '-' * 86 + '\n'
    header = '\nXrefs to %s from %s:\n'

    global localxrefs
    fmt = ''

    r = LocalXrefs()
    localxrefs = r

#    offsets = r.xrefs.keys()
#    offsets.sort()

    offsets = sorted(r.xrefs)

    if r.highlighted:
        ida_shims.msg(header % (r.highlighted, r.function))
        ida_shims.msg(delim)

        for ea in offsets:
            info = r.xrefs[ea]

            if not fmt:
                fmt = "%%s   %%s   %%-%ds   %%s\n" % (len(info['offset']) + 15)

            ida_shims.msg(fmt % (info['direction'], info['type'],
                                 info['offset'], info['text']))

        ida_shims.msg(delim)


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

        def activate(self, ctx):
            global localxrefs
            show_local_xrefs()
            add_to_namespace('__main__', 'localxrefs', localxrefs)
            return 1

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


class localizedxrefs_t(idaapi.plugin_t):
    flags = 0
    comment = "IDA Localized Xrefs"
    help = ""
    wanted_name = "Localized Xrefs"
    wanted_hotkey = ""
    menu_context = None
    menu_name = 'List local xrefs'
    action_name = 'localxrefs:action'
    menu_tab = 'Jump/'

    def init(self):
        if idaapi.IDA_SDK_VERSION >= 700:
            action_desc = idaapi.action_desc_t(self.action_name,
                                               self.menu_name,
                                               LocalXrefHandler(),
                                               self.wanted_hotkey,
                                               'Localized Xrefs.',
                                               199)
            idaapi.register_action(action_desc)
            idaapi.attach_action_to_menu(
                self.menu_tab, self.action_name, idaapi.SETMENU_APP)
        else:
            self.menu_context = idaapi.add_menu_item(
                self.menu_tab, self.menu_name, "", 0, show_local_xrefs, (None,))
        return idaapi.PLUGIN_KEEP

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

    def run(self):
        pass


def PLUGIN_ENTRY():
    return localizedxrefs_t()


