# plugins

repository: stevemk14ebr/RETools
file: IdaScripts/plugins/string_from_selection.py
description: My reversing tools. Some custom, some not.
wanted_name: String From Selection
comment: Define a string from selection, useful for non-null terminated strings
created_at: 2020-05-27 21:10:30+00:00
pushed_at: 2024-01-03 21:54:17+00:00
url:
https://github.com/stevemk14ebr/RETools/blob/9501332bf02688d3a4b31d85a25349671767be98/IdaScripts/plugins/string_from_selection.py

content:
import idaapi
import ida_bytes
import ida_nalt

popup_action_names = []
class Hooks(idaapi.UI_Hooks):
    def __init__(self):
        idaapi.UI_Hooks.__init__(self)
 
    def finish_populating_widget_popup(self, form, popup):
        global popup_action_names
        form_type = idaapi.get_widget_type(form)
        if form_type == idaapi.BWN_DISASM:
            for action_name in popup_action_names:
                idaapi.attach_action_to_popup(form, popup, action_name, None)
 
class StringFromSelectionPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_HIDE | idaapi.PLUGIN_PROC
    comment = "Define a string from selection, useful for non-null terminated strings"
    help = "Select a region and right click to define as a string of that length"
    wanted_name = "String From Selection"
    wanted_hotkey = ""
    
    def init(self):
        idaapi.msg("String From Selection :: init\n")
        AddToPopup('definestr:add_action', 'Define string from selection', DefineStringAction(), '', None)
 
        self.hooks = Hooks()
        self.hooks.hook()
        return idaapi.PLUGIN_KEEP
 
    def run(self):
        idaapi.msg("String From Selection :: run\n")
 
    def term(self):
        idaapi.msg("String From Selection :: term\n")
        if self.hooks:
            self.hooks.unhook()
 
        idaapi.unregister_action('definestr:add_action')
 
class EA_Action(idaapi.action_handler_t):
    def __init__(self):
        idaapi.action_handler_t.__init__(self)
 
    def update(self, ctx):
        return idaapi.AST_ENABLE_FOR_WIDGET if ctx.form_type == idaapi.BWN_DISASM else idaapi.AST_DISABLE_FOR_WIDGET
 
class DefineStringAction(EA_Action):
    def activate(self, ctx):
        t0, t1, view = idaapi.twinpos_t(), idaapi.twinpos_t(), idaapi.get_current_viewer()
        if idaapi.read_selection(view, t0, t1):
            start, end = t0.place(view).toea(), t1.place(view).toea()
            end += idaapi.get_item_size(end)
        else:
            start = idaapi.get_screen_ea()
 
            if start == idaapi.BADADDR:
                print('String From Selection  :: Screen EA == idaapi.BADADDR')
                return 0
 
            end = start + idaapi.get_item_size(start)
 
        if start == idaapi.BADADDR:
            print('String From Selection :: Selection EA == idaapi.BADADDR')
            return 0
 
        if start == end:
            print('String From Selection :: Selection is of zero length, nothing to define')
            return 0
 
        ida_bytes.create_strlit(start, end-start, ida_nalt.STRTYPE_TERMCHR)
        return 1
 
def AddToPopup(action_name, display, handler, shortcut, tooltip, icon=None):
    global popup_action_names
 
    if tooltip == None:
        tooltip = action_name
 
    if idaapi.register_action(idaapi.action_desc_t(
                action_name,
                display,
                handler,
                shortcut,
                tooltip
            )):
 
        popup_action_names.append(action_name)
    else:
        print('String From Selection :: Error registering action %s' % (action_name))
 
def PLUGIN_ENTRY(*args, **kwargs):
    return StringFromSelectionPlugin()


repository: airbus-seclab/bincat
file: python/idabincat/bcplugin.py
wanted_name: BinCAT
comment: Interface to the BinCAT analyzer
created_at: 2017-05-29 15:31:22+00:00
pushed_at: 2025-02-25 16:54:39+00:00
description:
Binary code static analyser, with IDA integration. Performs value and taint analysis, type reconstruction, use-after-free and double-free detection

url:
https://github.com/airbus-seclab/bincat/blob/5d0ee3b56867059427eb0f4123c4d9de0b8059dd/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 ida_ida
import idaapi
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
        line = loglines.pop()
        if line == staging:
            n += 1
        else:
            if flush_staging():
                break
            staging = line
            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):
        procname = ida_ida.inf_get_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):
    """
    Runs BinCAT locally using QProcess.
    """

    def __init__(self, *args, **kwargs):
        Analyzer.__init__(self, *args, **kwargs)

    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")

        self.process = QtCore.QProcess()
        # Qprocess signal handlers
        self.process.error.connect(self.procanalyzer_on_error)
        self.process.stateChanged.connect(self.procanalyzer_on_state_change)
        self.process.started.connect(self.procanalyzer_on_start)
        self.process.finished.connect(self.procanalyzer_on_finish)

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

        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.process.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.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.process.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.process.readAllStandardOutput()))
        bc_log.info("---- stderr ----------------")
        bc_log.info(str(self.process.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: rcx/ida-scripts
file: plugins/force_lvar_width.py
description: Collection of IDA Pro/Hex-Rays configs, scripts, and plugins
wanted_name: Hex-Rays lvar width forcer
comment: Force lvar width plugin for Hex-Rays decompiler
created_at: 2018-08-01 03:12:29+00:00
pushed_at: 2020-03-20 20:45:25+00:00
url:
https://github.com/rcx/ida-scripts/blob/963d7a40330eed1d5e8ead8f4297e5a3555e2456/plugins/force_lvar_width.py

content:
""" Forcibly change the width of an lvar.
Useful for making an lvar smaller (which Hex-Rays does not let you do apparently)
Tags (for search engines):
IDA Pro Hex-Rays make stack variable smaller force variable size decrease width delete variable split up break up variable into smaller
"""

import idautils
import idaapi
import idc
import ida_hexrays

import traceback

force_width_actname = "forcelvarwidth:forcewidth"

class force_width_action_handler_t(idaapi.action_handler_t):
    def __init__(self, callback_info):
        idaapi.action_handler_t.__init__(self)
        self.callback_info = callback_info

    def activate(self, ctx):
        vdui = idaapi.get_widget_vdui(ctx.widget)
        self.callback_info.gui_action_callback(vdui)
        return 1

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


class hexrays_callback_info(object):

    def __init__(self):
        self.vu = None
        return

    def load(self):
        return

    def save(self):
        return

    def do_force_width(self, cfunc, insn):

        if insn.opname != 'if':
            return False

        cif = insn.details

        if not cif.ithen or not cif.ielse:
            return False

        idaapi.qswap(cif.ithen, cif.ielse)
        cond = idaapi.cexpr_t(cif.expr)
        notcond = idaapi.lnot(cond)

        cif.expr.swap(notcond)

        return True

    def gui_action_callback(self, vu):

        cfunc = vu.cfunc.__deref__()
        
        if not vu.get_current_item(idaapi.USE_KEYBOARD):
            print "Force lvar width: you don't have anything selected"
            return False
            
        badlv = vu.item.get_lvar()
        if not badlv:
            print "Force lvar width: you don't have an lvar selected"
            return False

        new_width = idc.AskLong(badlv.width, "Enter the new width for " + badlv.name)
        if new_width == None: # cancelled
            print "Force lvar width: operation cancelled"
            return False

        if new_width <= 0:
            print "Force lvar width: not allowed. Non-positive width will crash IDA"
            return False

        badlv.set_width(new_width)
        print 'Set the type in IDA (Y) for it to apply'
        idaapi.process_ui_action('hx:SetType')

        # vu.refresh_ctext()
        print 'Force lvar width: OK.'
        return True

    def event_callback(self, event, *args):

        if event == idaapi.hxe_populating_popup:
            widget, phandle, vu = args
            res = idaapi.attach_action_to_popup(vu.ct, None, force_width_actname)

        return 0

class ForceLvarWidth(idaapi.plugin_t):
    flags = idaapi.PLUGIN_PROC | idaapi.PLUGIN_HIDE
    wanted_hotkey = "Shift-W"
    comment = "Force lvar width plugin for Hex-Rays decompiler"
    help = "There is no one to help you now"
    wanted_name = "Hex-Rays lvar width forcer"

    def init(self):
        if idaapi.init_hexrays_plugin():
            i = hexrays_callback_info()
            idaapi.register_action(
                idaapi.action_desc_t(
                    force_width_actname,
                    "Force lvar width",
                    force_width_action_handler_t(i),
                    "Shift-W"))
            idaapi.install_hexrays_callback(i.event_callback)
            print 'Hex-Rays lvar width forcer by ecx86 loaded!'
        else:
            print 'Force lvar width: Hexrays is not available.'

    def term(self):
        pass

    def run(self, arg):
        pass

def PLUGIN_ENTRY():
    return ForceLvarWidth()


repository: trichimtrich/ofb
file: ofb.py
description: offset from base, idapython
url: https://github.com/trichimtrich/ofb/blob/035ec428071425493051a64cda7478e78eaf7bba/ofb.py
wanted_name: OFB
comment: OFB
created_at: 2021-02-07 11:46:35+00:00
pushed_at: 2021-06-15 21:35:35+00:00
content:
import idaapi
from PyQt5.Qt import QApplication
import re

NAME_TEMPLATE = '{name}_{offset}'


def _log(fmt, *args):
    print('[OFB] ' + fmt % args)


def _get_base(ea):
    if idaapi.is_debugger_on():
        x = idaapi.modinfo_t()
        _ = idaapi.get_module_info(ea, x)
        base = x.base
    else:
        base = idaapi.get_imagebase()
    return base


class ActionWrap(idaapi.action_handler_t):
    desc = None
    hotkey = None

    def do_register(self):
        self.do_unregister()
        res = idaapi.register_action(idaapi.action_desc_t(self.name, self.desc, self, self.hotkey))
        _log('Register action %s: %s', repr(self.name), res)


    def do_unregister(self):
        stt, _ = idaapi.get_action_state(self.name)
        if stt:
            res = idaapi.unregister_action(self.name)
            _log('Unregister action %s: %s', repr(self.name), res)

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


class ActionGetOffset(ActionWrap):
    name = 'get_offset_from_base'
    desc = 'Get offset from base'

    def activate(self, ctx):
        ea = ctx.cur_ea
        # print(hex(ida_kernwin.get_screen_ea()))
        base = _get_base(ea)
        off = ea - base
        _log('EA: %x. Base: %x => Offset: %x', ea, base, off)
        QApplication.clipboard().setText(hex(off))
        _log('Copied offset to clipboard')
        
        return 1
        
        
class ActionJumpOffset(ActionWrap):
    name = 'jump_offset_from_base'
    desc = 'Jump offset from base'
    hotkey = 'shift+G'

    def activate(self, ctx):
        ea = ctx.cur_ea
        base = _get_base(ea)
        off = idaapi.ask_addr(0, 'New offset')
        if off is None:
            return 1

        new_ea = base + off
        _log('Base: %x. New offset: %x => New addr: %x', base, off, new_ea)
        if not idaapi.jumpto(new_ea):
            idaapi.warning('Invalid offset to jump from base')

        return 1
        

class ActionRenameOffset(ActionWrap):
    name = 'rename_offset_from_base'
    desc = 'Rename offset from base'
    hotkey = 'shift+N'

    def __init__(self, template):
        super(ActionRenameOffset, self).__init__()

        # precheck template
        try:
            template.format(name='a', offset='1234')
            if template.count('{name}') != 1:
                raise Exception()
        except:
            idaapi.warning('Invalid name template {}, set to default'.format(repr(template)))
            template = '{name}_{offset}'

        self.template = template

        pattern = '^' + template + '$'
        pattern = pattern.replace('{name}', '(.+?)')
        pattern = pattern.replace('{offset}', '[a-z0-9]+')
        self.pattern = re.compile(pattern)


    def activate(self, ctx):
        ea = ctx.cur_ea
        cur_name = idaapi.get_name(ea)
        real_cur_name = None
        if cur_name:
            m = self.pattern.match(cur_name)
            if m:
                real_cur_name = m.group(1)
            else:
                real_cur_name = cur_name
        
        if real_cur_name is None:
            real_cur_name = ''

        real_new_name = idaapi.ask_str(real_cur_name, 0, self.template)
        if real_new_name and real_new_name != real_cur_name:
            base = _get_base(ea)
            off = ea - base
            new_name = self.template.format(
                name=real_new_name, 
                offset='%x' % off,
            )
            if not idaapi.set_name(ea, new_name):
                _log('Set name failed at %x : %s', ea, new_name)
    
        return 1
    

class OFBHook(idaapi.UI_Hooks):
    def __init__(self, *args):
        super(OFBHook, self).__init__(*args)
        self.actions = []


    def register_action(self, a):
        self.actions.append(a)
        a.do_register()


    def finish_populating_widget_popup(self, form, popup):
        if idaapi.get_widget_type(form) in (idaapi.BWN_DISASM, idaapi.BWN_DUMP, idaapi.BWN_PSEUDOCODE ):
            for a in self.actions:
                idaapi.attach_action_to_popup(form, popup, a.name, 'Hoho')

    def fini(self):
        for a in self.actions:
            a.do_unregister()


class OFB(idaapi.plugin_t):
    flags = idaapi.PLUGIN_HIDE
    comment = 'OFB'
    help = ''
    wanted_name = 'OFB'
    wanted_hotkey = ''

    
    def init(self):
        ui_hook = OFBHook()
        ui_hook.register_action(ActionGetOffset())
        ui_hook.register_action(ActionJumpOffset())
        ui_hook.register_action(ActionRenameOffset(NAME_TEMPLATE))

        _log('UI Hook: %s', ui_hook.hook())
        self.ui_hook = ui_hook
        return idaapi.PLUGIN_KEEP

        
    def run(self, *argv, **kargv):
        pass

        
    def term(self):
        if hasattr(self, 'ui_hook'):
            _log('Unhook UI: %s', self.ui_hook.unhook())
            self.ui_hook.fini()
        

def PLUGIN_ENTRY():
    return OFB()


repository: toshipiazza/ghidra-plugin-hexagon
file: GhidraBuild/IDAPro/Python/7xx/plugins/xml_importer.py
description: Hexagon decompiler for Ghidra
wanted_name: XML Importer
comment: Import XML PROGRAM file
created_at: 2022-01-15 03:46:25+00:00
pushed_at: 2022-02-15 08:27:42+00:00
url:
https://github.com/toshipiazza/ghidra-plugin-hexagon/blob/be29246fc234ece766a5719564bb55ccf69bc1e7/GhidraBuild/IDAPro/Python/7xx/plugins/xml_importer.py

content:
## ###
#  IP: GHIDRA
# 
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#  
#       http://www.apache.org/licenses/LICENSE-2.0
#  
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
##
#---------------------------------------------------------------------
# xmlimp.py - IDA XML Importer plugin
#---------------------------------------------------------------------
"""
Plugin for IDA to import a XML PROGRAM file into an existing open database.
This file must be placed in the IDA plugins directory.
The file idaxml.py must be placed in the IDA python directory.
"""

import ida_idaapi
import ida_pro
import idaxml
import idc
import sys

class XmlImporterPlugin(ida_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 idaxml.is_ida_version_supported():
            return ida_idaapi.PLUGIN_OK 
        else:
            return ida_idaapi.PLUGIN_SKIP


    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 = idc.set_ida_state(idc.IDA_STATUS_WORK)
        xml = idaxml.XmlImporter(idaxml.PLUGIN, arg)
        try:
            try:
                xml.import_xml()
            except idaxml.Cancelled:
                msg = "XML Import cancelled!"
                print "\n" + msg
                idc.warning(msg)
            except idaxml.MultipleAddressSpacesNotSupported:
                msg  = "XML Import cancelled!"
                msg += "\n\nXML Import does not currently support"
                msg += "\nimporting multiple address spaces."
                print "\n" + msg
                idc.warning(msg)
            except:
                msg = "***** Exception occurred: XML Importer failed! *****"
                print "\n" + msg + "\n", sys.exc_type, sys.exc_value
                idc.warning(msg)
        finally:
            xml.cleanup()
            idc.set_ida_state(st)


    def term(self):
        pass


def PLUGIN_ENTRY():
    return XmlImporterPlugin()


repository: Air14/SourceSync
file: SourceSync/Ida/SourceSync.py
description: Set of plugins and library for dynamic pdb generation and synchronisation
wanted_name: SourceSync plugin
comment: This is comment
created_at: 2024-05-03 19:27:12+00:00
pushed_at: 2024-05-03 19:35:46+00:00
url:
https://github.com/Air14/SourceSync/blob/a9d3473e24a23eadaaaa8706f99877285909dda2/SourceSync/Ida/SourceSync.py

content:
import os
import sys
sys.path.append(os.path.join(os.path.dirname(__file__), "srcsync"))

import grpc
import threading
import ida_hexrays
import ida_idp
import ida_bytes
import ida_funcs
import idaapi
import ctypes
import configparser
import PdbGeneratorPy
from srcsync.TypeExtractor import TypeExtractor
from srcsync.PEDataExtractor import PEDataExtractor
from srcsync.SymbolExtractor import SymbolExtractor
from srcsync.FunctionDataExtractor import FunctionDataExtractor
from srcsync.SourceSyncServer import SourceSyncServer
from srcsync.DecompilerSynchronizer_pb2_grpc import add_DecompilerSynchronizerServicer_to_server
from concurrent import futures
from PyQt5 import QtCore, QtWidgets

HookCb = ctypes.WINFUNCTYPE(
    # return type
    ctypes.c_int,      # int idaapi

    # argument types
    ctypes.c_void_p,   # void       *user_data
    ctypes.c_int,      # int         notification_code
    ctypes.c_void_p,   # va_list     va 
)

class HOOK_TYPES:
    # Hook to the processor module.
    HT_IDP = 0
    # Hook to the user interface.
    HT_UI = 1
    # Hook to the debugger.
    HT_DBG = 2
    # Hook to the database events.
    HT_IDB = 3
    # Internal debugger events.
    HT_DEV = 4
    # Custom/IDA views notifications.
    HT_VIEW = 5
    # Output window notifications.
    HT_OUTPUT = 6
    HT_LAST = 7

def HookToNotificationPoint(Dll, HookType, CallBack, UserData):
    hook = Dll.hook_to_notification_point
    hook.argtypes = [
        ctypes.c_int,      # hook_type_t hook_type
        HookCb,            # hook_cb_t  *cb
        ctypes.c_void_p,   # void       *user_data
    ]

    hook(HookType, CallBack, UserData)

def RemoveHookToNotificationPoint(Dll, HookType, CallBack, UserData):
    removeHook = Dll.unhook_from_notification_point
    removeHook.argtypes = [
        ctypes.c_int,      # hook_type_t hook_type
        HookCb,            # hook_cb_t  *cb
        ctypes.c_void_p,   # void       *user_data
    ]

    removeHook(HookType, CallBack, UserData)

def GetHostIpAndPort():
    for location in (os.path.dirname(os.path.realpath(idaapi.get_path(idaapi.PATH_TYPE_IDB))), os.environ["USERPROFILE"]):
        configPath = os.path.join(location, '.srcsync')
        if os.path.exists(configPath):
            config = configparser.ConfigParser()
            config.read(configPath)
            if config.has_section('INTERFACE'):
                return config.get('INTERFACE', 'host') + ":" + config.get('INTERFACE', 'port')

    return "localhost:5111"

class CheckBoxActionHandler(idaapi.action_handler_t):
    def __init__(self, Cb):
        idaapi.action_handler_t.__init__(self)
        self.Cb = Cb

    def activate(self, Ctx):
        self.Cb.toggle()
        return 1

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

class SourceSyncForm(idaapi.PluginForm):
    def __init__(self):
        super().__init__()
        self.InitIdbEventHook()
        
    def InitIdbEventHook(self):
        def IdbEventCallback(UserData, NotificationCode, VaList):
            VaList = ctypes.cast(VaList, ctypes.POINTER(ctypes.c_ulonglong))
        
            if self.SourceSyncServer is None:
                return 0

            if NotificationCode == ida_idp.renamed:
                address = VaList[0]
                if not ida_bytes.is_loaded(address):
                    return 0
        
                if function := ida_funcs.get_func(address):
                    self.SourceSyncServer.OnFunctionChanged(function.start_ea)
                else:
                    self.SourceSyncServer.OnPublicSymbolChanged()
        
            elif NotificationCode == ida_idp.local_types_changed:
                self.SourceSyncServer.OnTypesChanged()
        
            elif NotificationCode == ida_idp.ti_changed:
                address = VaList[0]
                if ida_funcs.get_func(address) or not ida_bytes.is_loaded(address):
                    return 0
        
                self.SourceSyncServer.OnPublicSymbolChanged()
        
            return 0
        
        self.Dll = ctypes.windll["ida64.dll"]

        # Need to keep a ref around, or the function gets garbage collected
        self.IdbEventCallback = HookCb(IdbEventCallback)

        # Need to keep a ref around, or the param gets garbage collected
        self.UserData = ctypes.c_long(10)

    def StartServer(self):
        typeExtractor = TypeExtractor()
        self.SourceSyncServer = SourceSyncServer(typeExtractor, PEDataExtractor(), SymbolExtractor(typeExtractor), FunctionDataExtractor(typeExtractor))
        serverAddress = GetHostIpAndPort()
        
        def SourceSyncServerThread():
            self.Server = grpc.server(futures.ThreadPoolExecutor())
            add_DecompilerSynchronizerServicer_to_server(self.SourceSyncServer, self.Server)
            self.Server.add_insecure_port(serverAddress)

            self.Server.start()
            self.Server.wait_for_termination()

        sourceSyncServerThread = threading.Thread(target=SourceSyncServerThread)
        sourceSyncServerThread.start()

    def StopServer(self):
        if self.Server is not None:
            self.Server.stop(None)

    def InstallHexraysHooks(self):
        class LvarHooks(ida_hexrays.Hexrays_Hooks):
            def lvar_name_changed(_, *Args):
                vdui = Args[0]
                self.SourceSyncServer.OnFunctionChanged(vdui.cfunc.entry_ea)
                return 0
            
            def lvar_type_changed(_, *Args):
                vdui = Args[0]
                self.SourceSyncServer.OnFunctionChanged(vdui.cfunc.entry_ea)
                return 0
            
            def lvar_type_changed(_, *Args):
                vdui = Args[0]
                self.SourceSyncServer.OnFunctionChanged(vdui.cfunc.entry_ea)
                return 0

        self.LvarHooks = LvarHooks()
        self.LvarHooks.hook()

    def RemoveHexraysHooks(self):
        self.LvarHooks.unhook()

    def CbChangeSourceSyncState(self, State):
        if State == QtCore.Qt.Checked:
            self.StartServer()
            self.InstallHexraysHooks()
            HookToNotificationPoint(self.Dll, HOOK_TYPES.HT_IDB, self.IdbEventCallback, ctypes.byref(self.UserData))
        else:
            RemoveHookToNotificationPoint(self.Dll, HOOK_TYPES.HT_IDB, self.IdbEventCallback, ctypes.byref(self.UserData))
            self.RemoveHexraysHooks()
            self.StopServer()
            
    def CbRestartSourceSyncServer(self):
        self.StopServer()
        self.StartServer()

    def CbGeneratePdb(self):
        print("[SourceSync] Generating pdb")
        typeExtractor = TypeExtractor()
        typeExtractor.GatherData(ExecuteSync = False)
        
        symbolExtractor = SymbolExtractor(typeExtractor)
        publicSymbolsData = symbolExtractor.GetPublics(ExecuteSync = False)
        globalSymbolsData = symbolExtractor.GetGlobals(ExecuteSync = False)
        
        functionDataExtractor = FunctionDataExtractor(typeExtractor)
        functionsData = functionDataExtractor.GetFunctionsData()
        
        enumsData = typeExtractor.GetEnumsData()
        structsData = typeExtractor.GetStructsData()
        complexTypes = typeExtractor.GetComplexTypesData()
        
        peDataExtractor = PEDataExtractor()
        pdbInfo = peDataExtractor.GetPdbInfo()
        sectionsData = peDataExtractor.GetSectionsData()
        if len(sectionsData) == 0:
            print("[SourceSync] Failed to get sections")       
            return

        if ida_idp.ph.id == ida_idp.PLFM_386 and ida_idp.ph.flag & ida_idp.PR_USE64:
            cpuArchitectureType = PdbGeneratorPy.CpuArchitectureType.X86_64
        else:
            cpuArchitectureType = PdbGeneratorPy.CpuArchitectureType.X86
            
        pdbGenerator = PdbGeneratorPy.PdbGenerator(
            complexTypes, structsData, enumsData, functionsData,
            pdbInfo, sectionsData, publicSymbolsData, globalSymbolsData, cpuArchitectureType)
        
        if pdbGenerator.Generate():
            print("[SourceSync] Pdb generated")       
        else:
            print("[SourceSync] Failed to generate pdb")       

    def OnCreate(self, Form):
        parent = self.FormToPyQtWidget(Form)
        
        self.SyncBox = QtWidgets.QCheckBox('Synchronization enable')
        self.SyncBox.move(20, 20)
        self.SyncBox.stateChanged.connect(self.CbChangeSourceSyncState)
        
        self.RestartButton = QtWidgets.QPushButton('Restart', parent)
        self.RestartButton.setToolTip('Restart SourceSync server')
        self.RestartButton.clicked.connect(self.CbRestartSourceSyncServer)
        
        self.GeneratePdbButton = QtWidgets.QPushButton('Generate Pdb', parent)
        self.GeneratePdbButton.setToolTip('Generate pdb for current idb inside idb directory')
        self.GeneratePdbButton.clicked.connect(self.CbGeneratePdb)

        layout = QtWidgets.QGridLayout()
        layout.addWidget(self.SyncBox)
        layout.addWidget(self.RestartButton)
        layout.addWidget(self.GeneratePdbButton)
        layout.setColumnStretch(4, 1)
        layout.setRowStretch(4, 1)
        parent.setLayout(layout)
        
        # Enable source sync server by default
        self.SyncBox.toggle()
        
        actionDesc = idaapi.action_desc_t(
            'SyncToogle:action',
            'Toggle syncing',
            CheckBoxActionHandler(self.SyncBox),
            'Ctrl+Shift+D',
            'Toggle syncing',
            203)

        idaapi.register_action(actionDesc)
        idaapi.attach_action_to_toolbar("AnalysisToolBar", 'SyncToogle:action')
        
    def OnClose(self, Form):
        self.SyncBox.toggle()
        idaapi.unregister_action('SyncToogle:action')
        idaapi.detach_action_from_toolbar("DebugToolBar", 'SyncToogle:action')
        
        global sourceSyncForm
        sourceSyncForm = None
        
    def Show(Self):
        return idaapi.PluginForm.Show(Self, "SourceSync", options=idaapi.PluginForm.WOPN_PERSIST)

class SourceSyncPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_PROC
    comment = "This is comment"
    help = "This is help"
    wanted_name = "SourceSync plugin"
    wanted_hotkey = "Alt-Shift-D"
    global sourceSyncForm
    sourceSyncForm = None
    
    def init(self):
        return idaapi.PLUGIN_KEEP

    def run(self, Arg):
        if ida_idp.ph.id != ida_idp.PLFM_386:
            print("[SourceSync] Only x86_64/x86 architecture is supported.")
            return

        if "portable executable" not in idaapi.get_file_type_name().lower():
            print("[SourceSync] Only PE files are supported.")
            return
        
        global sourceSyncForm
        if not sourceSyncForm:
            sourceSyncForm = SourceSyncForm()
            sourceSyncForm.Show()
    
    def term(self):
        pass

def PLUGIN_ENTRY():
    return SourceSyncPlugin()


repository: captainGeech42/daas
file: decompile.py
description: Decompilation as a Service (powered by Hex-Rays)
url: https://github.com/captainGeech42/daas/blob/78a382b314a99ea650061c76fab3d28e7cdfc735/decompile.py
wanted_name: Ida Batch Decompile
comment: Batch Decompile
created_at: 2021-02-03 02:52:09+00:00
pushed_at: 2021-02-05 19:00:47+00:00
content:
#! /usr/bin/env python
# -*- coding: UTF-8 -*-
# Author : <github.com/tintinweb>

"""
This script was written by tintinweb, and is from https://github.com/tintinweb/ida-batch_decompile
"""

"""
IdaBatchDecompile Plugin and Script adds annotation and batch decompilation functionality to IDA Pro

* requires hexrays decompiler plugin

Usage:

* as idascript in ida gui mode: IDA Pro -> File/Script file... -> IdaDecompileBatch ...
* as idascript in ida cmdline mode: ida(w|w64) -B -M -S"<path_to_this_script> \"--option1\" \"--option2\"", "<target>"
 * see --help for options
* as Plugin: follow ida documentation on how to add python plugins

"""
import sys
import json
import glob
import subprocess
import shutil
import os
import tempfile
from optparse import OptionParser

import idaapi
import idautils
from idc import *

if idaapi.IDA_SDK_VERSION >= 700:
    import ida_idaapi
    import ida_kernwin
    from idaapi import *
    from idc import *

import logging

logger = logging.getLogger(__name__)


class IdaLocation(object):
    """ Wrap idautils Function
    """

    def __init__(self, location):
        self.at = location
        # self.name = GetFunctionName(location)
        self.name = GetFuncOffset(location)
        self.start = 0
        self.end = 0
        self.func_offset = 0
        try:
            _func = idaapi.get_func(location)
            self.start = _func.startEA
            self.end = _func.endEA  # ==FindFuncEnd(location)
            self.func_offset = self.start - self.at
        except Exception, e:
            logger.exception(e)
        if not self.name:
            self.indirect = True
        else:
            self.indirect = False

    def __repr__(self, *args, **kwargs):
        return "<Function %r at 0x%x (0x%x::0x%x)>" % (self.name, self.at,
                                                       self.start, self.end)

    def get_xrefs(self):
        return (IdaLocation(x.frm) for x in idautils.XrefsTo(self.at))

    def get_coderefs(self):
        return (IdaLocation(frm) for frm in idautils.CodeRefsTo(self.at, 0))

    def as_dict(self):
        return {'at': self.at, 'name': self.name}

    def decompile(self):
        """ decompile function
        """
        try:
            return idaapi.decompile(self.at)
        except idaapi.DecompilationFailure, e:
            return repr(str(e))
        text = str(idaapi.decompile(self.at)).strip()
        '''
        sprintf:
        Python>for w in idaapi.decompile(0x00001578 ).lvars: print w.name
            s
            format
            result
        '''
        # decompile.arguments
        # for w in idaapi.decompile(0x00001EF0 ).lvars: print w.name
        if not grep:
            return text.split('\n')
        # return all lines 
        return [line.strip() for line in text.split('\n') if grep in line]

    def get_function_args(self):
        # find the stack frame
        stack = GetFrame(self.start)
        stack_size = GetStrucSize(stack)
        # figure out all of the variable names
        # base is either ' s' ... saved register or ' r' ... return address
        base = GetMemberOffset(stack, ' s')
        if base == -1:
            base = GetMemberOffset(stack, ' r')
        if base == -1:
            # no ' s' no ' r' assume zero
            base == 0
        stack_vars = []

        for memberoffset in xrange(stack_size):
            previous = stack_vars[-1] if len(stack_vars) else None
            var_name = GetMemberName(stack, memberoffset)
            if not var_name or (previous and var_name == previous.get("name")):
                # skip that entry, already processed
                continue

            offset = GetMemberOffset(stack, var_name) - base
            size = GetMemberSize(stack, memberoffset)
            if previous:
                diff = offset - previous['offset']
                previous['diff_size'] = diff
            stack_vars.append({'name': var_name,
                               'offset': offset,
                               'offset_text': '[bp%Xh]' % offset if offset < 0 else '[bp+%Xh]' % offset,
                               'size': size,
                               'diff_size': size})
        return stack_size, stack_vars


class IdaHelper(object):
    """ Namespace for ida helper functions
    """

    @staticmethod
    def get_functions():
        return (IdaLocation(f) for f in idautils.Functions())

    @staticmethod
    def get_imports():
        for i in xrange(0, idaapi.get_import_module_qty()):
            name = idaapi.get_import_module_name(i)
            if name:
                yield name

    @staticmethod
    def decompile_full(outfile):
        return idaapi.decompile_many(outfile, None, 0)

    @staticmethod
    def annotate_xrefs():
        stats = {'annotated_functions': 0, 'errors': 0}
        for f in IdaHelper.get_functions():
            try:
                function_comment = GetFunctionCmt(f.start, 0)
                if '**** XREFS ****' in function_comment:
                    logger.debug("[i] skipping function %r, already annotated." % f.name)
                    continue
                xrefs = [x.name for x in f.get_coderefs()]
                comment = []
                if function_comment:
                    comment.append(function_comment)
                comment.append("***** XREFS *****")
                comment.append("* # %d" % len(xrefs))
                comment.append(', '.join(xrefs))
                comment.append("*******************")
                SetFunctionCmt(f.start, '\n'.join(comment), 0)
                stats['annotated_functions'] += 1
            except Exception as e:
                print ("Annotate XRefs: %r"%e)
                stats['errors'] += 1
        print "[+] stats: %r" % stats
        print "[+] Done!"

    @staticmethod
    def annotate_functions_with_local_var_size():
        stats = {'annotated_functions': 0, 'errors': 0}
        for f in IdaHelper.get_functions():
            try:
                function_comment = GetFunctionCmt(f.start, 0)
                if '**** Variables ****' in function_comment:
                    logger.debug("[i] skipping function %r, already annotated." % f.name)
                    continue
                size, stack_vars = f.get_function_args()
                comment = []
                if function_comment:
                    comment.append(function_comment)
                comment.append("**** Variables ****")
                comment.append("* stack size: %s" % size)
                for s in stack_vars:
                    comment.append(json.dumps(s))
                comment.append("*******************")
                SetFunctionCmt(f.start, '\n'.join(comment), 0)
                stats['annotated_functions'] += 1
            except Exception, e:
                print ("Annotate Funcs: %r" % e)
                stats['errors'] += 1
        print "[+] stats: %r" % stats
        print "[+] Done!"


class IdaDecompileBatchController(object):
    def __init__(self):
        self.is_windows = sys.platform.startswith('win')
        self.is_ida64 = GetIdbPath().endswith(".i64")  # hackhackhack - check if we're ida64 or ida32
        logger.debug("[+] is_windows: %r" % self.is_windows)
        logger.debug("[+] is_ida64: %r" % self.is_ida64)
        self.my_path = os.path.abspath(__file__)
        self.temp_path = None
        self._init_target()
        # settings (form)
        # todo: load from configfile if available.
        self.output_path = None
        self.chk_annotate_stackvar_size = False
        self.chk_annotate_xrefs = False
        self.chk_decompile_imports = False
        self.chk_decompile_imports_recursive = False
        self.chk_decompile_alternative = False
        # self.ida_home = idaapi.idadir(".")
        self.ida_home = GetIdaDirectory()
        # wait for ida analysis to finish
        self.wait_for_analysis_to_finish()
        if not idaapi.init_hexrays_plugin():
            logger.warning("forcing hexrays to load...")
            self.load_plugin_decompiler()
        if not idaapi.init_hexrays_plugin():
            raise Exception("hexrays decompiler is not available :(")

    def _init_target(self):
        self.target_path = idc.GetInputFilePath()
        self.target_file = idc.GetInputFile()
        self.target_dir = os.path.split(self.target_path)[0]
        logger.debug("reinitializing target: %r" % self.target_file)

    def init_tempdir(self):
        self.temp_path = self.temp_path or tempfile.mkdtemp(prefix="idbc_")
        logger.debug("[i] using tempdir: %r" % self.temp_path)

    def remove_tempdir(self):
        if not self.temp_path:
            return
        logger.debug("[i] removing tempdir: %r" % self.temp_path)
        shutil.rmtree(self.temp_path)
        self.temp_path = None

    def wait_for_analysis_to_finish(self):
        logger.debug("[+] waiting for analysis to finish...")
        idaapi.autoWait()
        idc.Wait()
        logger.debug("[+] analysis finished.")

    def load_plugin_decompiler(self):
        # load decompiler plugins (32 and 64 bits, just let it fail)
        logger.debug("[+] trying to load decompiler plugins")
        if self.is_ida64:
            # 64bit plugins
            idc.RunPlugin("hexx64", 0)
        else:
            # 32bit plugins
            idc.RunPlugin("hexrays", 0)
            idc.RunPlugin("hexarm", 0)
        logger.debug("[+] decompiler plugins loaded.")

    def run(self):
        files_decompiled = []
        self._init_target()

        if self.chk_decompile_imports:
            self.init_tempdir()
            if self.chk_decompile_imports_recursive:
                pass
            for image_type, image_name, image_path in self.enumerate_import_images():
                try:
                    self.exec_ida_batch_decompile(target = image_path, output = self.output_path,
                                                  annotate_stackvar_size = self.chk_annotate_stackvar_size,
                                                  annotate_xrefs = self.chk_annotate_xrefs,
                                                  imports = self.chk_decompile_imports,
                                                  recursive = self.chk_decompile_imports_recursive,
                                                  experimental_decomile_cgraph = self.chk_decompile_alternative)
                    files_decompiled.append(image_path)
                except subprocess.CalledProcessError, cpe:
                    logger.warning("[!] failed to decompile %r - %r" % (image_path, cpe))

            self.remove_tempdir()

        if self.chk_annotate_stackvar_size:
            self.annotate_stack_variable_size()
        if self.chk_annotate_xrefs:
            self.annotate_xrefs()

        if self.chk_decompile_alternative:
            raise NotImplemented("Not yet implemented")
            pass
        else:
            self.decompile_all(self.output_path)
            files_decompiled.append(self.target_file)

        logger.info("[+] finished decompiling: %r" % files_decompiled)
        logger.info("    output dir: %s"%self.output_path if self.output_path else self.target_dir)
        return files_decompiled

    def annotate_stack_variable_size(self):
        logger.debug("[+] annotating function stack variables")
        IdaHelper.annotate_functions_with_local_var_size()
        logger.debug("[+] done.")

    def annotate_xrefs(self):
        logger.debug("[+] annotating function xrefs")
        IdaHelper.annotate_xrefs()
        logger.debug("[+] done.")

    def file_is_decompilable(self, path):
        with open(path, 'rb') as ftest:
            magic = ftest.read(4)
            if magic == 'MZ\x90\x00':
                return 'pe/dos'
            elif magic == "\x7fELF":
                return 'elf'
        return None

    def enumerate_import_images(self):
        for import_name in IdaHelper.get_imports():
            logger.debug("[i] trying to find image for %r" % import_name)
            for image_path in glob.glob(os.path.join(self.target_dir, import_name) + '*'):
                image_type = self.file_is_decompilable(image_path)
                if image_type:
                    logger.debug("[i] got image %r as %r" % (image_path, image_type))
                    yield image_type, os.path.split(image_path)[1], image_path
                    # I do not think there's any need to check other files with the same name ?!
                    break

    def enumerate_files(self, recursive=False):
        for root, dirs, files in os.walk(self.target_dir):
            for name in files:
                fpath = os.path.join(root, name)
                logger.debug("[+] checking %r" % fpath)
                try:
                    ftype = self.file_is_decompilable(fpath)
                    if ftype:
                        logger.debug("[+] is candidate %r" % [fpath, ftype])
                        yield ftype, name, fpath
                except IOError: pass

    def decompile_all(self, outfile=None):
        outfile = self._get_suggested_output_filename(outfile or self.target_path)
        logger.warning(outfile)
        logger.debug("[+] trying to decompile %r as %r" % (self.target_file,
                                                           os.path.split(outfile)[1]))
        IdaHelper.decompile_full(outfile)
        logger.debug("[+] finished decompiling %r as %r" % (self.target_file,
                                                            os.path.split(outfile)[1]))

    def _get_suggested_output_filename(self, target):
        # /a/b/c/d/e/bin.ext
        # target is a directory
        if os.path.isdir(target):
            fname, fext = os.path.splitext(self.target_file)
            return '%s.c' % os.path.join(target, fname)
        # target is not a directory
        root, fname = os.path.split(target)
        if fname:
            fname, fext = os.path.splitext(fname)  # bin,ext
        else:
            fname, fext = os.path.splitext(self.target_file)

        # obsolete
        # suggested_outpath = '%s.c'%os.path.join(root,fname)
        # if not os.path.exists(suggested_outpath):
        #    return suggested_outpath
        return '%s.c' % os.path.join(root, fname)

    def exec_ida_batch_decompile(self, target, output, annotate_stackvar_size, annotate_xrefs, imports, recursive,
                                 experimental_decomile_cgraph):
        logger.debug("[+] batch decompile %r" % target)
        # todo: pass commandlines,
        # todo parse commandline
        script_args = ['--output=%s' % output]
        if annotate_stackvar_size:
            script_args.append("--annotate-stackvar-size")
        if annotate_xrefs:
            script_args.append("--annotate-xrefs")
        if imports:
            script_args.append("--imports")
        if recursive:
            script_args.append("--recursive")
        if experimental_decomile_cgraph:
            script_args.append("--experimental-decompile-cgraph")

        script_args = ['\\"%s\\"' % a for a in script_args]
        command = "%s %s" % (self.my_path, ' '.join(script_args))
        self._exec_ida_batch(target, command)

    def _exec_ida_batch(self, target, command):
        # build exe path
        if self.is_windows:
            ida_exe = os.path.join(self.ida_home, 'idaw64.exe' if self.is_ida64 else 'idaw.exe')
        else:
            ida_exe = os.path.join(self.ida_home, 'idal64' if self.is_ida64 else 'idal')
        '''
        https://www.hex-rays.com/products/ida/support/idadoc/417.shtml
        -B  ..  Batch mode
        -M  ..  disable mouse
        -c  ..  create new database
        -o  ..  database output path
        -S  ..  execute script
        '''
        #temp_path = os.path.join(self.temp_path, os.path.splitext(os.path.split(target)[1])[0] + '.idb')
        cmd = [ida_exe, '-B', '-M', '-c', '-o"%s"'%self.temp_path if self.temp_path else '', '-S"%s"' % command, '"' + target + '"']
        logger.debug(' '.join(cmd))
        logger.debug('[+] executing: %r' % cmd)
        #return 0
        # TODO: INSECURE!
        return subprocess.check_call(' '.join(cmd), shell=True)


class TestEmbeddedChooserClass(Choose,Choose2):
    """
    A simple chooser to be used as an embedded chooser
    """
    def __init__(self, title, nb = 5, flags=0):
        Choose.__init__(self,
                         title,
                         [["Type", 10], ["Name", 10], ["Path", 30]],
                         flags=flags)
        Choose2.__init__(self,
                         title,
                         [ ["Type", 10], ["Name", 10],  ["Path", 30] ],
                         embedded=True, width=50, height=10, flags=flags)
        self.n = 0
        self.items = []
        self.icon = 5
        self.selcount = 0

        self.selected = []

    def OnClose(self):
        pass

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

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

    def OnRefresh(self, n):
        print "refresh %s"%n

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

    def getSelected(self):
        for idx in self.selected:
            yield self.items[idx-1]

    def addItem(self, e):
        if e not in self.items:
            self.items.append(e)


class DecompileBatchForm(Form):
    """
    Form to prompt for target file, backup file, and the address
    range to save patched bytes.
    """

    def __init__(self, idbctrl, enumerate_imports=True, enumerate_other=False):
        self.idbctrl = idbctrl
        self.EChooser = TestEmbeddedChooserClass("Batch Decompile", flags=Choose2.CH_MULTI)
        self.propagateItems(enumerate_imports=enumerate_imports, enumerate_other=enumerate_other)
        Form.__init__(self,
                      r"""Ida Batch Decompile ...
{FormChangeCb}
<##Target    :{target}>
<##OutputPath:{outputPath}>
<##Annotate StackVar Size:{chkAnnotateStackVars}>
<##Annotate Func XRefs   :{chkAnnotateXrefs}>
<##Process Imports       :{chkDecompileImports}>
<##Cgraph (experimental) :{chkDecompileAlternative}>{cGroup1}>


<##Scan Target Directory:{btnLoad}> <##Recursive:{chkDecompileImportsRecursive}>{cGroup2}>
<##Decompile!:{btnProcessFiles}>
<Please select items to decompile:{cEChooser}>


""", {
                          'target': Form.FileInput(swidth=50, open=True, value=idbctrl.target_path),
                          'outputPath': Form.DirInput(swidth=50, value=idbctrl.output_path),
                          'cGroup1': Form.ChkGroupControl(("chkAnnotateStackVars", "chkAnnotateXrefs",
                                                           "chkDecompileImports",
                                                           "chkDecompileAlternative")),
                          'cGroup2': Form.ChkGroupControl(("chkDecompileImportsRecursive", )),
                          'FormChangeCb': Form.FormChangeCb(self.OnFormChange),
                          'btnLoad':  Form.ButtonInput(self.OnButtonLoad),
                          'btnProcessFiles': Form.ButtonInput(self.OnButtonProcess),
                          'cEChooser': Form.EmbeddedChooserControl(self.EChooser),
                      })
        self.Compile()

    def propagateItems(self, enumerate_imports=False, enumerate_other=False):
        self.EChooser.addItem([self.idbctrl.file_is_decompilable(self.idbctrl.target_path),
                               os.path.split(self.idbctrl.target_path)[1],
                               self.idbctrl.target_path])

        if enumerate_imports:
            for candidate in self.idbctrl.enumerate_import_images():
                self.EChooser.addItem(list(candidate))
        if enumerate_other:
            for candidate in self.idbctrl.enumerate_files(recursive=self.chkDecompileImportsRecursive.checked):
                self.EChooser.addItem(list(candidate))

    def OnButtonProcess(self, code=0):
        ### process selected files
        if not len(list(self.EChooser.getSelected())):
            logger.warning("[!] Aborting. Please select at least one item from the list!")
            return

        self.idbctrl.target = self.target.value
        outputPath = self.GetControlValue(self.outputPath)
        if outputPath == '' or os.path.exists(outputPath):
            self.idbctrl.output_path = outputPath
        else:
            logger.warning("[!!] output path not valid! %r" % outputPath)
            self.idbctrl.output_path = None

        self.idbctrl.chk_annotate_stackvar_size = self.chkAnnotateStackVars.checked
        self.idbctrl.chk_decompile_imports = self.chkDecompileImports.checked
        self.idbctrl.chk_decompile_imports_recursive = self.chkDecompileImportsRecursive.checked
        self.idbctrl.chk_annotate_xrefs = self.chkAnnotateXrefs.checked
        self.idbctrl.chk_decompile_alternative = self.chkDecompileAlternative.checked
        logger.debug("[+] config updated")

        files_decompiled = []
        decompile_main_binary = False

        self.idbctrl.init_tempdir()
        for _type, name, image_path in self.EChooser.getSelected():
            if image_path is self.idbctrl.target_path:
                decompile_main_binary = True
                continue
            try:
                self.idbctrl.exec_ida_batch_decompile(target=image_path, output=outputPath,
                                              annotate_stackvar_size=self.idbctrl.chk_annotate_stackvar_size,
                                              annotate_xrefs=self.idbctrl.chk_annotate_xrefs,
                                              imports=self.idbctrl.chk_decompile_imports,
                                              recursive=self.idbctrl.chk_decompile_imports_recursive,
                                              experimental_decomile_cgraph=self.idbctrl.chk_decompile_alternative)
                files_decompiled.append(image_path)
            except subprocess.CalledProcessError, cpe:
                logger.warning("[!] failed to decompile %r - %r" % (image_path, cpe))

        self.idbctrl.remove_tempdir()
        ## process current file
        if decompile_main_binary:
            # well, loop here even though we know it can only
            logger.debug("[+] decompiling current file...")
            files_decompiled += self.idbctrl.run()  # decompile main binary
            logger.info("[+] finished decompiling: %r" % files_decompiled)
            logger.info("    output dir: %s" % self.idbctrl.output_path if self.idbctrl.output_path else self.idbctrl.target_dir)

    def OnButtonLoad(self, code=0):
        self.Close(0)
        self.propagateItems(enumerate_other=True, enumerate_imports=True)
        self.Execute()

    def OnFormChange(self, fid):
        # Set initial state
        INIT = -1
        BTN_OK = -2

        if fid == INIT:
            self.EnableField(self.target, False)
            self.EnableField(self.outputPath, True)
            self.EnableField(self.chkDecompileAlternative, False)

        elif fid == BTN_OK:
            # just return
            return True

        # Toggle backup checkbox
        elif fid == self.chkAnnotateStackVars.id:
            self.chkAnnotateStackVars.checked = not self.chkAnnotateStackVars.checked
        elif fid == self.chkDecompileImports.id:
            self.chkDecompileImports.checked = not self.chkDecompileImports.checked
        elif fid == self.chkDecompileImportsRecursive.id:
            self.chkDecompileImportsRecursive.checked = not self.chkDecompileImportsRecursive.checked
        elif fid == self.chkDecompileAlternative.id:
            self.chkDecompileAlternative.checked = not self.chkDecompileAlternative.checked
        elif fid == self.chkAnnotateXrefs.id:
            self.chkAnnotateXrefs.checked = not self.chkAnnotateXrefs.checked

        return False


if idaapi.IDA_SDK_VERSION >= 700:
    class IdaDecompileUiActionHandler(idaapi.action_handler_t):

        def __init__(self, caller):
            idaapi.action_handler_t.__init__(self)
            self.caller = caller

        def activate(self, ctx):
            self.caller.menu_config()
            return 1

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

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


class IdaDecompileBatchPlugin(idaapi.plugin_t):
    """ IDA Plugin Base"""
    flags = idaapi.PLUGIN_FIX
    comment = "Batch Decompile"
    help = "github.com/tintinweb"
    wanted_name = "Ida Batch Decompile"
    wanted_hotkey = ""
    wanted_menu = "File/Produce file/", "{} ...".format(wanted_name)
    wanted_menu_id = 'tintinweb:batchdecompile'

    def init(self):
        NO_HOTKEY = ""
        SETMENU_INS = 0
        NO_ARGS = tuple()

        logger.debug("[+] %s.init()" % self.__class__.__name__)
        self.menuitems = []

        logger.debug("[+] setting up menus for ida version %s" % idaapi.IDA_SDK_VERSION)

        if idaapi.IDA_SDK_VERSION >= 700:
            # >= 700
            action_desc = idaapi.action_desc_t("tintinweb:batchdecompile:load", self.wanted_name, IdaDecompileUiActionHandler(self))
            idaapi.register_action(action_desc)
            idaapi.attach_action_to_menu(''.join(self.wanted_menu), "tintinweb:batchdecompile:load", idaapi.SETMENU_APP)

        else:
            menu = idaapi.add_menu_item(self.wanted_menu[0],
                                        self.wanted_menu[1],
                                        NO_HOTKEY,
                                        SETMENU_INS,
                                        self.menu_config,
                                        NO_ARGS)

            self.menuitems.append(menu)

        return idaapi.PLUGIN_KEEP

    def run(self, arg=None):
        logger.debug("[+] %s.run()" % self.__class__.__name__)

    def term(self):
        logger.debug("[+] %s.term()" % self.__class__.__name__)
        if idaapi.IDA_SDK_VERSION < 700:
            for menu in self.menuitems:
                idaapi.del_menu_item(menu)

    def menu_config(self):
        logger.debug("[+] %s.menu_config()" % self.__class__.__name__)
        self.idbctrl._init_target() # force target reinit
        DecompileBatchForm(self.idbctrl).Execute()

    def set_ctrl(self, idbctrl):
        logger.debug("[+] %s.set_ctrl(%r)" % (self.__class__.__name__, idbctrl))
        self.idbctrl = idbctrl


def PLUGIN_ENTRY(mode=None):
    """ check execution mode:

        a) as Plugin, return plugin object
        b) as script as part of a batch execution, do not spawn plugin object
     """
    logging.basicConfig(level=logging.DEBUG,
                        format="[%(name)s/%(process)s][%(levelname)-10s] [%(module)s.%(funcName)-14s] %(message)s")
    logger.setLevel(logging.DEBUG)
    # always wait for analysis to finish
    logger.debug("[+] initializing IdaDecompileBatchPlugin")
    # create our controller interface
    idbctrl = IdaDecompileBatchController()
    # parse cmdline
    if mode == '__main__':
        # cmdline mode
        if len(idc.ARGV) > 1:
            # cmdline batch mode
            logger.debug("[+] Mode: commandline")
            parser = OptionParser()
            parser.add_option("-o", "--output", dest="output",
                              help="output path")
            parser.add_option("-S", "--annotate-stackvar-size",
                              action="store_true", default=False,
                              help="Generate stack variable size annotations")
            parser.add_option("-X", "--annotate-xrefs",
                              action="store_true", default=False,
                              help="Generate xref annotations")
            parser.add_option("-I", "--imports",
                              action="store_true", default=False,
                              help="try to decompile files referenced in IAT")
            parser.add_option("-R", "--recursive",
                              action="store_true", default=False,
                              help="Recursive decompile files/imports")
            parser.add_option("-Z", "--experimental-decompile-cgraph",
                              action="store_true", default=False,
                              help="[experimental] decompile funcs referenced in calltree manually")

            options, args = parser.parse_args(idc.ARGV[1:])
            # set options
            idbctrl.output_path = options.output
            idbctrl.chk_annotate_stackvar_size = options.annotate_stackvar_size
            idbctrl.chk_annotate_xrefs = options.annotate_xrefs
            idbctrl.chk_decompile_imports = options.imports
            idbctrl.chk_decompile_imports_recursive = options.recursive
            idbctrl.chk_decompile_alternative = options.experimental_decompile_cgraph
            # set all the idbctrl checkboxes and files
            idbctrl.run()
            idc.Exit(0)
            # return

        logger.debug("[+] Mode: commandline w/o args")
        # PluginMode
        plugin = IdaDecompileBatchPlugin()
        plugin.set_ctrl(idbctrl=idbctrl)
        plugin.init()
        logger.info("[i] %s loaded, see Menu: %s" % (IdaDecompileBatchPlugin.wanted_name,
                                                     IdaDecompileBatchPlugin.wanted_menu))
        #plugin.menu_config()
        return plugin

    else:
        logger.debug("[+] Mode: plugin")
        # PluginMode
        plugin = IdaDecompileBatchPlugin()
        plugin.set_ctrl(idbctrl=idbctrl)
        return plugin


if __name__ == '__main__':
    PLUGIN_ENTRY(mode=__name__)


repository: idapython/src
file: examples/misc/cvt64/py_cvt64_sample.py
description: IDAPython project for Hex-Ray's IDA Pro
wanted_name: IDAPython: CVT64 sample
comment: IDAPython: An example how to implement CVT64 functionality
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/misc/cvt64/py_cvt64_sample.py

content:
"""
summary: add 64-bit (.idb->.i64) conversion capabilities to custom plugins

description:
  For more infortmation see SDK/plugins/cvt64_sample example

level: advanced
"""

import idaapi
import ida_idaapi
import ida_netnode

SAMPLE_NETNODE_NAME = "$ cvt64 py_sample netnode"
DEVICE_INDEX   = idaapi.BADADDR     # -1
IDPFLAGS_INDEX = idaapi.BADADDR     # -1
HASH_COMMENT   = "Comment"
HASH_ADDRESS   = "Address"

#--------------------------------------------------------------------------
class idp_listener_t(idaapi.IDP_Hooks):
    def __init__(self):
        idaapi.IDP_Hooks.__init__(self)

    def ev_cvt64_hashval(self, node, tag, name, data):
        helper = idaapi.netnode(SAMPLE_NETNODE_NAME)
        if helper == node and tag == ida_netnode.htag:
            if name == HASH_COMMENT:
                comment = helper.hashstr(name)
                helper.hashset_buf(name, comment)
                return 1
            if name == HASH_ADDRESS:
                address = helper.hashval_long(name)
                if address == ida_idaapi.BADADDR32:
                    address = ida_idaapi.BADADDR
                    helper.hashset_idx(name, address)
                return 1
        return 0

    def ev_cvt64_supval(self, node, tag, idx, data):
        helper = idaapi.netnode(SAMPLE_NETNODE_NAME)
        if helper == node:
            if tag == ida_netnode.stag and idx == ida_idaapi.BADADDR32:
                helper.supset(DEVICE_INDEX, data)
                return 1
            if tag == ida_netnode.atag and len(data):
                if idx == ida_idaapi.BADADDR32:
                    idx = IDPFLAGS_INDEX
                val = int.from_bytes(data, 'little')
                if val == ida_idaapi.BADADDR32:
                    val = ida_idaapi.BADADDR
                helper.altset(idx, val)
                return 1
        return 0

#--------------------------------------------------------------------------
# This class is instantiated once per each opened database.
class cvt64_ctx_t(idaapi.plugmod_t):

    def __init__(self):
        self.prochook = idp_listener_t()
        self.prochook.hook()

    def __del__(self):
        self.prochook.unhook()

    def run(self, arg):
        pass

#--------------------------------------------------------------------------
# This class is instantiated when IDA loads the plugin.
class cvt64_sample_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_MULTI | idaapi.PLUGIN_MOD
    comment = "IDAPython: An example how to implement CVT64 functionality"
    wanted_name = "IDAPython: CVT64 sample"
    wanted_hotkey = ""
    help = ""

    def init(self):
        return cvt64_ctx_t()

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


repository: gilboz/ida_kernelcache_ng
file: ida_kernelcache/ida_plugin.py
wanted_name: IDA KernelCache NG
comment: Based on ida_kernelcache_ng by @_bazad! Now maintained by gilboz
created_at: 2024-12-15 15:35:59+00:00
pushed_at: 2025-05-07 15:57:34+00:00
url:
https://github.com/gilboz/ida_kernelcache_ng/blob/f19e5ceeed0eaa493aa2a0f756c4835af418c525/ida_kernelcache/ida_plugin.py

content:
"""
Register a new IDA Plugin using the python api.
Updated to conform to the new plugin API: https://docs.hex-rays.com/developer-guide/idapython/how-to-create-a-plugin

See: https://hex-rays.com/blog/scriptable-plugins/ and loader.hpp in the IDA SDK
See: https://hex-rays.com/blog/augmenting-ida-ui-with-your-own-actions/ and the Pacxplorer plugin by Lipner..
"""
import idaapi
from ida_kernelcache.kernelcache import KernelCache
from ida_kernelcache.ida_helpers.hooks import IDPHooks


class IDAKernelCachePluginMod(idaapi.plugmod_t):

    def __init__(self):
        super().__init__()
        print('[ida_kernelcache_ng] plugmod constructor called!')
        self.kc = KernelCache()
        self.idp_hooks = IDPHooks()
        # self.idp_hooks.hook()
        print('[ida_kernelcache_ng] placed procmod hooks..')

    def __del__(self):
        print('[ida_kernelcache_ng] plugmod destructor called!')
        if self.idp_hooks is not None:
            print('[ida_kernelcache_ng] unhooking procmod..')
            # self.idp_hooks.unhook()

    def run(self, arg):
        print('[ida_kernelcache_ng] plugmod run called starting to process!!')
        self.kc.process()


class IDAKernelCachePlugin(idaapi.plugin_t):
    """
    TODO: add more menu items to improve UI/UX
    TODO: cache data in the IDB for persistency?
    """
    wanted_name = "IDA KernelCache NG"
    # wanted_hotkey = "Meta-K"

    flags = idaapi.PLUGIN_MULTI | idaapi.PLUGIN_MOD

    comment = "Based on ida_kernelcache_ng by @_bazad! Now maintained by gilboz"
    help = "Helps with iOS KernelCache reversing"

    def init(self) -> idaapi.plugmod_t | None:
        """
        TODO: when I set the PLUGIN_PROC flag the plugin is initialized too early, IDA still doesn't detect the input file as a kernelcache and we alawys skip it
        """
        # We only want to be loaded when analyzing kernelcaches
        if not KernelCache.is_supported():
            print('[ida_kernelcache_ng] not a kernelcache, skipping..')
            return None  # None is the same as PLUGIN_SKIP

        print('[ida_kernelcache_ng] input file is supported loading plugin..')
        return IDAKernelCachePluginMod()


def PLUGIN_ENTRY():
    return IDAKernelCachePlugin()


repository: opensrcsec/ida-linux-alternatives
file: linux_alternatives.py
description: IDA Pro Python plugin to analyze and annotate Linux kernel alternatives
comment: Analyze and annotate linux kernel alternatives
created_at: 2021-11-28 14:51:38+00:00
pushed_at: 2021-12-15 14:46:37+00:00
url:
https://github.com/opensrcsec/ida-linux-alternatives/blob/7a5aaea3cc0526d2e244dbb0860e82b117aa4a2e/linux_alternatives.py

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

from linux_alternatives_lib.lib import *

PLUGIN_NAME = "Linux Alternatives"
WINDOW_NAME = "Alternatives"


class Alternatives_viewer_t(PluginForm):
    def __init__(self, alternatives, column_width=150):
        super(Alternatives_viewer_t, self).__init__()
        self.alternatives = alternatives
        self.column_width = column_width

    def OnCreate(self, form):
        self.parent = self.FormToPyQtWidget(form)
        self.PopulateForm()

    def PopulateForm(self):
        # Create layout
        layout = QtWidgets.QVBoxLayout()

        # Table View
        self.table = QtWidgets.QTableWidget()
        self.table.setSortingEnabled(True)
        self.table.setRowCount(len(self.alternatives["rows"]))
        self.table.setColumnCount(len(self.alternatives["header"]))
        for i in range(len(self.alternatives["header"])):
            self.table.setColumnWidth(i, self.column_width)

        self.table.setHorizontalHeaderLabels(self.alternatives["header"])
        for i, row in enumerate(self.alternatives["rows"]):
            for j, elem in enumerate(row):
                fmt = "%s" if isinstance(elem, str) else "0x%04x"
                self.table.setItem(i, j, QtWidgets.QTableWidgetItem(fmt % elem))

        self.table.cellDoubleClicked.connect(self.jumpto)

        layout.addWidget(self.table)
        self.parent.setLayout(layout)

    def jumpto(self, row, column):
        try:
            jumpto(int(self.table.item(row, max(column, 1)).text(), 16))
        except:
            pass

    def OnClose(self, form):
        self.alternatives = None


class File_loader_t(ida_kernwin.action_handler_t):
    def __init__(self, action_name, node):
        ida_kernwin.action_handler_t.__init__(self)
        self.action_name = action_name
        self.node = node

    def activate(self, ctx):
        filepath = ida_kernwin.ask_file(False, None, self.action_name)
        self.node[0].supset(self.node[1], filepath)
        print("Loaded %s file" % filepath)

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


class Remover_t(ida_kernwin.action_handler_t):
    def __init__(self, reset, action_name):
        ida_kernwin.action_handler_t.__init__(self)
        self.reset = reset
        self.action_name = action_name

    def activate(self, ctx):
        self.reset()

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


class Patch_input_t(ida_kernwin.action_handler_t):
    def __init__(self, action_name, cpufeat_node):
        ida_kernwin.action_handler_t.__init__(self)
        self.action_name = action_name
        self.cpufeat_node = cpufeat_node
        self.current_flags = None

    def activate(self, ctx):
        prompt = "Enter comma-separated list of CPU features"
        defval = self.current_flags if self.current_flags else ""

        patch_features_str = ida_kernwin.ask_str(defval, 0, prompt)
        # Ignore when nothing has been specified
        if not patch_features_str:
            return

        self.current_flags = patch_features_str.upper()

        patcher = Alternative_patcher_t(self.cpufeat_node)
        patcher.patch(self.current_flags)

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


class Linux_alternatives_t(plugin_t):
    flags = 0
    comment = "Analyze and annotate linux kernel alternatives"
    help = ""
    wanted_name = PLUGIN_NAME
    wanted_hotkey = "Alt-F9"

    TOPMENU = "Edit"

    cpufeat_name = "cpufeatures.h"
    cpufeat_node = [None, 0]

    cpu_flags = None
    alt_instr_struct = None

    cpufeatures_action_name = None

    remove_action_name = "Remove alternative comments"

    patch_input_name = "patch_input"
    patch_input_action_name = "Patch selected alternatives"
    patch_input_node = [None, 0]

    def __init__(self):
        self.view = None
        self.alternatives = {}

        self.cpufeatures_action_name = "Import %s file" % self.cpufeat_name

    def init(self):
        ida_kernwin.create_menu(PLUGIN_NAME, PLUGIN_NAME, "%s/" % self.TOPMENU)

        self.cpufeat_node[0] = ida_netnode.netnode("$ %s" % self.cpufeat_name, self.cpufeat_node[1], True)
        self._register_cpufeatures_action()

        self.patch_input_node[0] = ida_netnode.netnode("$ %s" % self.patch_input_name, self.patch_input_node[1], True)
        self._register_patch_input_action()

        return PLUGIN_KEEP

    def run(self, arg):
        self.reset()
        print("Running %s plugin..." % PLUGIN_NAME)

        alt_gen = Alternative_generator_t(self.cpufeat_node)
        struct_metadata = alt_gen.alt_instr_struct.get_struct_metadata()

        alternatives = alt_gen.gen_alternatives(alt_gen.add_alternatives_cmts)
        self._register_remove_action()

        self.alternatives["header"] = ["index"] + [name for name, _, _ in struct_metadata]
        for _, rows in alternatives.items():
            for row in rows:
                self.alternatives["rows"].append(row)
        self.display_alternatives()

    def term(self):
        self._reset_alternatives()

        if self.view:
            self.view.Close(ida_kernwin.WCLS_DONT_SAVE_SIZE)

        ida_kernwin.unregister_action(self.remove_action_name)
        ida_kernwin.unregister_action(self.cpufeatures_action_name)

    def reset(self):
        self.term()
        self.init()

    def _reset_alternatives(self):
        # Remove applied alternatives comments
        for row in self.alternatives.get("rows", []):
            ea, size = row[1], row[4]

            delete_extra_cmts(ea, E_PREV)
            delete_extra_cmts(ea + size, E_PREV)

        self.alternatives["header"] = []
        self.alternatives["rows"] = []

    def _register_cpufeatures_action(self):
        action_name = self.cpufeatures_action_name
        desc = ida_kernwin.action_desc_t(action_name, action_name, File_loader_t(action_name, self.cpufeat_node))
        ida_kernwin.register_action(desc)
        ida_kernwin.attach_action_to_menu("%s/%s/" % (self.TOPMENU, PLUGIN_NAME), action_name, ida_kernwin.SETMENU_INS)

    def _register_patch_input_action(self):
        action_name = self.patch_input_action_name
        desc = ida_kernwin.action_desc_t(action_name, action_name, Patch_input_t(action_name, self.cpufeat_node))
        ida_kernwin.register_action(desc)
        ida_kernwin.attach_action_to_menu("%s/%s/" % (self.TOPMENU, PLUGIN_NAME), action_name, ida_kernwin.SETMENU_INS)

    def _register_remove_action(self):
        action_name = self.remove_action_name
        desc = ida_kernwin.action_desc_t(action_name, action_name, Remover_t(self.reset, action_name))
        ida_kernwin.register_action(desc)
        ida_kernwin.attach_action_to_menu("%s/%s/" % (self.TOPMENU, PLUGIN_NAME), action_name, ida_kernwin.SETMENU_INS)

    def display_alternatives(self):
        self.view = Alternatives_viewer_t(self.alternatives)
        self.view.Show("Alternatives")


def PLUGIN_ENTRY():
    return Linux_alternatives_t()


repository: nihilus/IDASimulator
file: src/plugins/idasimulator.py
wanted_name: IDA Simulator
comment: IDA Simulator Plugin
created_at: 2014-09-02 23:21:16+00:00
pushed_at: 2014-09-02 23:28:13+00:00
description:
IDASimulator is a plugin that extends IDA's conditional breakpoint support, making it easy to augment / replace complex executable code inside a debugged process with Python code. Specifically, IDASimulator makes use of conditional breakpoints in the IDA debugger to hijack the execution flow of a process and invoke Python handler functions whenever particular code blocks are executed. With support for multiple target architectures, it handles details such as register initialization, memory allocation, pointers, function arguments and return values seamlessly and transparently, making it easy to replace, modify and subvert existing functionality (or lack thereof) in the target process. IDASimulator also includes the IDASim python module, on which IDASimulator is based. This allows for all of the features of IDASimulator to be integrated into more complex IDAPython scripts. IDASimulator currently supports the x86, x86_64, ARM and MIPS32 architectures. Porting to other architectures is very easy.

url:
https://github.com/nihilus/IDASimulator/blob/06e385e79fd57b860e0d7ee8089a769a2f03e829/src/plugins/idasimulator.py

content:
import os
import pickle
import inspect
import idaapi
import idautils
import idc
import idasim

IDASIM = None

class IDASimConfiguration:
	'''
	Responsible for loading, saving and cleaning the idasimulator configuration file.
	Configuration data is a dictionary stored in pickle format:

		cfg = {
			'/path/to/database.idb' : {
						'handlers'   : [enabled handlers],
						'startup'    : 'startup script',
						'startname'  : named location to set breakpoint for setting startup values,
						'membase'    : membase value
			}
		}
	'''
	CONFIG_FILE = 'idasimulator.cfg'

	def __init__(self, sim):
		'''
		Class constructor.

		@sim - idasimulator_t class instance.

		Returns None.
		'''
		self.cfg = None
		self.sim = sim
		self.idb = idc.GetIdbPath()
		self.confile = os.path.join(idaapi.get_user_idadir(), self.CONFIG_FILE)

	def _load_config_file(self):
		'''
		Loads the entire configuration file, cleaning out stale config entries in the process.

		Returns the entire configuration dictionary.
		'''
		cfg = {}
		stale = []

		try:
			cfg = pickle.load(open(self.confile, "rb"))

			# The IDB path is used as the configuration key. If a IDB no longer exists, its config data is useless.
			# Mark any IDB's that no longer exist.
			for (idb, config) in cfg.iteritems():
				if not os.path.exists(idb):
					stale.append(idb)

			# Delete any stale config entries and save the data back to the config file.
			if len(stale) > 0:
				for idb in stale:
					del cfg[idb]
				self._save_config_file(cfg)
		except Exception, e:
			pass

		return cfg

	def _save_config_file(self, fdata):
		'''
		Saves data to the config file, in pickle format.
		
		@fdata - Configuration file dictionary.

		Returns None.
		'''
		try:
			pickle.dump(fdata, open(self.confile, "wb"))
		except Exception, e:
			print "Failed to save %s: %s" % (self.confile, str(e))

	def _load_config_data(self, idb_path):
		'''
		Loads configuration data for this IDB from the config file.

		Returns this IDB's configuration data.
		'''
		data = {}

		if os.path.exists(self.confile):
			cfgdata = self._load_config_file()
			if cfgdata.has_key(idb_path):
				data = cfgdata[idb_path]
		return data

	def _populate_config_data(self, data):
		'''
		Populates the current running configuration from data.
		
		@data - Configuration dictionary.

		Returns None.
		'''
		for name in data['handlers']:
			self.sim.EnableHandler(name)

		self.sim.SetInitValues(data['startname'], data['startup'])
		self.sim.idasim.mmu.base(data['membase'])

	def _save_config_data(self, idb_path, fdata):
		'''
		Saves the current running configuration to disk.

		@idb_path - Path to the IDB file.
		@fdata    - Configuration file dictionary.

		Returns None.
		'''
		fdata[idb_path] = {}
		fdata[idb_path]['handlers'] = self.sim.EnabledHandlers()
		fdata[idb_path]['membase'] = self.sim.idasim.mmu.base()

		(start_name, start_script) = self.sim.GetInitValues()
		fdata[idb_path]['startname'] = start_name
		fdata[idb_path]['startup'] = start_script

		self._save_config_file(fdata)

	def Load(self):
		'''
		Loads the saved configuration data into the running configuration.

		Returns None.
		'''
		data = self._load_config_data(self.idb)
		if data:
			self._populate_config_data(data)

	def Save(self):
		'''
		Saves the running configuration to disk.
	
		Returns None.
		'''
		fdata = self._load_config_file()
		self._save_config_data(self.idb, fdata)


class IDASimFunctionChooser(idaapi.Choose2):
	'''
	Primary IDASimulator UI.
	'''

	def __init__(self, sim):
		idaapi.Choose2.__init__(self, "IDA Simulator", [ 	
									["Handler", 20 | Choose2.CHCOL_PLAIN], 
									["Name", 15 | Choose2.CHCOL_PLAIN], 
									["Description", 30 | Choose2.CHCOL_PLAIN], 
									["Status", 10 | Choose2.CHCOL_PLAIN], 
								])
		self.icon = 41
		self.sim = sim
		self.save_cmd = None
		self.quit_cmd = None
		self.goto_cmd = None
		self.reset_cmd = None
		self.mbase_cmd = None
		self.toggle_cmd = None
		self.config_cmd = None
		self.enable_all_cmd = None
		self.disable_all_cmd = None

		self.PopulateItems()

	def PopulateItems(self):
		'''
		Populates the chooser window with named locations that have registered handlers.
		'''
		self.items = []

		for (name, info) in self.sim.functions.iteritems():
			addr = idc.LocByName(info['function'])

			if addr != idc.BADADDR:
				if self.sim.IsSimulated(name):
					status = "Enabled"
				else:
					status = "Disabled"
				
				self.items.append([name, info['function'], self.sim.GetHandlerDesc(name), status, addr])

	def OnSelectLine(self, n):
		'''
		Invoked when the user double-clicks on a selection in the chooser.
		'''
		self.sim.ToggleHandler(self.items[n][0])
		# Not sure why, but the displayed items aren't refreshed if PopulateItems isn't called here.
		# Interestingly, this is NOT required when OnSelectLine is invoked via the OnCommand method.
		self.PopulateItems()
		self.Refresh()

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

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

	def OnDeleteLine(self, n):
		'''
		Invoked when a user deletes a selection from the chooser.
		'''
		return n

	def OnRefresh(self, n):
		'''
		Refreshes the display.
		'''
		self.sim.Refresh()
		self.PopulateItems()
        	return n

	def OnCommand(self, n, cmd_id):
		'''
		Handles custom right-click commands.
		'''
		if self.sim.idasim is not None:
			if cmd_id == self.reset_cmd:
				self.reset()
			elif cmd_id == self.goto_cmd:
				idc.Jump(self.items[n][-1])
			elif cmd_id == self.toggle_cmd:
				self.OnSelectLine(n)
			elif cmd_id == self.enable_all_cmd:
				self.enable_all()
			elif cmd_id == self.disable_all_cmd:
				self.disable_all()
			elif cmd_id == self.mbase_cmd:
				self.set_mbase()
			elif cmd_id == self.config_cmd:
				self.configure_form()
			elif cmd_id == self.save_cmd:
				self.sim.config.Save()
			elif cmd_id == self.quit_cmd:
				self.quit_idasim()
		return 1

	def OnClose(self):
		'''
		Save the current settings when the chooser window is closed.
		'''
		if self.sim.idasim is not None:
			self.sim.config.Save()
		return None

	def quit_idasim(self):
		'''
		Quits IDASimulator, disabling everything.
		'''
		self.sim.config.Save()
		self.sim.Cleanup(closegui=False)

	def set_mbase(self):
		'''
		Sets the memory base address for the IDASimMMU instance.
		'''
		mbase = AskAddr(self.sim.idasim.mmu.base(), "Configure base memory allocation address")
		if mbase != idc.BADADDR:
			if mbase == 0:
				mbase = idc.BADADDR
			self.sim.idasim.mmu.base(mbase)

	def configure_form(self):
		'''
		Displays the configuration form for setting up startup register values.
		'''
		script_file = AskFile(0, '*.py', 'Select a script to run on process init/attach.')
		if script_file:
			self.sim.SetInitValues(None, open(script_file, 'rb').read())

	def enable_all(self):
		'''
		Enables all handlers.
		'''
		for i in range(0, len(self.items)):
			self.sim.EnableHandler(self.items[i][0])
		self.Refresh()

	def disable_all(self):
		'''
		Disable all handlers.
		'''
		for i in range(0, len(self.items)):
			self.sim.DisableHandler(self.items[i][0])
		self.Refresh()

	def reset(self):
		'''
		Resets all settings to the defaults.
		'''
		if idc.AskYN(0, "Are you sure you want to undo all changes and reset?") == 1:
			self.sim.Reset()
			self.Refresh()

	def show(self):
		'''
		Displays the chooser, initializes the custom right-click options.
		'''
		if self.Show(modal=False) < 0:
			return False
	
		self.toggle_cmd = self.AddCommand("Enable / disable selected handler")
		self.enable_all_cmd = self.AddCommand("Enable all handlers")
		self.disable_all_cmd = self.AddCommand("Disable all handlers")
		self.config_cmd = self.AddCommand("Load startup script")
		self.mbase_cmd = self.AddCommand("Set MMU base address")
		self.reset_cmd = self.AddCommand("Reset to defaults")
		self.goto_cmd = self.AddCommand("Jump to selected name")
		self.save_cmd = self.AddCommand("Save settings")
		self.quit_cmd = self.AddCommand("Quit")
		return True


class idasimulator_t(idaapi.plugin_t):
	'''
	Primary IDASimulator plugin class.
	'''

	flags = 0
	comment = "IDA Simulator Plugin"
	help = "Simulate excutable logic in Python"
	wanted_name = "IDA Simulator"
	wanted_hotkey = ""

	def init(self):
		'''
		Initialize some default values for class variables.
		'''
		self.gui = None
		self.idasim = None
		self.config = None
		self.functions = {}
		self.startup_script = ''
		self.startup_name = None
		self.stubs = True
		self.menu_context = idaapi.add_menu_item("Options/", "Simulate functions and code blocks...", "Alt-0", 0, self.run, (None,))
		return idaapi.PLUGIN_KEEP

	def term(self):
		'''
		Cleanup IDASimulator and breakpoints when terminated.
		'''
		self.Cleanup()
		idaapi.del_menu_item(self.menu_context)
		return None

	def run(self, arg):
		'''
		Initialize IDASimulator and chooser GUI, if not already initialized.
		'''
		global IDASIM

		if IDASIM is None:
			IDASIM = idasim.IDASim()
			print "%s enabled." % self.wanted_name
		
		self.idasim = IDASIM
		
		self.__parse_libraries()
	
		self.config = IDASimConfiguration(self)
		self.config.Load()
			
		self.gui = IDASimFunctionChooser(self)
		self.gui.show()

	def GetInitValues(self):
		'''
		Returns the named initialization location and the array of startup Python statements.
		'''
		return (self.startup_name, self.startup_script)

	def SetInitValues(self, name=None, script=''):
		'''
		Sets the named initialization location and the array of startup Python statements.

		@name  - Named location.
		@lines - Array of tuples (register name, Python statement).

		Returns the named initialization location and the array of startup Python statements.
		'''
		self.idasim.ExecuteOnStart(None, None, disable=True)

		if not name:
			disable = True
		else:
			disable = False
		
		self.startup_name = name
		self.startup_script = script
		self.idasim.ExecuteOnStart(self.startup_script, self.startup_name, disable=disable)
		
		return (self.startup_name, self.startup_script)

	def IsSimulated(self, name):
		'''
		Checks if a named location has an active IDASimulator handler.

		@name - Named location.

		Returns True if a handler is active, False if not.
		'''
		if self.functions.has_key(name):
			return self.functions[name]['enabled']
		else:
			return False

	def GetHandlerDesc(self, name):
		'''
		Get a handler description for a given named location.

		@name - Handler name.

		Returns the handler description, if it exists. Else, returns None.
		'''
		if name == self.startup_name:
			return 'Initialization handler.'
		else:
			return self.functions[name]['description']

	def ToggleHandler(self, name):
		'''
		Enables/disables the handler for the named location.

		@name - Named location.

		Returns None.
		'''
		if self.IsSimulated(name):
			self.DisableHandler(name)
		else:
			self.EnableHandler(name)

	def EnableHandler(self, name):
		'''
		Enables the handler for the named location.

		@name - Named location.

		Returns None.
		'''
		existing_handler = self.idasim.FunctionHandler.GetHandler(self.functions[name]['function'])
		if existing_handler:
			self.DisableHandler(self.__get_handler_name(existing_handler.im_class.__name__, existing_handler.__name__))

		self.idasim.FunctionHandler.RegisterHandler(self.functions[name]['function'], self.functions[name]['handler'], self.stubs)
		self.functions[name]['enabled'] = True
	
	def EnabledHandlers(self):
		'''
		Returns a list of all named locations that have an active handler.
		'''
		return [name for (name, info) in self.functions.iteritems() if info['enabled']]

	def DisableHandler(self, name):
		'''
		Disables the handler for the named location.

		@name - Named location.
		
		Returns None.
		'''
		self.idasim.FunctionHandler.UnregisterHandler(self.functions[name]['function'], self.stubs)
		self.functions[name]['enabled'] = False

	def Refresh(self):
		'''
		Refreshes the internal list of supported handlers.
		'''
		if self.idasim is not None:
			self.__parse_libraries()
		
			for name in self.EnabledHandlers():
				if name != self.startup_name:
					self.idasim.FunctionHandler.RegisterHandler(self.functions[name]['function'], self.functions[name]['handler'], self.stubs)	
					
	def Reset(self):
		'''
		Resets all IDASimulator settings to the defaults.
		'''
		self.SetInitValues(None, None)
		self.idasim.Cleanup()
		self.idasim.mmu.base(idc.BADADDR)
		self.__parse_libraries()

	def Cleanup(self, closegui=True):
		'''
		Cleans up all IDASimulator changes and disables the plugin.
		'''
		global IDASIM

		try:
			if closegui and self.gui is not None:
				self.gui.Close()
		except:
			pass

		try:
			if self.idasim is not None:
				self.idasim.Cleanup()
		except:
			pass

		IDASIM = None
		self.gui = None
		self.idasim = None
		self.functions = {}
		print "%s disabled." % self.wanted_name

	def __get_handler_name(self, class_name, method_name):
		'''
		Builds a handler key name from the class and method names.
		'''
		return class_name + '.' + method_name

	def __generate_handler_entry(self, instance, method, name=None):
		'''
		Creates a single handler dictionary entry for the provided class instance and method.
		'''
		if not name:
			name = method

		handler = getattr(instance, method)
		class_name = instance.__class__.__name__
		handler_name = self.__get_handler_name(class_name, method)

		entry = {}
		entry[handler_name] = {}

		entry[handler_name]['class'] = class_name
		entry[handler_name]['handler'] = handler
		entry[handler_name]['function'] = name
		entry[handler_name]['enabled'] = False

		existing_handler = self.idasim.FunctionHandler.GetHandler(name)
		if existing_handler:
			if self.__get_handler_name(existing_handler.im_class.__name__, existing_handler.__name__) == handler_name:
				entry[handler_name]['enabled'] = True
		try:
			entry[handler_name]['description'] = handler.__doc__.strip().split('\n')[0].strip()
		except:
			entry[handler_name]['description'] = 'Simulates the ' + name + ' function.'

		return entry

	def __parse_library(self, lib):
		'''
		Parses a loaded library for all handlers.

		@lib - Class instance.
		
		Returns a dictionary of handlers.
		'''
		ignore = ['__init__', '__del__', '__enter__', '__exit__']
		handlers = {}
		instance = lib(IDASIM)

		for (name, obj) in inspect.getmembers(lib, inspect.ismethod):
			if name not in ignore:
				handlers.update(self.__generate_handler_entry(instance, name))
		
		self.functions.update(handlers)
		return handlers

	def __parse_libraries(self):
		'''
		Loads/reloads and parses all IDASimulator handlers.
		'''
		import idasimlib
		reload(idasimlib)
		self.functions = {}

		for module_name in dir(idasimlib):
			# Don't process modules whose file names begin with a double underscore
			if not module_name.startswith('__'):
				try:
					module = getattr(idasimlib, module_name)
					reload(module)
					for (class_name, class_obj) in inspect.getmembers(module, inspect.isclass):
						# Don't process classes whose names begin with an underscore
						if not class_name.startswith('_'):
							self.__parse_library(getattr(module, class_name))
				except Exception, e:
					print "WARNING: Failed to load %s: %s" % (module_name, str(e))
					continue

def PLUGIN_ENTRY():
	return idasimulator_t()


repository: dtzxporter/IDA-XMMCpy
file: ida_xmmcpy.py
description: A plugin to copy an XMMWORD to a ready to use C++ intrinsics
wanted_name: XMMCpy
comment: XMMCpy plugin for IDA Pro
created_at: 2019-04-12 17:37:45+00:00
pushed_at: 2019-04-12 17:39:39+00:00
url:
https://github.com/dtzxporter/IDA-XMMCpy/blob/49c663b9773302835eeda195e226624fc110372d/ida_xmmcpy.py

content:
import idaapi
import idautils
import idc
import operator

if idaapi.IDA_SDK_VERSION >= 690:
    from PyQt5.Qt import QApplication
else:
    from PySide.QtGui import QApplication


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 Exporter(Kp_Menu_Context):
    def activate(self, ctx):
        self.plugin.export_xmm()
        return 1


p_initialized = False


class IDAXMMCpy_Plugin_t(idaapi.plugin_t):
    comment = "XMMCpy plugin for IDA Pro"
    help = "todo"
    wanted_name = "XMMCpy"
    wanted_hotkey = "Ctrl-Shift-E"
    flags = idaapi.PLUGIN_KEEP

    def init(self):
        
        try:
            Exporter.register(self, "XMMCpy")
        except:
            pass

        return idaapi.PLUGIN_KEEP

    def term(self):
        pass

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

    def export_xmm_float(self):
        result_float = "_mm_set_ps(%f, %f, %f, %f)" % (idc.GetFloat(get_screen_ea() + 12),
                                                            idc.GetFloat(get_screen_ea() + 8), idc.GetFloat(get_screen_ea() + 4), idc.GetFloat(get_screen_ea()))
        print("XMMCpy Result: %s" % result_float)
        QApplication.clipboard().setText(result_float)

    def export_xmm(self):
        result = "_mm_set_epi32(0x%x, 0x%x, 0x%x, 0x%x)" % (idc.Dword(get_screen_ea() + 12),
                                                            idc.Dword(get_screen_ea() + 8), idc.Dword(get_screen_ea() + 4), idc.Dword(get_screen_ea()))
        
        print("XMMCpy Result: %s" % result)
        QApplication.clipboard().setText(result)


def PLUGIN_ENTRY():
    return IDAXMMCpy_Plugin_t()


repository: PAGalaxyLab/vxhunter
file: firmware_tools/vxhunter_ida_py3.py
description: ToolSet for VxWorks Based Embedded Device Analyses
wanted_name: VxHunter
comment: VxHunter plugin for IDA Pro
created_at: 2017-12-01 03:13:12+00:00
pushed_at: 2024-09-24 07:02:21+00:00
url:
https://github.com/PAGalaxyLab/vxhunter/blob/5984b985963e2e370b536a239771110865d44ef0/firmware_tools/vxhunter_ida_py3.py

content:
# coding=utf-8
import logging
import struct
import idc
import ida_nalt
import ida_bytes
import ida_funcs
import ida_kernwin
import idaapi
import time


default_check_count = 100

known_address = [0x80002000, 0x10000, 0x1000, 0xf2003fe4, 0x100000, 0x107fe0]

function_name_key_words = ['bzero', 'usrInit', 'bfill']

# VxWorks 5.5
vx_5_sym_types = [
    # 0x00,      # Undefined Symbol
    # 0x01,      # Global (external)
    # 0x02,      # Local Absolute
    0x03,      # Global Absolute
    0x04,      # Local .text
    0x05,      # Global .text
    0x06,      # Local Data
    0x07,      # Global Data
    0x08,      # Local BSS
    0x09,      # Global BSS
    0x12,      # Local Common symbol
    0x13,      # Global Common symbol
    0x40,      # Local Symbols related to a PowerPC SDA section
    0x41,      # Global Symbols related to a PowerPC SDA section
    0x80,      # Local symbols related to a PowerPC SDA2 section
    0x81,      # Local symbols related to a PowerPC SDA2 section
]

# VxWorks 6.8
vx_6_sym_types = [
    # 0x00,  # Undefined Symbol
    # 0x01,  # Global (external)
    # 0x02,  # Local Absolute
    0x03,  # Global Absolute
    0x04,  # Local .text
    0x05,  # Global .text
    0x08,  # Local Data
    0x09,  # Global Data
    0x10,  # Local BSS
    0x11,  # Global BSS
    0x20,  # Local Common symbol
    0x21,  # Global Common symbol
    0x40,  # Local Symbols
    0x41,  # Global Symbols
]

need_create_function = [0x04, 0x05]


class VxTarget(object):
    def __init__(self, firmware, vx_version=5, is_big_endian=False, logger=None):
        """
        :param firmware: data of firmware
        :param vx_version: 5 = VxWorks 5.x; 6= VxWorks 6.x
        :param is_big_endian: True = big endian; False = little endian
        :param logger: logger for the target (default: None)
        """
        self.is_big_endian = is_big_endian
        self._vx_version = vx_version
        self.symbol_table_start = None
        self.symbol_table_end = None
        self._string_table = []
        self._symbol_table = []
        self.symbols = []
        self.load_address = None
        self._firmware = firmware
        self._has_symbol = None
        if self._vx_version == 5:
            self._symbol_interval = 16
        elif self._vx_version == 6:
            self._symbol_interval = 20
        self.start_time = None
        self._performance_status = []

        if logger is None:
            self.logger = logging.getLogger(__name__)
            self.logger.setLevel(logging.INFO)
            console_handler = logging.StreamHandler()
            console_format = logging.Formatter('[%(levelname)-8s][%(module)s.%(funcName)s] %(message)s')
            console_handler.setFormatter(console_format)
            self.logger.addHandler(console_handler)
        else:
            self.logger = logger
        self._firmware_info = {
            "load_address": -1,
            "has_symbol": False,
            "vx_version": self._vx_version,
            "big_endian": self.is_big_endian,
            "symbol_table_start": None,
            "symbol_table_end": None,
            "symbols": None
        }

        self.prepare()

    def reset_timer(self):
        self.start_time = time.time()

    def start_timer(self):
        if self.start_time:
            return False
        else:
            self.start_time = time.time()
            return self.start_time

    def get_timer(self):
        if self.start_time:
            return time.time() - self.start_time
        return False

    def prepare(self):
        """ Trying to find symbol from image.

        :return: True if found symbol, False otherwise.
        """
        self.find_symbol_table()
        if self._has_symbol is False:
            return None
        self.logger.debug("has_symbol: {}".format(self._has_symbol))
        self.get_symbol_table()

    def _check_vxworks_endian(self):
        """ Get image endian from image file.

        :return:
        """
        data1 = self._firmware[self.symbol_table_start + 4:self.symbol_table_start + 4 + self._symbol_interval]
        data2 = self._firmware[self.symbol_table_start + 4 + self._symbol_interval:self.symbol_table_start +
                                                                                   4 + self._symbol_interval * 2]
        if data1[0:2] == data2[0:2]:
            self.logger.info("VxWorks endian: Big endian.")
            self.is_big_endian = True
        elif data1[2:4] == data2[2:4]:
            self.logger.info("VxWorks endian: Little endian.")
            self.is_big_endian = False
        else:
            self.logger.info("VxWorks endian unknown. Assuming little endian.")
            self.is_big_endian = False

    def _check_symbol_format(self, offset):
        """ Check offset is symbol table.

        :param offset: offset from image.
        :return: True if offset is symbol table, False otherwise.
        """
        start_offset = offset
        end_offset = offset + self._symbol_interval * default_check_count
        if end_offset > len(self._firmware):
            return False

        check_data = self._firmware[start_offset:end_offset]
        is_big_endian = True
        is_little_endian = True
        # check symbol data match struct
        for i in range(default_check_count):
            check_data_1 = check_data[i * self._symbol_interval:(i + 1) * self._symbol_interval]
            if len(check_data_1) < self._symbol_interval:
                self.logger.debug("check_data_1 length is too small: {}".format(len(check_data_1)))
                break

            if self._check_symbol_format_simple(check_data_1) is False:
                return False

        if self._vx_version == 5:
            self.logger.debug("Check VxWorks 5 symbol format")
            # check is big endian
            for i in range(9):
                check_data_1 = check_data[4 + i * self._symbol_interval:6 + i * self._symbol_interval]
                data2 = check_data[4 + (i + 1) * self._symbol_interval:6 + (i + 1) * self._symbol_interval]
                if check_data_1 != data2:
                    self.logger.debug("VxWorks binary is not big endian.")
                    is_big_endian = False
                    break

            # check is little endian
            for i in range(9):
                check_data_1 = check_data[6 + i * self._symbol_interval:8 + i * self._symbol_interval]
                data2 = check_data[6 + (i + 1) * self._symbol_interval:8 + (i + 1) * self._symbol_interval]
                if check_data_1 != data2:
                    self.logger.debug("VxWorks binary is not little endian.")
                    is_little_endian = False
                    break

            if is_big_endian and is_little_endian:
                return False

            return is_big_endian ^ is_little_endian

        return True

    def _check_symbol_format_simple(self, data):
        """ Check single symbol format is correct.

        :param data: single symbol data.
        :return: True if data is symbol, False otherwise.
        """
        if self._vx_version == 5:
            # Check symbol type is valid
            sym_type = data[14]
            if sym_type not in vx_5_sym_types:
                return False

            # symbol should end with '\x00'
            if data[15] != 0:
                return False

            # Check symbol group is '\x00\x00'
            if data[12:14] != b'\x00\x00':
                return False

            # symbol_name point should not be zero
            if data[4:8] == b'\x00\x00\x00\x00':
                return False

            # symbol value point should not be zero
            if data[8:12] == b'\x00\x00\x00\x00':
                return False

            return True

        elif self._vx_version == 6:
            # Check symbol type is valid
            sym_type = data[18]
            if sym_type not in vx_6_sym_types:
                return False

            # symbol should end with '\x00'
            if data[19] != 0:
                return False

            # Check symbol group is '\x00\x00'
            if data[16:18] != b'\x00\x00':
                return False

            # symbol_name point should not be zero
            if data[4:8] == b'\x00\x00\x00\x00':
                return False

            # TODO: Need handle this problem
            # sometime data[8:12] will be '\x00\x00\x00\x00'
            # if data[8:12] == '\x00\x00\x00\x00':
            #     return False
            return True

        return False

    def find_symbol_table(self):
        """ Find symbol table from image.

        :return:
        """
        self.reset_timer()
        for offset in range(len(self._firmware) - self._symbol_interval):
            if self.symbol_table_start is None:
                # Get first data valid the symbol_format
                if not self._check_symbol_format_simple(self._firmware[offset: offset + self._symbol_interval]):
                    continue

                elif self._check_symbol_format(offset):
                    self.logger.info("symbol table start offset: {:010x}".format(offset))
                    self.symbol_table_start = offset
                    self._has_symbol = True
                    self._firmware_info["has_symbol"] = True
                    break
            else:
                break
        performance_data = "Find symbol table takes {:.3f} seconds".format(self.get_timer())
        self._performance_status.append(performance_data)
        self.logger.debug(performance_data)

        if self.symbol_table_start:
            self.reset_timer()
            for i in range(self.symbol_table_start, len(self._firmware), self._symbol_interval):
                check_data = self._firmware[i:i + self._symbol_interval]

                if len(check_data) < self._symbol_interval:
                    self.logger.debug("check_data length is too small: {}".format(check_data))
                    break

                if self._check_symbol_format_simple(check_data):
                    self.symbol_table_end = i + self._symbol_interval
                    self.logger.debug("self.symbol_table_end: {:010x}".format(self.symbol_table_end))

                else:
                    self.logger.info("Symbol table end offset: {:010x}".format(self.symbol_table_end))
                    break

        else:
            self.logger.error("Didn't find symbol table in this image")
            self._has_symbol = False

        performance_data = "Analyzer symbol table takes {:.3f} seconds".format(self.get_timer())
        self._performance_status.append(performance_data)
        self.logger.debug(performance_data)

    def get_symbol_table(self):
        """ get symbol table data.

        :return: True if get symbol table data successful, False otherwise.
        """
        if self.symbol_table_start and self.symbol_table_end:
            self._check_vxworks_endian()

        else:
            return False

        for i in range(self.symbol_table_start, self.symbol_table_end, self._symbol_interval):
            symbol_name_addr = self._firmware[i + 4:i + 8]
            symbol_dest_addr = self._firmware[i + 8:i + 12]
            symbol_flag = self._firmware[i + self._symbol_interval - 2]
            if self.is_big_endian:
                unpack_format = '>I'
            else:
                unpack_format = '<I'
            symbol_name_addr = int(struct.unpack(unpack_format, symbol_name_addr)[0])
            symbol_dest_addr = int(struct.unpack(unpack_format, symbol_dest_addr)[0])
            self.logger.debug("symbol_name_addr: {}; symbol_dest_addr: {}".format(symbol_name_addr, symbol_dest_addr))
            self._symbol_table.append({'symbol_name_addr': symbol_name_addr, 'symbol_name_length': None, 'symbol_dest_addr': symbol_dest_addr, 'symbol_flag': symbol_flag, 'offset': i})
        self.logger.debug("len(self._symbol_table): %s".format(len(self._symbol_table)))
        self._symbol_table = sorted(self._symbol_table, key=lambda x: x['symbol_name_addr'])
        for i in range(len(self._symbol_table) - 1):
            self._symbol_table[i]['symbol_name_length'] = self._symbol_table[i + 1]['symbol_name_addr'] - \
                                                          self._symbol_table[i]['symbol_name_addr']
        self.logger.debug("len(self._symbol_table): {}".format(len(self._symbol_table)))
        return True

    @staticmethod
    def _is_printable(c):
        """ Check Char is printable.

        :param c: char to check.
        :return: True if char is printable, False otherwise.
        """
        if type(c) is int:
            return 32 <= c <= 126
        elif type(c) is str:
            return 32 <= ord(c) <= 126

    def _is_func_name(self, string):
        """ Check target string is match function name format.

        :param string: string to check.
        :return: True if string is match function name format, False otherwise.
        """
        #
        bad_str = [b'\\', b'%', b'+', b',', b'&', b'/', b')', b'(', b'[', b']']
        # function name length should less than 512 byte
        if len(string) > 512:
            return False

        for data in bad_str:
            if data in string:
                return False

        for c in string:
            if self._is_printable(c) is False:
                return False
        return True

    def _get_prev_string_data(self, offset):
        """ Get previous string from giving offset.

        :param offset: offset of image.
        :return: string data, string start offset, string end offset.
        """
        while offset > 0:
            if self._firmware[offset] != 0:
                start_address = offset
                end_address = offset + 1
                while offset > 0:
                    if self._firmware[offset - 1] == 0:
                        start_address = offset
                        break
                    offset -= 1
                data = self._firmware[start_address:end_address]
                self.logger.debug("data: {}; start_address: {:010x}; end_address: {:010x}".format(data, start_address, end_address))
                return data, start_address, end_address
            else:
                offset -= 1
        self.logger.debug("Done looking for previous string data.")
        return None, None, None

    def _get_next_string_data(self, offset):
        """ Get next string from giving offset.

        :param offset: offset of image.
        :return: string data, string start offset, string end offset.
        """
        while offset < len(self._firmware):
            if self._firmware[offset] != 0:
                start_address = offset
                end_address = offset
                while offset <= len(self._firmware):
                    offset += 1
                    if self._firmware[offset] == 0:
                        end_address = offset
                        break
                data = self._firmware[start_address:end_address]
                return data, start_address, end_address
            else:
                offset += 1
        return None, None, None

    def find_string_table_by_key_function_index(self, key_offset):
        """ Find string table by VxWorks key function name offset in VxWorks image.

        :param key_offset: key function name offset in VxWorks image.
        :return:
        """
        self.logger.debug("Attempting to find string table by key function index with offset {:010x}".format(key_offset))
        temp_str_tab_data = []
        if len(self._symbol_table) > default_check_count:
            count = default_check_count
        else:
            count = len(self._symbol_table)
        start_offset = key_offset
        end_offset = key_offset
        self.logger.debug("Initializing with start_offset = end_offset = {:010x}".format(key_offset))

        while start_offset > 0:
            if self._is_printable(self._firmware[start_offset]) is True:
                # get string from offset
                string, start_address, end_address = self._get_prev_string_data(start_offset)
                self.logger.debug("string: {}; start_address: {:010x}; end_address: {:010x}".format(string, start_address, end_address))
                # check string is function name
                if self._is_func_name(string) is False:
                    if len(temp_str_tab_data) < count:
                        self.logger.error("Can't find any string table with key index: {}".format(string))
                        return None, None
                    else:
                        self.logger.info("Found string table start address at {:010x}".format(start_address))
                        break
                else:
                    temp_str_tab_data.append((string, start_address, end_address))

                # get previous string from offset
                prev_string, prev_start_address, prev_end_address = self._get_prev_string_data(start_address - 1)
                self.logger.debug("prev_string: {}, prev_start_address: {:010x}, prev_end_address: {:010x}".format(prev_string, prev_start_address, prev_end_address))
                if prev_start_address:
                    # strings interval should less than 4
                    if 4 < (start_address - prev_end_address):
                        if len(temp_str_tab_data) < count:
                            self.logger.error("Can't find any string table with key index: {}".format(string))
                            return None, None
                        else:
                            self.logger.info("found string table start address at {:010x}".format(start_address))
                            break
                    else:
                        start_offset = start_address - 1
                        self.logger.debug("start_offset: {}".format(start_offset))
                else:
                    break
            else:
                start_offset -= 1

        while end_offset < len(self._firmware):
            # find first printable char
            if self._is_printable(self._firmware[end_offset]) is True:
                # get string from offset
                string, start_address, end_address = self._get_next_string_data(end_offset)
                # check string is function name
                if self._is_func_name(string) is False:
                    if len(temp_str_tab_data) < count:
                        temp_str_tab_data = []
                        end_offset = end_address
                        continue
                    else:
                        self.logger.info("found string table end at {:010x}".format(end_address))
                        break

                else:
                    temp_str_tab_data.append((string, start_address, end_address))

                # get next string from offset
                next_string, next_start_address, next_end_address = self._get_next_string_data(end_address)
                if next_start_address:
                    # strings interval should less than 4
                    if 4 < (next_start_address - end_address):
                        if len(temp_str_tab_data) < count:
                            self.logger.error("Can't find any string table with key index.")
                            return None, None
                        else:
                            self.logger.info("Found string table end at {:010x}".format(end_address))
                            break
                    else:
                        end_offset = end_address
            else:
                end_offset += 1

        temp_str_tab_data = sorted(temp_str_tab_data, key=lambda x: (x[1]))
        table_start_offset = temp_str_tab_data[0][1]
        table_end_offset = temp_str_tab_data[-1][2]
        self.logger.info("Found a string table at: {:010x} to {:010x}".format(table_start_offset, table_end_offset))
        return table_start_offset, table_end_offset

    def get_string_table(self, str_start_address, str_end_address):
        """ Get string table data from VxWorks image with string table start and end offset.

        :param str_start_address: string table start address.
        :param str_end_address: string table end address.
        :return:
        """
        self._string_table = []
        offset = str_start_address
        address = offset
        str_tab_data = []
        while offset <= str_end_address:
            if self._firmware[offset] == 0:
                while offset <= str_end_address:
                    offset += 1
                    if self._firmware[offset] != 0:
                        next_address = offset
                        string = self._firmware[address:next_address]
                        length = next_address - address
                        str_tab_data.append({'address': address, 'string': string, 'length': length})
                        offset = next_address
                        address = next_address
                        break
            else:
                offset += 1
        self._string_table = str_tab_data

    def _check_fix(self, func_index, str_index):
        """

        :param func_index:
        :param str_index:
        :return:
        """
        try:
            fault_count = 0
            self.logger.debug("Symbol table's first element: {}".format(self._symbol_table[0]))
            if len(self._symbol_table) <= default_check_count:
                count = len(self._symbol_table)
                self.logger.debug("Length of symbol table, {}, is less than default. Setting iteration count to actual length of table, {}.".format(len(self._symbol_table), count))
            else:
                count = default_check_count
                self.logger.debug("Length of symbol table, {}, is greater than default. Setting iteration count to default, {}.".format(len(self._symbol_table), count))
            for i in range(count):

                if (func_index >= len(self._symbol_table)) or (str_index >= len(self._string_table)):
                    self.logger.debug("_check_fix False: func_index greater than length of _symbol_table, or str_index greater than length of _string_table.")
                    return False
                self.logger.debug("str_index: {}; _string_table[str_index]: {}".format(str_index, self._string_table[str_index]))
                self.logger.debug("func_index: {}; _symbol_table[func_index]: {}".format(func_index, self._symbol_table[func_index]))
                if i == count - 1:
                    if fault_count < 10:
                        self.logger.debug("_check_fix True")
                        return True
                    else:
                        self.logger.debug("_check_fix False: Too many faults.")
                        return False

                if self._string_table[str_index]['length'] == self._symbol_table[func_index]['symbol_name_length']:
                    func_index += 1
                    str_index += 1
                    self.logger.debug("_check_fix continue")

                elif self._symbol_table[func_index]['symbol_name_length'] < self._string_table[str_index]['length']:
                    # Sometime Symbol name might point to mid of string.
                    fault_count += 1
                    func_index += 1
                else:
                    self.logger.debug("_check_fix False: symbol_name_length from func_index larger than length from str_index.")
                    return False
        except Exception as e:
            self.logger.debug(e)
            raise

    def find_loading_address(self):
        """ Find VxWorks image load address by automatic analysis.

        :return: Load address if found, None otherwise.
        """

        # Search function keyword in firmware to locate the function string tables.
        self.reset_timer()
        if self._has_symbol is False:
            return None

        for key_word in function_name_key_words:
            prefix_keyword = b'\x00_' + key_word + b'\x00'
            key_word = b'\x00' + key_word + b'\x00'
            if key_word in self._firmware is False and prefix_keyword in self._firmware is False:
                self.logger.info("Firmware does not contain a function named {}".format(key_word))
                return None

        performance_data = "Search function keyword in firmware takes {:.3f} seconds".format(self.get_timer())
        self._performance_status.append(performance_data)
        self.logger.debug(performance_data)

        # Search function keyword in firmware to locate the function string tables.
        self.reset_timer()
        try:
            key_function_index = self._firmware.index(b'\x00' + function_name_key_words[0] + b'\x00')
            self.logger.debug("key_function_index: {}".format(key_function_index))
        except Exception as err:
            # Handler _ prefix symbols
            self.logger.debug("Exception occurred while loading key_function_index: {}. Continuing...".format(err))
            key_function_index = self._firmware.index(b'\x00_' + function_name_key_words[0] + b'\x00')
            self.logger.debug("key_function_index: {}".format(key_function_index))

        str_start_address, str_end_address = self.find_string_table_by_key_function_index(key_function_index)
        self.get_string_table(str_start_address, str_end_address)
        performance_data = "Get function string table takes {:.3f} seconds".format(self.get_timer())
        self._performance_status.append(performance_data)
        self.logger.debug(performance_data)

        # TODO: Need improve performance
        self.reset_timer()
        self.logger.info("Starting loading address analysis")
        for str_index in range(len(self._string_table)):
            for func_index in range(len(self._symbol_table)):
                self.logger.debug("self._string_table[str_index]['length']: {}".format(self._string_table[str_index]['length']))
                self.logger.debug("self._symbol_table[func_index]['symbol_name_length']: {}".format(self._symbol_table[func_index]['symbol_name_length']))
                if self._string_table[str_index]['length'] == self._symbol_table[func_index]['symbol_name_length']:
                    if self._check_fix(func_index, str_index) is True:
                        self.logger.debug("self._symbol_table[func_index]['symbol_name_addr']: {}".format(self._symbol_table[func_index]['symbol_name_addr']))
                        self.logger.debug("self._string_table[str_index]['address']: %s" % self._string_table[str_index]['address'])
                        self.load_address = self._symbol_table[func_index]['symbol_name_addr'] - \
                                            self._string_table[str_index]['address']
                        self._firmware_info["load_address"] = self.load_address
                        self.logger.info('load address is {:010x}'.format(self.load_address))
                        performance_data = "Analyze loading address takes {:.3f} seconds".format(
                            self.get_timer())
                        self._performance_status.append(performance_data)
                        self.logger.debug(performance_data)
                        return self.load_address
                else:
                    continue

        self.logger.error("We didn't find load address in this firmware, sorry!")
        performance_data = "Analyze loading address takes {:.3f} seconds".format(self.get_timer())
        self._performance_status.append(performance_data)
        self.logger.debug(performance_data)

    def _check_load_address(self, address):
        """

        :param address:
        :return:
        """
        if not self._has_symbol:
            return False
        if len(self._symbol_table) > default_check_count:
            self.logger.debug("Length of symbol table greater than default. Setting iteration count to default of {}.".format(default_check_count))
            count = default_check_count
        else:
            count = len(self._symbol_table)
        self.logger.debug("symbol_table length is {}".format(count))
        for i in range(count):
            offset = self._symbol_table[i]['symbol_name_addr'] - address
            if offset <= 0:
                return False
            # TODO: Need improve, currently use string point to check.
            string, str_start_address, str_end_address = self._get_next_string_data(offset)
            if str_start_address != offset:
                self.logger.info("String {} at offset {} didn't match symbol table.".format(string, offset))
                return False
        self.logger.info('Load address is {:010x}'.format(address))
        return True

    def quick_test(self):
        """ Using known load address list to match VxWorks image.

        :return: Load address if match known address, None otherwise.
        """
        if self._has_symbol is False:
            return None
        self.logger.debug("has_symbol: {}".format(self._has_symbol))
        for address in known_address:
            if self._check_load_address(address):
                self.load_address = address
                self._firmware_info["load_address"] = self.load_address
                return self.load_address
            else:
                self.logger.info('Load address is not {:010x}'.format(address))

    def cleanup(self):
        """ Clean up variables.

        :return:
        """
        self.is_big_endian = False
        self.symbol_table_start = None
        self.symbol_table_end = None
        self._string_table = []
        self._symbol_table = []
        self.load_address = None
        self._has_symbol = None

    def get_string_from_firmware_by_offset(self, string_offset):
        symbol_name = ""
        while True:
            if self._firmware[string_offset] != 0x00:
                symbol_name += chr(self._firmware[string_offset])
                string_offset += 1

            else:
                break

        return symbol_name

    def get_symbols(self):
        self.symbols = []
        if self.load_address:
            for symbol in self._symbol_table:
                symbol_name_addr = symbol["symbol_name_addr"]
                symbol_dest_addr = symbol["symbol_dest_addr"]
                symbol_flag = symbol["symbol_flag"]
                symbol_name_firmware_addr = symbol_name_addr - self.load_address
                symbol_name = self.get_string_from_firmware_by_offset(symbol_name_firmware_addr)
                self.symbols.append({
                    "symbol_name": symbol_name,
                    "symbol_name_addr": symbol_name_addr,
                    "symbol_dest_addr": symbol_dest_addr,
                    "symbol_flag": symbol_flag
                })
            return self.symbols

        else:
            return None

    def get_firmware_info(self):
        symbols = self.get_symbols()
        self._firmware_info["symbols"] = symbols
        return self._firmware_info

    def get_performance_status(self):
        return self._performance_status

def demangle_function(demangle_string):
    function_return = None
    function_parameters = None
    function_name_end = len(demangle_string)

    # get parameters
    index = len(demangle_string) - 1
    if demangle_string[-1] == ')':
        # have parameters
        parentheses_count = 0
        while index >= 0:
            if demangle_string[index] == ')':
                parentheses_count += 1

            elif demangle_string[index] == '(':
                parentheses_count -= 1

            index -= 1

            if parentheses_count == 0:
                break

        function_parameters = demangle_string[index + 2:-1]
        function_name_end = index

    # get function name
    while index >= 0:
        if demangle_string[index] == ' ':
            break
        else:
            index -= 1
    function_name_start = index
    function_name = demangle_string[function_name_start + 1:function_name_end + 1]

    # get function return
    function_return = demangle_string[:function_name_start]
    return function_return, function_name, function_parameters


def is_vx_symbol_file(file_data, is_big_endian=True):
    # Check key function names
    for key_function in function_name_key_words:
        if key_function not in file_data:
            print("key function not found")
            return False

    if is_big_endian:
        return struct.unpack('>I', file_data[:4])[0] == len(file_data)

    else:
        return struct.unpack('<I', file_data[:4])[0] == len(file_data)


# --------------------------------------------------------------------------
# Plugin
# --------------------------------------------------------------------------
class AutoFixIDBForm(idaapi.Form):
    def __init__(self):
        self.invert = False
        self.vx_version = 5
        super(AutoFixIDBForm, self).__init__(
            r"""BUTTON YES* Start analyze

VxHunter Auto Fix IDB will fix IDB with VxWorks symbol table.

Please choose VxWorks main version
{FormChangeCb}
<VxWorks Main Version     :{c_vxversion}>
            """, {
                'FormChangeCb': self.FormChangeCb(self.OnFormChange),
                'c_vxversion': self.DropdownListControl(
                    items=("5", "6"),
                    readonly=True,
                    selval=0),

            }
        )

        self.Compile()

    def OnFormChange(self, fid):
        if fid == -2:
            self.vx_version = (5, 6)[self.GetControlValue(self.c_vxversion)]
            return 1


class FixCodeForm(idaapi.Form):
    def __init__(self):
        self.invert = False
        self.start_address = 0
        self.end_address = 0
        super(FixCodeForm, self).__init__(
            r"""BUTTON YES* Fix Code from image

VxHunter fix code will make all data as code from start_address to end_address. 

Please input start address and end address
{FormChangeCb}
<Start address     :{c_StartAddress}>
<End address     :{c_EndAddress}>
            """, {
                'FormChangeCb': self.FormChangeCb(self.OnFormChange),
                'c_StartAddress': self.NumericInput(value=self.start_address, swidth=40, tp=self.FT_ADDR),
                'c_EndAddress': self.NumericInput(value=self.end_address, swidth=40, tp=self.FT_ADDR),
            }
        )
        self.Compile()

    def OnFormChange(self, fid):
        if fid == -2:
            self.start_address = self.GetControlValue(self.c_StartAddress)
            self.end_address = self.GetControlValue(self.c_EndAddress)
            return 1


class FixAsciiForm(idaapi.Form):
    def __init__(self):
        self.invert = False
        self.string_address = 0
        super(FixAsciiForm, self).__init__(
            r"""BUTTON YES* Fix Code from image

VxHunter fix ascii string. 

Please input string table start address.
{FormChangeCb}
<String Address     :{c_Address}>
            """, {
                'FormChangeCb': self.FormChangeCb(self.OnFormChange),
                'c_Address': self.NumericInput(value=self.string_address, swidth=40, tp=self.FT_ADDR),
            }
        )
        self.Compile()

    def OnFormChange(self, fid):
        if fid == -2:
            self.string_address = self.GetControlValue(self.c_Address)
            return 1


class VxHunter_Plugin_t(idaapi.plugin_t):
    comment = "VxHunter plugin for IDA Pro"
    help = ""
    wanted_name = "VxHunter"
    wanted_hotkey = ""
    flags = idaapi.PLUGIN_KEEP

    def init(self):
        # register popup menu handlers
        try:
            # Register Auto Fix IDB handler
            VxHunterMCFixIDB.register(self, "Auto Fix IDB With symbol table")
            # Register Fix Code handler
            VxHunterMCFixCode.register(self, "Fix Code from start address to end address")
            # Register Fix Ascii handler
            VxHunterMCFixAscii.register(self, "Fix Ascii string table with giving address")
            # Register Load symbol file handler
            VxHunterMCLoadSymbolFile.register(self, "Load VxWorks symbol file")

        except Exception as err:
            print("Got Error!!!: %s" % err)

        # setup popup menu
        if idaapi.IDA_SDK_VERSION >= 700:
            # Add menu IDA >= 7.0
            idaapi.attach_action_to_menu("Edit/VxHunter/", VxHunterMCFixIDB.get_name(), idaapi.SETMENU_APP)
            idaapi.attach_action_to_menu("Edit/VxHunter/", VxHunterMCFixCode.get_name(), idaapi.SETMENU_APP)
            idaapi.attach_action_to_menu("Edit/VxHunter/", VxHunterMCFixAscii.get_name(), idaapi.SETMENU_APP)
            idaapi.attach_action_to_menu("Edit/VxHunter/", VxHunterMCLoadSymbolFile.get_name(), idaapi.SETMENU_APP)
        else:
            # add Vxhunter menu
            menu = idaapi.add_menu_item("Edit/VxHunter/", "Auto Fix IDB1", "", 1, self.handler_auto_fix_idb, None)
            if menu is not None:
                pass

        print("=" * 80)
        return idaapi.PLUGIN_KEEP

    def term(self):
        pass

    # null handler
    def menu_null(self):
        pass

    def handler_auto_fix_idb(self):
        form = AutoFixIDBForm()
        ok = form.Execute()
        if ok == 1:
            vx_version = int(form.vx_version)
            print("vx_version:%s" % vx_version)
            firmware_path = idaapi.get_input_file_path()
            firmware = open(firmware_path, 'rb').read()
            target = VxTarget(firmware=firmware, vx_version=vx_version)
            # target.logger.setLevel(logging.DEBUG)
            target.quick_test()

            if target.load_address:
                print("Load Address is:%s" % target.load_address)
            else:
                target.find_loading_address()
                if target.load_address:
                    print("Load Address is:%s" % target.load_address)
            if not target.load_address:
                return
            symbol_table_start = target.symbol_table_start
            symbol_table_end = target.symbol_table_end
            load_address = target.load_address
            self.fix_vxworks_idb(load_address, vx_version, symbol_table_start, symbol_table_end)
        form.Free()

    def handler_fix_code(self):
        form = FixCodeForm()
        ok = form.Execute()
        if ok == 1:
            start_address = int(form.start_address)
            end_address = int(form.end_address)
            self.fix_code(start_address, end_address)
        form.Free()

    def handler_fix_ascii(self):
        form = FixAsciiForm()
        ok = form.Execute()
        if ok == 1:
            string_address = int(form.string_address)
            self.fix_ascii(string_address)

        form.Free()

    def handler_load_symbol_file(self):
        self.load_symbol_file()

    @staticmethod
    def fix_vxworks_idb(load_address, vx_version, symbol_table_start, symbol_table_end):
        current_image_base = idaapi.get_imagebase()
        symbol_interval = 16
        if vx_version == 6:
            symbol_interval = 20
        symbol_table_start += load_address
        symbol_table_end += load_address
        ea = symbol_table_start
        shift_address = load_address - current_image_base
        while shift_address >= 0x70000000:
            idaapi.rebase_program(0x70000000, 0x0008)
            shift_address -= 0x70000000
        idaapi.rebase_program(shift_address, 0x0008)
        while ea < symbol_table_end:
            # for VxWorks 6 unknown symbol format
            if idc.get_wide_byte(ea + symbol_table_end - 2) == 3:
                ea += symbol_interval
                continue
            offset = 4
            idc.create_strlit(idc.get_wide_dword(ea + offset), idc.BADADDR)
            sName = idc.get_strlit_contents(idc.get_wide_dword(ea + offset), -1, ida_nalt.STRTYPE_C)
            print("Found %s in symbol table" % sName)
            if sName:
                sName_dst = idc.get_wide_dword(ea + offset + 4)
                if vx_version == 6:
                    sName_type = idc.get_wide_byte(ea + 0x12)
                else:
                    sName_type = idc.get_wide_byte(ea + 0x0e)

                idc.set_name(sName_dst, sName.decode('utf-8'), idc.SN_CHECK)
                if sName_type in need_create_function:
                    print("Start fix Function %s at %s" % (sName, hex(sName_dst)))
                    idc.create_insn(sName_dst)  # might not need
                    ida_funcs.add_func(sName_dst, idc.BADADDR)
            ea += symbol_interval
        print("Fix function by symbol table finish.")
        print("Start IDA auto analysis, depending on the size of the firmware this might take a few minutes.")
        idaapi.auto_wait()

    @staticmethod
    def fix_code(start_address, end_address):
        # Todo: There might be some data in the range of codes.
        offset = start_address
        while offset <= end_address:
            offset = ida_bytes.next_addr(offset)
            flags = ida_bytes.get_full_flags(offset)
            if not ida_bytes.is_code(flags):
                # Todo: Check should use MakeCode or MakeFunction
                # idc.create_insn(offset)
                ida_funcs.add_func(offset)

    @staticmethod
    def get_prev_ascii_string_address(address):
        """

        :param address: must be current ascii string start address.
        :return:
        """
        prev_string_start_address = address
        # string table interval should less than 5 bytes.
        if idc.get_wide_dword(address - 5) == 0:
            return None
        else:
            prev_string_start_address -= 5
            # TODO: Need handle short string.
            while idaapi.get_byte(prev_string_start_address) != 0:
                prev_string_start_address -= 1
            return prev_string_start_address + 1

    @staticmethod
    def get_next_ascii_string_address(address):
        """

        :param address: must be current ascii string start address.
        :return:
        """
        next_string_start_address = address
        # find current string end address
        while idaapi.get_byte(next_string_start_address) != 0:
            next_string_start_address += 1

        # string table interval should less than 5 bytes.
        # TODO: need handle short string.
        if idc.get_wide_dword(next_string_start_address + 1) == 0:
            return None

        while idaapi.get_byte(next_string_start_address) == 0:
            next_string_start_address += 1

        return next_string_start_address

    def get_string_table_start_address(self, address):
        string_table_start_address = address
        # find current string start address
        while idaapi.get_byte(string_table_start_address - 1) != 0:
            string_table_start_address -= 1

        while self.get_prev_ascii_string_address(string_table_start_address):
            string_table_start_address = self.get_prev_ascii_string_address(string_table_start_address)

        return string_table_start_address

    def fix_ascii(self, address):
        string_table_start_address = self.get_string_table_start_address(address)
        string_address = string_table_start_address
        while True:
            if string_address:
                print("Start Make string at address: %s" % hex(string_address))
                idc.create_strlit(string_address, idc.BADADDR)
                string_address = self.get_next_ascii_string_address(string_address)
            else:
                break

    def load_symbols(self, file_data, is_big_endian=True):
        symbol_list = []
        if is_big_endian:
            unpack_format = '>I'
        else:
            unpack_format = '<I'

        symbol_count = struct.unpack(unpack_format, file_data[4:8])[0]
        print("symbol_count: %s" % symbol_count)
        symbol_offset = 8
        string_table_offset = 8 + 8 * symbol_count
        print("string_table_offset: %s" % string_table_offset)
        # get symbols
        for i in range(symbol_count):
            offset = i * 8
            symbol_data = file_data[symbol_offset + offset:symbol_offset + offset + 8]
            flag = ord(symbol_data[0])
            string_offset = struct.unpack(unpack_format, '\x00' + symbol_data[1:4])[0]
            string_offset += string_table_offset
            print("string_offset: %s" % string_offset)
            symbol_name = ""
            while True:
                if file_data[string_offset] != '\x00':
                    symbol_name += file_data[string_offset]
                    string_offset += 1

                else:
                    break
            print("symbol_name: %s" % symbol_name)
            symbol_address = struct.unpack(unpack_format, symbol_data[-4:])[0]
            symbol_list.append([flag, symbol_name, symbol_address])
            # Find TP-Link device loading address with symbols
            if "wrs_kernel_text_start" in symbol_name:
                load_address = symbol_address
                current_image_base = idaapi.get_imagebase()
                shift_address = load_address - current_image_base
                while shift_address >= 0x70000000:
                    idaapi.rebase_program(0x70000000, 0x0008)
                    shift_address -= 0x70000000
                idaapi.rebase_program(shift_address, 0x0008)

        # load symbols
        for symbol_data in symbol_list:
            flag, symbol_name, symbol_address = symbol_data
            idc.set_name(symbol_address, symbol_name.decode('utf-8'), idc.SN_CHECK)
            if flag == 0x54:
                if symbol_name:
                    print("Start fix Function %s at %s" % (symbol_name, hex(symbol_address)))
                    idc.create_insn(symbol_address)  # might not need
                    ida_funcs.add_func(symbol_address, idc.BADADDR)

    def load_symbol_file(self):
        symbol_file_path = ida_kernwin.ask_file(0, "*", "Please chose the VxWorks symbol file")
        print("symbol_file_path: {}".format(symbol_file_path))
        symbol_file_data = open(symbol_file_path, 'rb').read()
        if is_vx_symbol_file(symbol_file_data):
            self.load_symbols(symbol_file_data)
            idaapi.auto_wait()

        else:
            return

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


try:
    class VxHunterMenuContext(idaapi.action_handler_t):

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

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

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

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

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

        @classmethod
        def update(self, ctx):
            try:
                if ctx.form_type == idaapi.BWN_DISASM:
                    return idaapi.AST_ENABLE_FOR_FORM
                else:
                    return idaapi.AST_DISABLE_FOR_FORM
            except:
                # Add exception for main menu on >= IDA 7.0
                return idaapi.AST_ENABLE_ALWAYS

    # context menu for Fix idb
    class VxHunterMCFixIDB(VxHunterMenuContext):
        def activate(self, ctx):
            self.plugin.handler_auto_fix_idb()
            return 1

    class VxHunterMCFixCode(VxHunterMenuContext):
        def activate(self, ctx):
            self.plugin.handler_fix_code()
            return 1

    class VxHunterMCFixAscii(VxHunterMenuContext):
        def activate(self, ctx):
            self.plugin.handler_fix_ascii()
            return 1

    class VxHunterMCLoadSymbolFile(VxHunterMenuContext):
        def activate(self, ctx):
            self.plugin.handler_load_symbol_file()
            return 1


except Exception as err:
    # TODO: Add some handle later.
    print("err: {}".format(err))


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


repository: NeatMonster/AMIE
file: amie.py
description: A Minimalist Instruction Extender for the ARM architecture and IDA Pro
url: https://github.com/NeatMonster/AMIE/blob/8ca946a3b6ab0193d4f4ad2092f58b25f23a8210/amie.py
wanted_name: AMIE
comment: A Minimalist Instruction Extender
created_at: 2018-11-14 09:39:32+00:00
pushed_at: 2024-08-16 15:49:13+00:00
content:
import inspect
import json
import os
import re
import struct
import sys

import ida_allins
import ida_bytes
import ida_hexrays
import ida_idaapi
import ida_idp
import ida_kernwin
import ida_lines
import ida_segregs
import ida_ua


class Arch(object):
    @staticmethod
    def output_insn(outctx, mne, dst_reg, src_reg=None):
        outctx.out_custom_mnem(mne.upper(), 16)
        outctx.out_register(dst_reg)
        if src_reg:
            outctx.out_printf(", ")
            outctx.out_register(src_reg)
        outctx.flush_outbuf()
        return True

    def __init__(self, file_name):
        cur_file = inspect.getsourcefile(lambda: 0)
        cur_path = os.path.dirname(os.path.abspath(cur_file))

        def unicode_to_str(data, ignore_dicts=False):
            if isinstance(data, unicode):
                return data.encode("utf-8")
            if isinstance(data, list):
                return [unicode_to_str(item, True) for item in data]
            if isinstance(data, dict) and not ignore_dicts:
                return {unicode_to_str(key): unicode_to_str(value, True)
                        for key, value in data.iteritems()}
            return data

        with open(os.path.join(cur_path, file_name), "r") as fd:
            hook = unicode_to_str if sys.version_info[0] < 3 else None
            self.data = json.loads(fd.read(), object_hook=hook)

        self.regs, self.regs_enc = {}, {}
        for regs_group, regs in self.data["registers"].items():
            if regs_group == "encodings":
                self.regs_enc = regs
            else:
                self.regs.update(regs)

        self.insns, self.insns_enc = {}, {}
        for insns_group, insns in self.data["instructions"].items():
            if insns_group == "encodings":
                self.insns_enc = insns
            else:
                self.insns.update(insns)

    def matches(self, op, reg_op):
        op_str = "{:0{size}b}".format(op, size=len(reg_op))[: len(reg_op)]
        for i in range(len(op_str)):
            if op_str[i] != reg_op[i] and reg_op[i] != "x":
                return False
        return True

    def find_reg_enc(self, enc_type, ops):
        for cp_reg_name, reg_ops in self.regs_enc[enc_type].items():
            for i in range(len(ops)):
                if not self.matches(ops[i], reg_ops[i]):
                    break
            else:
                return cp_reg_name
        print("[AMIE] %s encoding not found for: %s" % (enc_type, ops))

    def find_insn_enc(self, enc_type, insn_bs, insn_mne):
        for insn_name, insn_encs in self.insns_enc[enc_type].items():
            for tmpl_name, insn_enc in insn_encs.items():
                match = self.matches(insn_bs, insn_enc[0])
                for enc in insn_enc[1:]:
                    match = match and not self.matches(insn_bs, enc)
                if not match:
                    continue

                for tmpls in self.insns[insn_name]["templates"].values():
                    for tmpl in tmpls:
                        mne = tmpl.split(" ")[0]
                        mne = re.sub("{[^}]}", "", mne)
                        mne = re.sub("<[^>]>", "", mne)
                        if mne.startswith(insn_mne):
                            return insn_name, tmpl_name
        print("[AMIE] %s encoding not found for: %s" % (enc_type, insn_mne))

    def output(self, outctx):
        raise False

    def hint(self, ea, tag, val):
        val = val.strip()
        if not val:
            return None

        if tag == ida_lines.SCOLOR_REG:
            if val.split(".")[0] in self.regs:
                register = self.regs[val.split(".")[0]]
                return register["long_name"], register["purpose"]

        elif tag == ida_lines.SCOLOR_INSN:
            insn = ida_ua.insn_t()
            ida_ua.decode_insn(insn, ea)
            insn_bs = ida_bytes.get_bytes(insn.ea, insn.size)
            fmt = {2: "<H", 4: "<I", 8: "<Q"}
            insn_bs = struct.unpack(fmt.get(len(insn_bs)), insn_bs)[0]

            enc_type = "A64"
            if isinstance(self, AArch32):
                enc_type = "A32"
                if ida_segregs.get_sreg(ea, 20) > 0:
                    enc_type = "T32"

            insn_enc = enc_type, insn_bs, insn.get_canon_mnem()
            insn_enc = self.find_insn_enc(*insn_enc)
            if not insn_enc:
                return
            insn_name, tmpl_name = insn_enc
            insn = self.insns[insn_name]

            desc = insn["authored"]
            if tmpl_name and tmpl_name in insn["templates"]:
                desc += "\n\n" + "\n".join(insn["templates"][tmpl_name])
            return insn["heading"], desc

        elif tag == ida_lines.SCOLOR_KEYWORD:
            if val in self.data["keywords"]:
                return val, self.data["keywords"][val]


class AArch32(Arch):
    def __init__(self):
        super(AArch32, self).__init__("aarch32.json")

    def decode_mcr_mrc(self, insn):
        cp = insn.ops[0].specflag1
        op1 = insn.ops[0].value
        rt = insn.ops[1].reg
        crn = insn.ops[1].specflag1
        crm = insn.ops[1].specflag2
        op2 = insn.ops[2].value

        ops = [cp, op1, crn, crm, op2]
        gp_reg = ida_idp.ph.regnames[rt]
        cp_reg = self.find_reg_enc("MCR|MRC", ops)
        return cp_reg, gp_reg

    def output(self, outctx):
        insn = outctx.insn
        mne = ida_ua.print_insn_mnem(insn.ea)

        if insn.itype in [ida_allins.ARM_mcr, ida_allins.ARM_mrc]:
            cp_reg, gp_reg = self.decode_mcr_mrc(insn)
            if cp_reg:
                if insn.itype == ida_allins.ARM_mcr:
                    return Arch.output_insn(outctx, mne, cp_reg, gp_reg)
                else:
                    return Arch.output_insn(outctx, mne, gp_reg, cp_reg)

        if insn.itype in [ida_allins.ARM_mcrr, ida_allins.ARM_mrrc]:
            cp = insn.ops[0].specflag1
            op3 = insn.ops[0].value
            rt = insn.ops[1].reg
            rt2 = insn.ops[1].specflag1
            crm = insn.ops[1].specflag2

            ops = [cp, op3, crm]
            gp_reg = "%s%s%s:%s%s%s" % (
                ida_idp.ph.regnames[rt2],
                ida_lines.SCOLOR_OFF,
                ida_lines.SCOLOR_REG,
                ida_lines.SCOLOR_ON,
                ida_lines.SCOLOR_REG,
                ida_idp.ph.regnames[rt],
            )
            cp_reg = self.find_reg_enc("MCRR|MRRC", ops)
            if cp_reg:
                if insn.itype == ida_allins.ARM_mcrr:
                    return Arch.output_insn(outctx, mne, cp_reg, gp_reg)
                else:
                    return Arch.output_insn(outctx, mne, gp_reg, cp_reg)

        return False


class AArch64(Arch):
    def __init__(self):
        super(AArch64, self).__init__("aarch64.json")

    def decode_msr_mrs(self, insn):
        x = ida_bytes.get_wide_dword(insn.ea)
        op0 = 2 + ((x >> 19) & 1)
        if insn.itype == ida_allins.ARM_msr:
            i, xt = 0, insn.ops[4].reg
        else:
            xt, i = insn.ops[0].reg, 1
        op1 = insn.ops[i].value
        crn = insn.ops[i + 1].reg
        crm = insn.ops[i + 2].reg
        op2 = insn.ops[i + 3].value

        ops = [op0, op1, crn, crm, op2]
        gp_reg = ida_idp.ph.regnames[xt]
        cp_reg = self.find_reg_enc("MSR|MRS", ops)
        return cp_reg, gp_reg

    def output(self, outctx):
        insn = outctx.insn
        mne = ida_ua.print_insn_mnem(insn.ea)

        if insn.itype in [ida_allins.ARM_msr, ida_allins.ARM_mrs]:
            if insn.ops[2].type > 0:
                cp_reg, gp_reg = self.decode_msr_mrs(insn)
                if cp_reg:
                    if insn.itype == ida_allins.ARM_msr:
                        return Arch.output_insn(outctx, mne, cp_reg, gp_reg)
                    else:
                        return Arch.output_insn(outctx, mne, gp_reg, cp_reg)

            else:
                fields = {5: "SPSel", 6: "DAIFSet", 7: "DAIFClr"}
                pstatefield = fields.get(insn.ops[0].value, None)
                if pstatefield:
                    outctx.out_custom_mnem(mne.upper(), 16)
                    outctx.out_register(pstatefield)
                    outctx.out_printf(", ")
                    imm = ida_ua.print_operand(insn.ea, 1)
                    outctx.out_printf(imm)
                    outctx.flush_outbuf()
                    return True

        if insn.itype == ida_allins.ARM_sys:
            op1 = insn.ops[0].value
            crn = insn.ops[1].reg
            crm = insn.ops[2].reg
            op2 = insn.ops[3].value
            xt = insn.ops[4].reg

            ops = [1, op1, crn, crm, op2]
            gp_reg = ida_idp.ph.regnames[xt]
            cp_reg = self.find_reg_enc("SYS", ops)
            if cp_reg:
                mne, op = cp_reg.split()
                if xt == 160:
                    return Arch.output_insn(outctx, mne, op)
                else:
                    return Arch.output_insn(outctx, mne, op, gp_reg)

        return False


class AMIE(ida_idaapi.plugin_t, ida_idp.IDP_Hooks, ida_kernwin.UI_Hooks):
    flags = ida_idaapi.PLUGIN_PROC | ida_idaapi.PLUGIN_DRAW
    comment = "A Minimalist Instruction Extender"
    help = ""
    wanted_name = "AMIE"
    wanted_hotkey = ""
    version = "0.8.6"

    @staticmethod
    def extract_item(viewer, stripped=False):
        cnt = ida_kernwin.get_custom_viewer_place(viewer, True)[1]
        line = ida_kernwin.get_custom_viewer_curline(viewer, True)

        tags = [ida_lines.SCOLOR_ON, ida_lines.SCOLOR_OFF]
        expr = re.compile("[A-Z0-9_]")

        def find_pos(line, pos, inc):
            while 0 <= pos < len(line):
                if not stripped and line[pos] in tags:
                    break
                if stripped and not expr.match(line[pos]):
                    break
                pos += inc
            return pos

        pos = ida_lines.tag_advance(line, cnt)
        if pos < 0 or pos >= len(line):
            return
        while line[pos] in [ida_lines.SCOLOR_ON, ida_lines.SCOLOR_OFF]:
            pos += 2
            if pos < 0 or pos >= len(line):
                return

        prev_pos, next_pos = find_pos(line, pos, -1), find_pos(line, pos, +1)
        if stripped:
            return None, line[prev_pos + 1:next_pos].strip()
        return line[prev_pos + 1], line[prev_pos + 2:next_pos].strip()

    @staticmethod
    def format_hint(title, content):
        def tag_append(s, t):
            return ida_lines.SCOLOR_ON + t + s + ida_lines.SCOLOR_OFF + t

        text = " " + "\n ".join([tag_append(line, ida_lines.SCOLOR_LOCNAME)
                                 for line in title.split("\n")])
        if content:
            text += "\n"
            for line in content.split("\n"):
                text += "\n " + tag_append(line, ida_lines.SCOLOR_AUTOCMT)
        return text, len(text.split("\n"))

    def __init__(self):
        ida_idp.IDP_Hooks.__init__(self)
        ida_kernwin.UI_Hooks.__init__(self)
        self.hexrays_support = False

    def init(self):
        info = ida_idaapi.get_inf_structure()
        if info.procName != "ARM":
            return ida_idaapi.PLUGIN_SKIP

        ida_kernwin.UI_Hooks.hook(self)
        ida_idp.IDP_Hooks.hook(self)
        print("[AMIE] Plugin v%s initialized" % AMIE.version)
        return ida_idaapi.PLUGIN_KEEP

    def term(self):
        ida_idp.IDP_Hooks.unhook(self)
        ida_kernwin.UI_Hooks.unhook(self)
        if self.hexrays_support:
            ida_hexrays.remove_hexrays_callback(self.hxe_callback)
            ida_hexrays.term_hexrays_plugin()
        print("[AMIE] Plugin v%s terminated" % AMIE.version)

    def run(self, _):
        return False

    def database_inited(self, *_):
        info = ida_idaapi.get_inf_structure()
        if info.is_64bit():
            print("[AMIE] Using AArch64 architecture")
            self.arch = AArch64()
        else:
            print("[AMIE] Using AArch32 architecture")
            self.arch = AArch32()

    def plugin_loaded(self, plugin_info):
        if plugin_info.name == "Hex-Rays Decompiler":
            if ida_hexrays.init_hexrays_plugin():
                self.hexrays_support = True
                ida_hexrays.install_hexrays_callback(self.hxe_callback)
                print("[AMIE] Hex-Rays decompiler is supported")

    def ev_out_insn(self, outctx):
        return self.arch.output(outctx)

    def get_custom_viewer_hint(self, viewer, place):
        widget_type = ida_kernwin.get_widget_type(viewer)
        if widget_type != ida_kernwin.BWN_DISASM:
            return

        flags = ida_bytes.get_flags(place.toea()) if place else 0
        if not ida_bytes.is_code(flags):
            return

        item = AMIE.extract_item(viewer)
        if not item:
            return
        tag, val = item

        hint = self.arch.hint(place.toea(), tag, val)
        if not hint:
            return

        return AMIE.format_hint(*hint)

    def hxe_callback(self, event, *args):
        if event == ida_hexrays.hxe_maturity:
            cfunc, new_maturity = args
            if new_maturity != ida_hexrays.CMAT_FINAL:
                return 0
            plugin = self

            class Visitor(ida_hexrays.ctree_visitor_t):
                def visit_expr(self, e):
                    if not e.x or e.x.op != ida_hexrays.cot_helper:
                        return 0

                    insn = ida_ua.insn_t()
                    ida_ua.decode_insn(insn, e.ea)

                    def make_reg(cp_reg):
                        reg = ida_hexrays.carg_t()
                        reg.op = ida_hexrays.cot_helper
                        reg.helper = cp_reg
                        reg.exflags = ida_hexrays.EXFL_ALONE
                        return reg

                    if e.x.helper in ["__mcr", "__mrc"]:
                        cp_reg = plugin.arch.decode_mcr_mrc(insn)[0]
                        if cp_reg:
                            if e.x.helper == "__mcr":
                                e.x.helper = "_WriteStatusReg"
                                val = ida_hexrays.carg_t()
                                e.a[2].swap(val)
                                e.a.clear()
                                e.a.push_back(make_reg(cp_reg))
                                e.a.push_back(val)

                            else:
                                e.x.helper = "_ReadStatusReg"
                                e.a.clear()
                                e.a.push_back(make_reg(cp_reg))

                    elif e.x.helper == "ARM64_SYSREG":
                        cp_reg = plugin.arch.decode_msr_mrs(insn)[0]
                        if cp_reg:
                            e.replace_by(make_reg(cp_reg))

                    return 0

            visitor = Visitor(ida_hexrays.CV_PARENTS)
            visitor.apply_to_exprs(cfunc.body, None)

        elif event == ida_hexrays.hxe_create_hint:
            hint = None

            e = args[0].item.e
            if e and e.op == ida_hexrays.cot_helper:
                hint = self.arch.hint(e.ea, ida_lines.SCOLOR_REG, e.helper)

            elif e and e.op == ida_hexrays.cit_asm:
                item = AMIE.extract_item(args[0].ct, True)
                if item and item[1]:
                    if item[1] in self.arch.regs:
                        tag = ida_lines.SCOLOR_REG
                    else:
                        insn = ida_ua.insn_t()
                        ida_ua.decode_insn(insn, e.ea)
                        if item[1].startswith(insn.get_canon_mnem()):
                            tag = ida_lines.SCOLOR_INSN
                        else:
                            tag = ida_lines.SCOLOR_KEYWORD
                    hint = self.arch.hint(e.ea, tag, item[1])

            if hint:
                return (1,) + AMIE.format_hint(*hint)

        return 0


def PLUGIN_ENTRY():
    return AMIE()


repository: gmh5225/myda
file: main_myda.py
description: General Purpose IDA Plugin
url: https://github.com/gmh5225/myda/blob/dc3d56d9f923bf313a61add2dfaa7ed5485e5537/main_myda.py
wanted_name: myda
created_at: 2023-11-23 10:13:26+00:00
pushed_at: 2023-11-20 15:42:47+00:00
content:
import idaapi
import os
from myda.vms import *

def add_menu():

    # create custom menus - top level and sub menu
    idaapi.create_menu("myda_menu", "myda", "")
    idaapi.create_menu("vm_submenu", "VMs", "myda/")

    ACTION_NAME_1 = "action1"
    
    d1 = idaapi.action_desc_t(
            ACTION_NAME_1,
            "Run the VMs script",
            VMs())

    idaapi.register_action(d1)

    idaapi.attach_action_to_menu(
            "myda/VMs/",
            ACTION_NAME_1,
            idaapi.SETMENU_INS)

class myda_plugmod_t(idaapi.plugmod_t):
    def run(self, arg):
        # on init, create the menu
        add_menu()
        return 0

class myda_plugin_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_MULTI
    help = "myda help"
    wanted_name = "myda"
    wanted_hotkey = ""

    def init(self):
        return myda_plugmod_t()

def PLUGIN_ENTRY():
    return myda_plugin_t()


repository: Donaldduck8/IDA-Scripts-and-Tools
file: better_annotator.py
description: A collection of plug-ins I've developed for IDA Pro.
wanted_name: BetterAnnotator
comment: BetterAnnotator
created_at: 2023-08-02 10:06:40+00:00
pushed_at: 2024-04-09 16:47:04+00:00
url:
https://github.com/Donaldduck8/IDA-Scripts-and-Tools/blob/4ec79fdcd2a2a43348282ee427944ed8a7593a57/better_annotator.py

content:
import idaapi
import ida_kernwin
import donald_ida_utils

import re
import json

from PyQt5 import QtGui, QtWidgets
from PyQt5.QtCore import Qt


SYMBOLS_TO_WORDS = {
    # ' ': '_',
    '!': 'exclamation_mark',
    '"': 'double_quote',
    '#': 'hash',
    '$': 'dollar',
    '%': 'percent',
    '&': 'ampersand',
    # "'": 'quote',
    '(': 'left_paren',
    ')': 'right_paren',
    '*': 'star',
    '+': 'plus',
    ',': 'comma',
    '-': 'dash',
    '.': 'dot',
    '/': 'slash',
    ':': 'colon',
    ';': 'semicolon',
    '<': 'less_than',
    '=': 'equals',
    '>': 'greater_than',
    '?': 'question_mark',
    '@': 'at',
    '[': 'left_bracket',
    '\\': 'backslash',
    ']': 'right_bracket',
    '^': 'caret',
    # '_': 'underscore',
    '`': 'backtick',
    '{': 'left_brace',
    '|': 'vertical_bar',
    '}': 'right_brace',
    '~': 'tilde'
}


class IDAStringFormatter:
    def __init__(
        self,
        data: dict,
        mode: str,
        maximum_string_length: int = 35,
        keep_original_capitalization: bool = False,
        transliterate_symbols: bool = False,
        comment_prefix: str = "[AUTO]",
    ):
        self.data = data
        self.mode = mode
        self.maximum_string_length = maximum_string_length
        self.keep_original_capitalization = keep_original_capitalization
        self.transliterate_symbols = transliterate_symbols
        self.comment_prefix = comment_prefix

    def transliterate_string(self, input_string: str) -> str:
        # Replace all non-alphanumeric characters with their word equivalent
        for symbol, word in SYMBOLS_TO_WORDS.items():
            input_string = input_string.replace(symbol, f"_{word}_")

        return input_string

    def sanitize_string(self, input_string: str) -> str:
        # Replace all non-alphanumeric characters with an underscore
        input_string = re.sub(r"[^\w]+", "_", input_string)

        # Remove leading, trailing, and multiple consecutive underscores
        input_string = re.sub(r"_+", "_", input_string).strip("_")

        return input_string

    def fix_duplicate_values(self, data: dict) -> dict:
        seen_values = []
        new_data = {}

        for key, value in data.items():
            if value in seen_values:
                duplicate_index = 2

                value_new = value + f"_x{duplicate_index}"

                while value_new in seen_values:
                    duplicate_index += 1
                    value_new = value + f"_x{duplicate_index}"

                value = value_new

            seen_values.append(value)
            new_data[key] = value

        return new_data

    def format_data(self):
        if self.mode == "enum":
            return self.format_enum()
        elif self.mode == "globals":
            return self.format_globals()
        elif self.mode == "comments":
            return self.format_comments()
        else:
            return None

    def format_enum(self):
        return self.format_tokens(enum=True)

    def format_globals(self):
        return self.format_tokens(enum=False)

    def format_tokens(self, enum: bool = False):
        data = self.data

        # Transliteration
        if self.transliterate_symbols:
            data = {
                k: self.transliterate_string(v) for k, v in data.items()
            }

        # Sanitize all strings
        data = {
            k: self.sanitize_string(v) for k, v in data.items()
        }

        # Fix all duplicate values
        data = self.fix_duplicate_values(data)

        # Truncate all strings to the maximum length
        data = {
            k: v[:self.maximum_string_length] for k, v in data.items()
        }

        # Convert all strings to lower-case
        if not self.keep_original_capitalization:
            if enum:
                data = {
                    k: v.lower() for k, v in data.items()
                }
            else:
                data = {
                    k: v.upper() for k, v in data.items()
                }
                
        # If string starts with a number, add "str_" prefix
        data = {
            k: f"str_{v}" if v[0].isdigit() else v for k, v in data.items()
        }

        return data

    def format_comments(self):
        data = self.data

        # Apply the comment prefix to all values:
        if self.comment_prefix is not None and self.comment_prefix != "":
            data = {
                k: f"{self.comment_prefix} {v}" for k, v in data.items()
            }

        return data

    def format_enum_string(self):
        data = self.format_enum()

        enum_s = "enum enum_name {\n"
        for key, value in data.items():
            enum_s += f"str_{value} = {hex(key)},\n"

        enum_s += "};"

        return enum_s


class BetterAnnotatorDialog(QtWidgets.QDialog):
    def __init__(self, parent):
        super(BetterAnnotatorDialog, self).__init__(parent)
        self.user_input = """{
    "0x18000154": "This is the decrypted string!"
}"""
        self.mode = "comments"
        self.log_widget = None
        self.data = None
        self.max_string_length = 35
        self.keep_original_capitalization = False
        self.transliterate_symbols = False
        self.comment_prefix = "[AUTO]"
        self.populate_form()

    def populate_form(self):
        self.setWindowTitle("Better Annotator")
        self.resize(800, 600)
        self.layout = QtWidgets.QVBoxLayout(self)
        self.top_layout = QtWidgets.QHBoxLayout()
        self.bottom_layout = QtWidgets.QHBoxLayout()
        self.bottom_layout.setAlignment(Qt.AlignRight | Qt.AlignBottom)

        self.log_widget = QtWidgets.QLabel("Paste your generated JSON data here!")
        self.layout.addWidget(self.log_widget)
        self.text_edit = QtWidgets.QTextEdit()
        font = QtGui.QFont()
        font.setFamily("Consolas")
        font.setStyleHint(QtGui.QFont.Monospace)
        font.setFixedPitch(True)
        font.setPointSize(10)
        self.text_edit.setFont(font)
        metrics = QtGui.QFontMetrics(font)
        self.text_edit.setTabStopWidth(4 * metrics.width(" "))
        self.text_edit.insertPlainText(self.user_input)
        self.layout.addWidget(self.text_edit)

        # Connect the textChanged signal to the on_text_changed method
        self.text_edit.textChanged.connect(self.on_text_changed)

        # Create the table and set its properties
        self.table_widget = QtWidgets.QTableWidget()
        self.table_widget.setColumnCount(2)  # Two columns
        self.table_widget.setHorizontalHeaderLabels(["Key", "Value"])
        self.table_widget.horizontalHeader().setStretchLastSection(True)
        self.layout.addWidget(self.table_widget)

        # Radio buttons group
        self.radio_group_box = QtWidgets.QGroupBox("Mode")
        self.radio_layout = QtWidgets.QHBoxLayout()
        self.radio_globals = QtWidgets.QRadioButton("Globals")
        self.radio_enum = QtWidgets.QRadioButton("Enum")
        self.radio_comments = QtWidgets.QRadioButton("Comments")
        self.radio_layout.addWidget(self.radio_globals)
        self.radio_layout.addWidget(self.radio_enum)
        self.radio_layout.addWidget(self.radio_comments)
        self.radio_group_box.setLayout(self.radio_layout)
        self.layout.addWidget(self.radio_group_box)

        # Connect each radio button to a callback
        self.radio_globals.clicked.connect(self.radio_button_clicked)
        self.radio_enum.clicked.connect(self.radio_button_clicked)
        self.radio_comments.clicked.connect(self.radio_button_clicked)

        # Set "Comments" radio button checked by default
        self.radio_comments.setChecked(True)

        # Options group
        self.options_group_box = QtWidgets.QGroupBox("Options")
        self.options_layout = QtWidgets.QHBoxLayout()

        # Max string length
        self.max_length_label = QtWidgets.QLabel("Maximum string length:")
        self.max_length_spinbox = QtWidgets.QSpinBox()
        self.max_length_spinbox.setMinimum(1)
        self.max_length_spinbox.setMaximum(10000)
        self.max_length_spinbox.setValue(35)  # Default value
        self.max_length_spinbox.valueChanged.connect(self.on_max_string_length_changed)

        self.options_layout.addWidget(self.max_length_label)
        self.options_layout.addWidget(self.max_length_spinbox)

        # Checkbox for keeping original capitalization
        self.keep_capitalization_checkbox = QtWidgets.QCheckBox("Keep original capitalization")
        self.keep_capitalization_checkbox.stateChanged.connect(self.on_keep_capitalization_changed)
        self.options_layout.addWidget(self.keep_capitalization_checkbox)

        # Checkbox for transliterating symbols
        self.transliterate_symbols_checkbox = QtWidgets.QCheckBox("Transliterate symbols")
        self.transliterate_symbols_checkbox.stateChanged.connect(self.on_transliterate_symbols_changed)
        self.options_layout.addWidget(self.transliterate_symbols_checkbox)

        # Comment prefix text input
        self.comment_prefix_label = QtWidgets.QLabel("Comment prefix:")
        self.comment_prefix_lineedit = QtWidgets.QLineEdit()
        self.comment_prefix_lineedit.setText("[AUTO]")
        self.comment_prefix_lineedit.textChanged.connect(self.on_comment_prefix_changed)
        self.options_layout.addWidget(self.comment_prefix_label)
        self.options_layout.addWidget(self.comment_prefix_lineedit)

        self.options_group_box.setLayout(self.options_layout)
        self.layout.addWidget(self.options_group_box)

        self.ok_btn = QtWidgets.QPushButton("OK")
        self.ok_btn.setFixedWidth(100)
        self.ok_btn.clicked.connect(self.ok_btn_clicked)
        self.bottom_layout.addWidget(self.ok_btn)

        self.layout.addLayout(self.top_layout)
        self.layout.addLayout(self.bottom_layout)

        self.set_table()

    def on_max_string_length_changed(self):
        self.max_string_length = self.max_length_spinbox.value()

        self.set_table()

    def on_keep_capitalization_changed(self):
        self.keep_original_capitalization = self.keep_capitalization_checkbox.isChecked()

        self.set_table()

    def on_transliterate_symbols_changed(self):
        self.transliterate_symbols = self.transliterate_symbols_checkbox.isChecked()

        self.set_table()

    def on_comment_prefix_changed(self):
        self.comment_prefix = self.comment_prefix_lineedit.text()

        self.set_table()

    def on_text_changed(self):
        # This method will be called whenever the text in the QTextEdit changes.
        self.set_table()

    def set_table(self):
        text = self.text_edit.toPlainText()  # Get current text from QTextEdit

        self.data = None
        self.table_widget.setRowCount(0)  # Clear existing rows

        if len(text) == 0:
            self.log_widget.setText("Paste your generated JSON data here!")

        try:
            # Attempt to parse the text as JSON
            data = json.loads(text)
            if not isinstance(data, dict):
                self.log_widget.setText("Data is not a dictionary!")
                return

            data_normalized = {}

            for key, value in data.items():
                # Ensure every key is a number and every value is a string
                if (
                    isinstance(key, int)
                    or (isinstance(key, str) and key.isdigit())
                    or key.startswith("0x")
                ):
                    if isinstance(value, str):
                        if key.startswith("0x"):
                            data_normalized[int(key, 16)] = value
                        else:
                            data_normalized[int(key)] = value
                    else:
                        self.log_widget.setText(f"Value for key {key} is not a string!")
                        return
                else:
                    self.log_widget.setText(f"Key {key} is not a number!")
                    return

            data = data_normalized

            self.formatter = IDAStringFormatter(
                data,
                self.mode,
                self.max_string_length,
                self.keep_original_capitalization,
                self.transliterate_symbols,
                self.comment_prefix,
            )

            data = self.formatter.format_data()

            for key, value in data.items():
                row_position = self.table_widget.rowCount()
                self.table_widget.insertRow(row_position)
                self.table_widget.setItem(
                    row_position, 0, QtWidgets.QTableWidgetItem(hex(key))
                )
                self.table_widget.setItem(
                    row_position, 1, QtWidgets.QTableWidgetItem(value)
                )

            self.log_widget.setText("Parsed successfully!")

            self.data = data

        except json.JSONDecodeError as e:
            self.log_widget.setText(f"Failed to parse JSON, {e.msg}")
            self.table_widget.setRowCount(0)

    def radio_button_clicked(self):
        if self.radio_globals.isChecked():
            self.mode = "globals"
        elif self.radio_enum.isChecked():
            self.mode = "enum"
        elif self.radio_comments.isChecked():
            self.mode = "comments"

        self.set_table()

    def ok_btn_clicked(self):
        if self.data is None:
            self.close()
            return

        if self.mode == "comments":
            for addr, value in self.data.items():
                donald_ida_utils.add_pseudocode_comment(addr, value, prefix=None)

        elif self.mode == "globals":
            for addr, value in self.data.items():
                donald_ida_utils.define_and_rename_global(addr, value)
                donald_ida_utils.add_disassembly_comment(addr, value)

        elif self.mode == "enum":
            enum_s = self.formatter.format_enum_string()

            print(enum_s)
        else:
            return

        self.close()


def show_text_input_dialog():
    f = TextInputForm()
    f.Compile()
    ok = f.Execute()
    if ok == 1:
        return f.inp_str.value
    f.Free()
    return None


class TextInputForm(ida_kernwin.Form):
    def __init__(self):
        self.inp_str = ida_kernwin.Form.StringInput()
        form_str = "STARTITEM 0\nEnum Name\n\n  <##Enter the desired name:{inp_str}>"
        ida_kernwin.Form.__init__(self, form_str, {"inp_str": self.inp_str})


class BetterAnnotatorPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_FIX
    comment = "BetterAnnotator"
    help = "BetterAnnotator"
    wanted_name = "BetterAnnotator"
    wanted_hotkey = "Ctrl+Shift+A"
    dialog = None

    def init(self):
        print("BetterAnnotator :: Plugin Started")
        return idaapi.PLUGIN_KEEP

    def term(self):
        pass

    def run(self, arg):
        self.dialog = BetterAnnotatorDialog(None)
        self.dialog.show()


def generic_handler(callback):
    class Handler(idaapi.action_handler_t):
        def __init__(self):
            idaapi.action_handler_t.__init__(self)

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

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

    return Handler()


plugin = BetterAnnotatorPlugin()


def PLUGIN_ENTRY():
    global plugin
    return plugin


repository: jonpalmisc/ida_screenshot
file: screenshot.py
description: High-resolution screenshot capture plugin for IDA Pro
wanted_name: screenshot
comment: Screenshot Capture
created_at: 2024-05-31 13:41:44+00:00
pushed_at: 2024-05-31 13:41:55+00:00
url:
https://github.com/jonpalmisc/ida_screenshot/blob/c2f452b8a965ecd6b533d33630642b0ba12c484b/screenshot.py

content:
import ida_idaapi
import ida_kernwin

from PyQt5.QtCore import QPoint, QRect, QSize
from PyQt5.QtGui import QPixmap, QRegion
from PyQt5.QtWidgets import QApplication, QWidget


def scale_size(r: QRect, scale: int) -> QSize:
    s = QRect(r)

    s.setWidth(r.width() * scale)
    s.setHeight(r.height() * scale)

    return s.size()


def render_widget_img(widget: QWidget, scale: int) -> QPixmap:
    rect = widget.rect()

    img = QPixmap(scale_size(rect, scale))
    img.setDevicePixelRatio(scale)
    widget.render(img, QPoint(), QRegion(rect))

    return img


class screenshot_handler_t(ida_kernwin.action_handler_t):
    widget_only: bool
    save_to_file: bool

    def __init__(self, widget_only: bool, save_to_file: bool):
        ida_kernwin.action_handler_t.__init__(self)
        self.widget_only = widget_only
        self.save_to_file = save_to_file

    def activate(self, ctx):  # pyright: ignore
        if self.widget_only:
            cur_twidget = ida_kernwin.get_current_widget()
            target = ida_kernwin.PluginForm.TWidgetToPyQtWidget(cur_twidget)
        else:
            target = QApplication.activeWindow()

        if target is None:
            print("Could not find widget or window!")
            return 1

        scale = ida_kernwin.ask_long(2, "Screenshot scale multiplier:")
        if not scale or scale < 1:
            scale = 1

        img = render_widget_img(target, scale)
        if self.save_to_file:
            path = ida_kernwin.ask_file(True, "screenshot.png", "Save screenshot")
            if path:
                img.save(path)
        else:
            QApplication.clipboard().setImage(img.toImage())

        return 1

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


ACTION_CAPTURE_WIDGET_COPY = "screenshot:CaptureWidgetCopy"
ACTION_CAPTURE_WINDOW_COPY = "screenshot:CaptureWindowCopy"
ACTION_CAPTURE_WIDGET_SAVE = "screenshot:CaptureWidgetSave"
ACTION_CAPTURE_WINDOW_SAVE = "screenshot:CaptureWindowSave"


class screenshot_ui_hooks_t(ida_kernwin.UI_Hooks):
    popup_actions = []

    def finish_populating_widget_popup(self, widget, popup):  # pyright: ignore
        for action in [
            ACTION_CAPTURE_WIDGET_COPY,
            ACTION_CAPTURE_WINDOW_COPY,
            ACTION_CAPTURE_WIDGET_SAVE,
            ACTION_CAPTURE_WINDOW_SAVE,
        ]:
            ida_kernwin.attach_action_to_popup(
                widget,
                popup,
                action,
                f"Sc&reenshot/",
                ida_kernwin.SETMENU_APP,
            )


class screenshot_plugin_t(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_DRAW | ida_idaapi.PLUGIN_HIDE
    help = ""
    comment = "Screenshot Capture"
    wanted_name = "screenshot"
    wanted_hotkey = ""

    ui_hooks: screenshot_ui_hooks_t

    def init(self):
        self.ui_hooks = screenshot_ui_hooks_t()

        ida_kernwin.register_action(
            ida_kernwin.action_desc_t(
                ACTION_CAPTURE_WIDGET_COPY,
                "Copy wid~g~et screenshot to clipboard",
                screenshot_handler_t(True, False),
                None,
                "Copy a screenshot of the current widget to the clipboard",
            )
        )
        ida_kernwin.register_action(
            ida_kernwin.action_desc_t(
                ACTION_CAPTURE_WINDOW_COPY,
                "Copy ~w~indow screenshot to clipboard",
                screenshot_handler_t(False, False),
                None,
                "Copy a screenshot of the current window to the clipboard",
            )
        )
        ida_kernwin.register_action(
            ida_kernwin.action_desc_t(
                ACTION_CAPTURE_WIDGET_SAVE,
                "Save widget screenshot to file...",
                screenshot_handler_t(True, True),
                None,
                "Save a screenshot of the current widget to a file",
            )
        )
        ida_kernwin.register_action(
            ida_kernwin.action_desc_t(
                ACTION_CAPTURE_WINDOW_SAVE,
                "Save window screenshot to file...",
                screenshot_handler_t(False, True),
                None,
                "Save a screenshot of the current window to a file",
            )
        )

        self.ui_hooks.hook()
        return ida_idaapi.PLUGIN_KEEP

    def run(self):  # pyright: ignore
        pass

    def term(self):
        self.ui_hooks.unhook()


def PLUGIN_ENTRY():
    return screenshot_plugin_t()


repository: Goatman13/spu2c
file: spu2c.py
description: CELL SPU to C plugin for IDA
url: https://github.com/Goatman13/spu2c/blob/3d0dc1b34f300a84b3bec50a7df7f0ea7329007f/spu2c.py
created_at: 2024-01-24 21:28:17+00:00
pushed_at: 2024-08-10 08:55:52+00:00
content:
# SPU To C

from ida_bytes import *
from idaapi import *
from idc import *
import idaapi
import ida_bytes
import idc

#Constants
MASK_ALLSET_128 = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF
MASK_ALLSET_96  = 0xFFFFFFFFFFFFFFFFFFFFFFFF
MASK_ALLSET_64  = 0xFFFFFFFFFFFFFFFF
MASK_ALLSET_32  = 0xFFFFFFFF
MASK_ALLSET_16  = 0xFFFF
UP   = 0
DOWN = 1

def get_reg(reg):

	if reg == 0:
		return "lr"
	elif reg == 1:
		return "sp"
	else:
		return "r{:d}".format(reg)

def get_reg_with_field(reg, field):

	if reg == 0:
		reg = "lr"
	elif reg == 1:
		reg = "sp"
	else:
		reg = "r{:d}".format(reg)
	return reg + "[{:d}]".format(field)

def get_preferred_reg(reg):

	return get_reg_with_field(reg, 0)

def get_channel(ca):

	if ca < 31:
		ca_tbl = ["event_status", "event_mask", "event_ack", "signal_notify1", "signal_notify2", "ch5", "ch6", "decrementer", "decrementer", "multisource_sync_req",
		"ch10", "event_mask", "tag_mask", "machine_status", "srr0", "srr0", "ls_address", "mfc_eah", "mfc_eal", "mfc_size",
		"tag_id", "mfc_cmd", "tag_mask", "tag_update", "tag_status", "list_stall_status", "list_stall_ack", "atomic_status", "out_mailbox", "in_mailbox",
		"intr_out_mailbox"]
		return ca_tbl[ca]
	elif ca == 69:
		ca = "set_bkmk_tag"
	elif ca == 70:
		ca = "perf_monitor_start_event"
	elif ca == 71:
		ca = "perf_monitor_stop_event"
	elif ca == 74:
		ca = "rng"
	else:
		ca = "ch{:d}".format(ca)
	return ca
	
def sign_extend_imm10(_16, value):

	if value & 0x200 == 0x200:
		value = (0xFFFFFE00 | value & 0x1FF)
	else:
		value &= 0x1FF
	if _16 == 1:
		value &= 0xFFFF
	return value

def sign_extend_imm16(value):

	if value & 0x8000 == 0x8000:
		return 0xFFFF0000 | value & 0xFFFF
	else:
		return value & 0xFFFF

def imm10_to_signed_string(value):

	sign = ""
	imm = value & 0x3FF
	if (imm > 0x1FF):
		imm = ~imm
		imm &= 0x1FF
		imm += 1
		sign = "-"
	return sign + "0x{:X}".format(imm)

def imm16_to_signed_string(value):

	sign = ""
	imm = value & 0xFFFF
	if (imm > 0x7FFF):
		imm = ~imm
		imm &= 0x7FFF
		imm += 1
		sign = "-"
	return sign + "0x{:X}".format(imm)

def exp2(j):
	i = 0
	r = 1
	while i < j:
		r *= 2
		i += 1
	return "{:.1f}".format(r)

def check_abort_shufb(addr, end, msk_reg, direction):
	
	while addr != end:
		opcode   = get_wide_dword(addr)
		if opcode >> 31 == 1:
			test_reg = (opcode >> 21) & 0x7F
		else:
			test_reg = opcode & 0x7F
		if test_reg == msk_reg:
			return 1
		if direction == DOWN:
			addr += 4
		else:
			addr -= 4
	return 0

def resolve_mask(opcode, full_string, msk):

	field  = 0
	result = 0
	space  = " "
	new_line    = "\n"
	rt = get_reg((opcode >> 21) & 0x7F)
	rb = get_reg((opcode >> 14) & 0x7F)
	ra = get_reg((opcode >> 7) & 0x7F)
	while field <= 15:
		if field > 9:
			space = ""
		if field == 15:
			new_line = ""
		x = (msk >> (15 - field) * 8) & 0xFF
		#print("X = 0x{:02X}".format(x))
		if x < 0x80:
			if x & 0x10 == 0x00:
				full_string += rt + " byte[{:d}] ".format(field,) + space + "= byte[{:d}] from ".format(x&0xF) + ra + new_line
				result |= ((x & 0x0F | 0xA0) << ((15 - field) * 8))
			else: 
				full_string += rt + " byte[{:d}] ".format(field,) + space + "= byte[{:d}] from ".format(x&0xF) + rb + new_line
				result |= ((x & 0x0F | 0xB0) << ((15 - field) * 8))
		else:
			if x < 0xC0:
				full_string += rt + " byte[{:d}] ".format(field,) + space + "= 0x00" + new_line
				result |= (0x00 << ((15 - field) * 8))
			elif x >= 0xC0 and x < 0xE0:
				full_string += rt + " byte[{:d}] ".format(field,) + space + "= 0xFF" + new_line
				result |= (0xFF << ((15 - field) * 8))
			else:
				full_string += rt + " byte[{:d}] ".format(field,) + space + "= 0x80" + new_line
				result |= (0x80 << ((15 - field) * 8))
		field += 1
	#print("MASK 0x{:032X}".format(result))
	#print(full_string)
	return full_string

def shufb_patterns(addr, opcode):
	
	msk_reg = opcode & 0x7F
	limit = addr - 0x500
	msk = 0
	full_string = ".\n"
	while addr > limit:
		# addr + 4 because opcode at address with cref is ok.
		xref_test = get_first_fcref_to(addr+4)
		
		# branch target, we need to check xrefs.
		if xref_test != BADADDR:
			abort = 1
			
			# Test xrefs.
			# Todo 0018064 in pemu 3 elf.
			# When we are sure that current address is unreachable by
			# any other way than that branch, we are safe to go.
			#
			# Good test case for successful abort in pemu 5 00007374, etc.
			while xref_test != BADADDR:
				if xref_test > addr+4:
					direction = DOWN
				else:
					direction = UP
				abort = check_abort_shufb(addr+4, xref_test, msk_reg, direction)
				if abort != 0:
					print("shufb: Aborting due to CREF at 0x{:X}, from 0x{:X}".format(addr+4, xref_test))
					return 0
				xref_test = get_next_fcref_to(addr+4, xref_test)

		test_op = get_wide_dword(addr)
		if (test_op >> 31) & 1 == 1:
			target_reg = (test_op >> 21) & 0x7F
		else:
			target_reg = test_op & 0x7F
		
		if target_reg == msk_reg:
			print("shufb: Using opcode from 0x{:X} as a mask loader".format(addr))
			name = print_insn_mnem(addr)
			if name in ["ila", "il", "ilh", "ilhu"]:
				if name == "ila":
					msk = (test_op >> 7) & 0x3FFFF
				elif name == "il":
					msk = (test_op >> 7) & 0xFFFF
					msk = sign_extend_imm16(msk)
				elif name == "ilh":
					msk = (test_op >> 7) & 0xFFFF
					msk = msk | msk << 16
				else: # name == "ilhu":
					msk = (test_op >> 7) & 0xFFFF
					msk = msk << 16
				print("shufb: Pre shift mask from " + name + " opcode = 0x{:08X}".format(msk))
				msk = (msk | msk << 32 | msk << 64 | msk << 96)
			elif name in ["lqa", "lqr", "lqd"]:
				if is_off1(ida_bytes.get_flags(addr)) == 0 and name == "lqd":
					print("shufb: Can't resolve mask, unresolved lqd at 0x{:X}".format(addr))
					return 0
				msk_addr = get_operand_value(addr, 1);
				msk = get_wide_dword(msk_addr) << 96 | get_wide_dword(msk_addr+4) << 64 | get_wide_dword(msk_addr+8) << 32 | get_wide_dword(msk_addr+12)
				print("shufb: Using mask from " + name + " opcode, mask at 0x{:X}".format(msk_addr))
			elif name in ["cbd", "chd", "cwd", "cdd"]:
				print("shufb: Generating mask from Generate Controls instruction")
				base   = 0x101112131415161718191A1B1C1D1E1F
				if name == "cbd":
					shift = ((15 - (get_operand_value(addr, 1) & 0x0F)) * 8)
					base  = base & ~(0xFF << shift)
					msk   = 0x03 << shift
				elif name == "chd":
					shift = ((14 - (get_operand_value(addr, 1) & 0x0E)) * 8)
					base  = base & ~(MASK_ALLSET_16 << shift)
					msk   = 0x0203 << shift
				elif name == "cwd":
					shift = ((12 - (get_operand_value(addr, 1) & 0x0C)) * 8)
					base  = base & ~(MASK_ALLSET_32 << shift)
					msk   = 0x00010203 << shift
				else:
					shift = ((8 - (get_operand_value(addr, 1) & 0x08)) * 8)
					base  = base & ~(MASK_ALLSET_64 << shift)
					msk   = 0x0001020304050607 << shift
				msk  |= base
				ctrl_reg = get_reg((test_op >> 7) & 0x7F)
				if 	ctrl_reg != "sp":
					full_string += "shufb: WARNING!\n" + name + " at 0x{:X} is not using sp register as a base!\n".format(addr)
					full_string += "Mask can be inaccurate if " + ctrl_reg + "[32b][0] & 0x0F != 0\n"
			elif name == "orbi":
				tra    = (test_op >> 7) & 0x7F
				timm   = (test_op >> 14) & 0xFF
				timm   = (timm | timm << 8 | timm << 16 | timm << 24) & MASK_ALLSET_32
				timm   = (timm | timm << 32 | timm << 64 | timm << 96) & MASK_ALLSET_128
				taddr  = addr - 4
				tlimit = taddr - 0x300
				while taddr > tlimit:
					# taddr + 4 because opcode at address with cref is ok. 
					if get_first_fcref_to(taddr + 4) != BADADDR:
						# branch target, we need to abandon searching.
						# Result can be inaccurate.
						print("shufb: Aborting post orbi search due to CREF AT 0x{:X}".format(taddr))
						return 0
					test_op = get_wide_dword(taddr)
					if (test_op >> 31) & 1 == 1:
						# RRRR
						target_reg = (test_op >> 21) & 0x7F
					else:
						target_reg = test_op & 0x7F
					
					if target_reg == tra:
						print("shufb: Orbi is using opcode from 0x{:X} as a mask loader".format(taddr))
						name = print_insn_mnem(taddr)
						if name == "ila":
							msk = (test_op >> 7) & 0x3FFFF
							print("shufb: Orbi pre shift mask from ila opcode = 0x{:08X}".format(msk))
							msk  = (msk | msk << 32 | msk << 64 | msk << 96)
							msk |= timm
							break
						#Unsafe, disabled for now
						#elif name in ["lqa", "lqr"]:
						#	msk_addr = get_operand_value(taddr, 1);
						#	msk  = get_wide_dword(msk_addr) << 96 | get_wide_dword(msk_addr+4) << 64 | get_wide_dword(msk_addr+8) << 32 | get_wide_dword(msk_addr+12)
						#	msk |= timm
						#	print("shufb!!!!: Using mask from lq opcode at 0x{:X}".format(msk_addr))
						#	break
						else:
							print("shufb: Aborting, orbi use unsupported mask base opcode")
							return 0
					taddr -= 4
				if taddr == tlimit:
					print("shufb: Aborting in orbi path, opcode with mask not found")
					return 0					
			else:
				print("shufb: Can't resolve mask, unsupported opcode")
				return 0
			break
		addr -= 4
		if addr == limit:
			print("shufb: Aborting, opcode with mask not found")
			return 0
	print("shufb: Mask = 0x{:032X}".format(msk))
	return resolve_mask(opcode, full_string, msk)
		

def avgb(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt +"[16x8b] = (" + ra + " + " + rb + " + 1) >> 1 (sum before shift is 9 bits value)"

def absdb(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt +"[16x8b] = if (" + ra + " < " + rb + "): " + rb + " - " + ra + ", else: " + ra + " - " + rb

def andc(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt +"[4x32b] = " + ra + " & ~" + rb

def andbi(opcode):

	imm    = (opcode >> 14) & 0xFF
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt +"[16x8b] = " + ra + " & 0x{:02X}".format(imm)

def andhi(opcode):

	imm    = (opcode >> 14) & 0x3FF
	imm    = sign_extend_imm10(1, imm)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt +"[8x16b] = " + ra + " & 0x{:04X}".format(imm)

def andi(opcode):

	imm    = (opcode >> 14) & 0x3FF
	imm    = sign_extend_imm10(0, imm)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt +"[4x32b] = " + ra + " & 0x{:08X}".format(imm)

def orc(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt +"[4x32b] = " + ra + " | ~" + rb

def orbi(opcode):

	imm    = (opcode >> 14) & 0xFF
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt +"[16x8b] = " + ra + " | 0x{:X}".format(imm)

def orhi(opcode):

	imm    = (opcode >> 14) & 0x3FF
	imm    = sign_extend_imm10(1, imm)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt +"[8x16b] = " + ra + " | 0x{:X}".format(imm)

def ori(opcode):

	imm    = (opcode >> 14) & 0x3FF
	imm    = sign_extend_imm10(0, imm)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt +"[4x32b] = " + ra + " | 0x{:X}".format(imm)

def orx(opcode):

	ra     = (opcode >> 7) & 0x7F
	rt     = get_reg(opcode & 0x7F)
	return rt +"[32b][0] = " + get_reg_with_field(ra,0) + " | " + get_reg_with_field(ra,1) + " | " + get_reg_with_field(ra,2) + " | " + get_reg_with_field(ra,3) + " (lower 96 bits of " + rt + " = 0)"

def xorbi(opcode):

	imm    = (opcode >> 14) & 0xFF
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt +"[16x8b] = " + ra + " ^ 0x{:X}".format(imm)

def xorhi(opcode):

	imm    = (opcode >> 14) & 0x3FF
	imm    = sign_extend_imm10(1, imm)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt +"[8x16b] = " + ra + " ^ 0x{:X}".format(imm)

def xori(opcode):

	imm    = (opcode >> 14) & 0x3FF
	imm    = sign_extend_imm10(0, imm)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt +"[4x32b] = " + ra + " ^ 0x{:X}".format(imm)

def eqv(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt +"[4x32b] = " + ra + " ^ ~" + rb + " (If the bit in" + ra + " and " + rb + " are the same, the result bit is 1 else 0)"

def fsmbi(opcode):

	rt      = get_reg(opcode & 0x7F)
	i44     = (opcode >> 7)  & 0xF
	i43     = (opcode >> 11) & 0xF
	i42     = (opcode >> 15) & 0xF
	i41     = (opcode >> 19) & 0xF
	mask_44 = 0
	mask_43 = 0
	mask_42 = 0
	mask_41 = 0
	i = 0
	while i < 4:
		mask_temp = i44 & (1<<i)
		if mask_temp != 0:
			mask_44 = mask_44 | (0xFF << (i*8))
		i += 1
	i = 0
	while i < 4:
		mask_temp = i43 & (1<<i)
		if mask_temp != 0:
			mask_43 = mask_43 | (0xFF << (i*8))
		i += 1
	i = 0
	while i < 4:
		mask_temp = i42 & (1<<i)
		if mask_temp != 0:
			mask_42 = mask_42 | (0xFF << (i*8))
		i += 1
	i = 0
	while i < 4:
		mask_temp = i41 & (1<<i)
		if mask_temp != 0:
			mask_41 = mask_41 | (0xFF << (i*8))
		i += 1
	return rt + "[128b] = 0x{:08X}:{:08X}:{:08X}:{:08X}".format(mask_41,mask_42,mask_43,mask_44)

def xsbh(opcode):

	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt +"[8x16b] = SignExtend16(" + ra + " & 0xFF)"

def xshw(opcode):

	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt +"[4x32b] = SignExtend32(" + ra + " & 0xFFFF)"

def xswd(opcode):

	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt +"[2x64b] = SignExtend64(" + ra + " & 0xFFFFFFFF)"

def ilh(opcode):

	rt     = get_reg(opcode & 0x7F)
	imm    = (opcode >> 7) & 0xFFFF
	return rt +"[8x16b] = 0x{:X}".format(imm)

def ilhu(opcode):

	rt     = get_reg(opcode & 0x7F)
	imm    = (opcode >> 7) & 0xFFFF
	return rt +"[4x32b] = 0x{:04X}0000".format(imm)

def il(opcode):

	rt     = get_reg(opcode & 0x7F)
	val    = (opcode >> 7) & 0xFFFF
	imm    = sign_extend_imm16(val)
	if val > 0x7FFF:
		imms   = imm16_to_signed_string(val)
		return rt +"[4x32b] = 0x{:X}".format(imm) + " (" + imms + ")"
	return rt +"[4x32b] = 0x{:X}".format(imm)

def ila(opcode):

	rt     = get_reg(opcode & 0x7F)
	imm    = (opcode >> 7) & 0x3FFFF
	return rt +"[4x32b] = 0x{:X}".format(imm)

def iohl(opcode):

	rt     = get_reg(opcode & 0x7F)
	imm    = (opcode >> 7) & 0xFFFF
	return rt +"[4x32b] = " + rt + " | 0x{:X}".format(imm)

def ah(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt +"[8x16b] = " + rb + " + " + ra

def a(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt +"[4x32b] = " + rb + " + " + ra

def addx(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt +"[4x32b] = " + ra + " + " + rb + " + (" + rt + " & 1)"

def sfh(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt +"[8x16b] = " + rb + " - " + ra

def sf(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt +"[4x32b] = " + rb + " - " + ra

def sfx(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt +"[4x32b] = (" + ra + " - " + rb + ") - ((" + rt + " & 1) ^ 1)"

def ahi(opcode):

	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	imm    = (opcode >> 14) & 0x3FF
	imm    = imm10_to_signed_string(imm)
	return rt +"[8x16b] = " + ra + " + " + imm

def ai(opcode):

	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	imm    = (opcode >> 14) & 0x3FF
	imm    = imm10_to_signed_string(imm)
	return rt +"[4x32b] = " + ra + " + " + imm

def sfhi(opcode):

	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	imm    = (opcode >> 14) & 0x3FF
	imm    = imm10_to_signed_string(imm)
	return rt +"[8x16b] = " + imm + " - " + ra

def sfi(opcode):

	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	imm    = (opcode >> 14) & 0x3FF
	imm    = imm10_to_signed_string(imm)
	return rt +"[4x32b] = " + imm + " - " + ra

def bg(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return "[4x32b] if (u32)" + ra + " > (u32)" + rb + ": " + rt + " = 0, else " + rt + " = 0x00000001"

#todo
#def bgx(opcode):
#
#	rt     = opcode & 0x7F
#	ra     = (opcode >> 7) & 0x7F
#	rb     = (opcode >> 14) & 0x7F
#	return "[4x32b] if (u32)" + ra + " > (u32)" + rb + ": " + rt + " = 0, else " + rt + " = 0x00000001"

def mpy(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[4x32b] = (s16)(" + ra + " & 0xFFFF) * (s16)(" + rb + " & 0xFFFF)"

def mpyu(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[4x32b] = (" + ra + " & 0xFFFF) * (" + rb + " & 0xFFFF)"

# signed rc? 
def mpya(opcode):

	rt     = get_reg((opcode >> 21) & 0x7F)
	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rc     = get_reg(opcode & 0x7F)
	return rt + "[4x32b] = ((s16)(" + ra + " & 0xFFFF) * (s16)(" + rb + " & 0xFFFF)) + " + rc

# signed? 
def mpyh(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[4x32b] = ((((" + ra + " >> 16) & 0xFFFF) * (" + rb + " & 0xFFFF)) << 16) & 0xFFFF0000"

# signed? 
def mpys(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[4x32b] = sign_extend32(((" + ra + " & 0xFFFF) * (" + rb + " & 0xFFFF)) >> 16)"

def mpyhh(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[4x32b] = (s16)((" + ra + " >> 16) & 0xFFFF) * (s16)((" + rb + " >> 16) & 0xFFFF)"

# signed rt? 
def mpyhha(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[4x32b] = (s16)((" + ra + " >> 16) & 0xFFFF) * (s16)((" + rb + " >> 16) & 0xFFFF) + " + rt

def mpyhhu(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[4x32b] = ((" + ra + " >> 16) & 0xFFFF) * ((" + rb + " >> 16) & 0xFFFF)"

def mpyhhau(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[4x32b] = ((" + ra + " >> 16) & 0xFFFF) * ((" + rb + " >> 16) & 0xFFFF) + " + rt

def mpyi(opcode):

	imm    = (opcode >> 14) & 0x3FF
	imm    = imm10_to_signed_string(imm)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[4x32b] = (s16)(" + ra + " & 0xFFFF) * " + imm

def mpyui(opcode):

	imm    = (opcode >> 14) & 0x3FF
	imm    = sign_extend_imm10(1, imm)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[4x32b] = (" + ra + " & 0xFFFF) * 0x{:X}".format(imm)


####################
# Imm shift start: #
####################

def shlqbii(opcode):
	shift  = (opcode >> 14) & 7
	result = MASK_ALLSET_128 << shift
	result &= MASK_ALLSET_128
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	a      = result & MASK_ALLSET_32
	return rt + "[128b] = (" + ra + " << {:d}) & 0xFFFFFFFF:FFFFFFFF:FFFFFFFF:{:08X}".format(shift,a)

def shlqbyi(opcode):
	shift  = (opcode >> 14) & 0x1F
	shift  *= 8
	result = MASK_ALLSET_128 << shift
	result &= MASK_ALLSET_128
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	a      = (result >> 96) & MASK_ALLSET_32
	b      = (result >> 64) & MASK_ALLSET_32
	c      = (result >> 32) & MASK_ALLSET_32
	d      = result & MASK_ALLSET_32
	return rt + "[128b] = (" + ra + " << {:d}) & 0x{:08X}:{:08X}:{:08X}:{:08X}".format(shift,a,b,c,d)

def shli(opcode):
	shift  = (opcode >> 14) & 0x3F
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	a      = (MASK_ALLSET_32 << shift) & MASK_ALLSET_32
	return rt + "[4x32b] = (" + ra + " << {:d}) & 0x{:08X}".format(shift,a)

def shlhi(opcode):
	shift  = (opcode >> 14) & 0x1F
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	a      = (MASK_ALLSET_16 << shift) & MASK_ALLSET_16
	return rt + "[8x16b] = (" + ra + " << {:d}) & 0x{:04X}".format(shift,a)

#####################
# Imm rotate start: #
#####################

# Right arithm shift 4x32 by bit
def rotmai(opcode):
	shift  = (0 -(opcode >> 14)) & 0x3F
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	#fixme: arithm
	return rt + "[4x32b] = (" + ra + " >> {:d})".format(shift)

# Right arithm shift 8x16 by bit
def rotmahi(opcode):
	shift  = (0 -(opcode >> 14)) & 0x1F
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	#fixme: arithm
	return rt + "[8x16b] = (" + ra + " >> {:d})".format(shift)

# Right shift 128 by bit
def rotqmbii(opcode):
	shift  = (0 -(opcode >> 14)) & 0x7
	const  = MASK_ALLSET_128
	result = const >> shift
	result &= MASK_ALLSET_128
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	a = (result >> 96) & MASK_ALLSET_32
	return rt + "[128b] = (" + ra + " >> {:d}) & 0x{:08X}:FFFFFFFF:FFFFFFFF:FFFFFFFF".format(shift,a)

# Right shift 128 by byte
def rotqmbyi(opcode):
	shift  = (0 -(opcode >> 14)) & 0x1F
	shift *= 8
	result = MASK_ALLSET_128 >> shift
	result &= MASK_ALLSET_128
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	a = (result >> 96) & MASK_ALLSET_32
	b = (result >> 64) & MASK_ALLSET_32
	c = (result >> 32) & MASK_ALLSET_32
	d = result & MASK_ALLSET_32
	return rt + "[128b] = (" + ra + " >> {:d}) & 0x{:08X}:{:08X}:{:08X}:{:08X}".format(shift,a,b,c,d)

# Right shift 4x32 by bit
def rotmi(opcode):
	shift  = (0 -(opcode >> 14)) & 0x3F
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	a      = MASK_ALLSET_32 >> shift
	return rt + "[4x32b] = (" + ra + " >> {:d}) & 0x{:08X}".format(shift,a)

# Right shift 8x16 by bit
def rothmi(opcode):
	shift  = (0 -(opcode >> 14)) & 0x1F
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	a      = MASK_ALLSET_16 >> shift
	return rt + "[8x16b] = (" + ra + " >> {:d}) & 0x{:04X}".format(shift,a)

# Left rotate 128 by bit
def rotqbii(opcode):
	shift  = (opcode >> 14) & 0x7
	const  = MASK_ALLSET_128
	result1 = const << shift
	result2 = const >> (128-shift)
	result1 &= MASK_ALLSET_128
	result2 &= MASK_ALLSET_128
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	# fixme: maybe split in 4 rows? Ugly
	return rt + "[128b] = (" + ra + " << {:d}) & 0x{:032X} | (".format(shift,result1) + ra + " >> {:d}) & 0x{:032X}".format((128-shift),result2)

# Left rotate 128 by byte
def rotqbyi(opcode):
	shift  = (opcode >> 14) & 0xF
	shift  *= 8
	const  = 0xAAAAAAAABBBBBBBBCCCCCCCCDDDDDDDD
	result = const << shift | const >>( 128-shift)
	result &= MASK_ALLSET_128
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	a = (result >> 96) & MASK_ALLSET_32
	b = (result >> 64) & MASK_ALLSET_32
	c = (result >> 32) & MASK_ALLSET_32
	d = result & MASK_ALLSET_32
	return rt + "[128b] = " + ra + " : {:08X}:{:08X}:{:08X}:{:08X}".format(a,b,c,d)

# Left rotate 4x32 by bit
def roti(opcode):
	#fixme: shift is signed?
	shift  = (opcode >> 14) & 0x1F
	result1 = MASK_ALLSET_32 << shift
	result2 = MASK_ALLSET_32 >> (32-shift)
	result1 &= MASK_ALLSET_32
	result2 &= MASK_ALLSET_32
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[4x32b] = (" + ra + " << {:d}) & 0x{:08X} | (".format(shift,result1) + ra + " >> {:d}) & 0x{:08X}".format((32-shift),result2)

# Left rotate 8x16 by bit
def rothi(opcode):
	shift  = (opcode >> 14) & 0xF
	result1 = MASK_ALLSET_16 << shift
	result2 = MASK_ALLSET_16 >> (16-shift)
	result1 &= MASK_ALLSET_16
	result2 &= MASK_ALLSET_16
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	# we can have special case here:
	if shift == 8:
		return rt + "[8x16b] = byteswap16(" + ra + ")"
	return rt + "[8x16b] = (" + ra + " << {:d}) & 0x{:04X} | (".format(shift,result1) + ra + " >> {:d}) & 0x{:04X}".format((16-shift),result2)

#########################
# Non imm rotate start: #
#########################

# Right arithm shift 4x32 by bit from rb
def rotma(opcode):
	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[4x32b] = " + ra + " >> -(" + rb + ") & 0x3F"

# Right arithm shift 8x16 by bit from rb
def rotmah(opcode):
	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[8x16b] = " + ra + " >> -(" + rb + ") & 0x1F"

# Right shift 128 by bit from rb
def rotqmbi(opcode):
	rb     = get_preferred_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[128b] = " + ra + " >> -(" + rb + ") & 7"

# Right shift 128 by byte from rb
def rotqmbybi(opcode):
	rb     = get_preferred_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[128b] = " + ra + " >> ( -(" + rb + " >> 3) & 0x1F) * 8"

# Right shift 128 by byte from rb
def rotqmby(opcode):
	rb     = get_preferred_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[128b] = " + ra + " >> ( -(" + rb + ") & 0x1F) * 8"

# Right logical shift 4x32 by bit from rb
def rotm(opcode):
	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[4x32b] = " + ra + " >> -(" + rb + ") & 0x3F"

# Right logical shift 8x16 by bit from rb
def rothm(opcode):
	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[8x16b] = " + ra + " >> -(" + rb + ") & 0x1F"

# Left rotate 128 by bit from rb
def rotqbi(opcode):
	rb     = get_preferred_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[128b] = " + ra + " << (" + rb + " & 7) | " + ra + " >> 128 - (" + rb + " & 7)"

#fixme wtf
def rotqbybi(opcode):
	rb     = get_preferred_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[128b] = " + ra + " << ((" + rb + " >> 3) & 0xF) * 8 | " + ra + " >> 128 - (((" + rb + " >> 3) & 0xF) * 8)"


# Left rotate 128 by byte from rb
def rotqby(opcode):
	rb     = get_preferred_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[128b] = " + ra + " << (" + rb + " & 0xF) * 8 | " + ra + " >> 128 - ((" + rb + " & 0xF) * 8)"

# Left rotate 4x32 by bit from rb
def rot(opcode):
	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[4x32b] = " + ra + " << (" + rb + " & 0x1F) | " + ra + " >> 32 - (" + rb + " & 0x1F)"

# Left rotate 8x16 by bit from rb
def roth(opcode):
	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[8x16b] = " + ra + " << (" + rb + " & 0xF) | " + ra + " >> 16 - (" + rb + " & 0xF)"

###################
# Branches start: #
###################

# Warning! Code for branches assume that SPU program is PS3 version.
# This mean address is AND with 0x3FFFC.
# This should be changed if you are working with hardware different than PS3.

# Branch indirect (always)
def bi(opcode):
	ra     = get_preferred_reg((opcode >> 7) & 0x7F)
	return "PC = " + ra + " & 0x3FFFC"

# Branch indirect and link if external data
def bisled(addr, opcode):
	ra     = get_preferred_reg((opcode >> 7) & 0x7F)
	rt     = get_preferred_reg(opcode & 0x7F)
	return "if ext_data: PC = " + ra + "& 0x3FFFC, ", + rt + " = 0x{:05X}".format(addr + 4)

# Branch indirect and set link (always)
def bisl(addr, opcode):
	ra     = get_preferred_reg((opcode >> 7) & 0x7F)
	rt     = get_preferred_reg(opcode & 0x7F)
	return "PC = " + ra + " & 0x3FFFC, " + rt + " = 0x{:05X}".format(addr + 4)

# Branch indirect if zero
def biz(addr, opcode):
	ra     = get_preferred_reg((opcode >> 7) & 0x7F)
	rt     = get_preferred_reg(opcode & 0x7F)
	return "if " + rt + " == 0: PC = " + ra + " & 0x3FFFC, else: PC = 0x{:05X}".format(addr + 4)

# Branch indirect if not zero
def binz(addr, opcode):
	ra     = get_preferred_reg((opcode >> 7) & 0x7F)
	rt     = get_preferred_reg(opcode & 0x7F)
	return "if " + rt + " != 0: PC = " + ra + " & 0x3FFFC, else: PC = 0x{:05X}".format(addr + 4)

# Branch indirect if zero halfword
def bihz(addr, opcode):
	ra     = get_preferred_reg((opcode >> 7) & 0x7F)
	rt     = get_preferred_reg(opcode & 0x7F)
	return "if (" + rt + "[32b] & 0xFFFF) == 0: PC = " + ra + " & 0x3FFFC, else: PC = 0x{:05X}".format(addr + 4)

# Branch indirect if not zero halfword
def bihnz(addr, opcode):
	ra     = get_preferred_reg((opcode >> 7) & 0x7F)
	rt     = get_preferred_reg(opcode & 0x7F)
	return "if (" + rt + "[32b] & 0xFFFF) != 0: PC = " + ra + " & 0x3FFFC, else: PC = 0x{:05X}".format(addr + 4)

###################
# Compares start: #
###################

# Compare Equal Byte
def ceqb(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return "[16x8b] if " + ra + " == " + rb + ": " + rt + " = 0xFF, else 0x00"
# Compare Equal Halfword
def ceqh(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return "[8x16b] if " + ra + " == " + rb + ": " + rt + " = 0xFFFF, else 0x0000"

# Compare Equal Word
def ceq(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return "[4x32b] if " + ra + " == " + rb + ": " + rt + " = 0xFFFFFFFF, else 0x00000000"

# Compare Greater Than Byte
def cgtb(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return "[16x8b][signed] if " + ra + " > " + rb + ": " + rt + " = 0xFF, else 0x00"

# Compare Greater Than Halfword
def cgth(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return "[8x16b][signed] if " + ra + " > " + rb + ": " + rt + " = 0xFFFF, else 0x0000"

# Compare Greater Than Word
def cgt(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return "[4x32b][signed] if " + ra + " > " + rb + ": " + rt + " = 0xFFFFFFFF, else 0x00000000"

def clgtb(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return "[16x8b][unsigned] if " + ra + " > " + rb + ": " + rt + " = 0xFF, else 0x00"

def clgth(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return "[8x16b][unsigned] if " + ra + " > " + rb + ": " + rt + " = 0xFFFF, else 0x0000"

def clgt(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return "[4x32b][unsigned] if " + ra + " > " + rb + ": " + rt + " = 0xFFFFFFFF, else 0x00000000"

def ceqbi(opcode):

	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	imm    = (opcode >> 14) & 0xFF
	return "[16x8b] if " + ra + " == 0x{:X}: ".format(imm) + rt + " = 0xFF, else 0x00"

def ceqhi(opcode):

	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	imm    = (opcode >> 14) & 0xFF
	imm    = sign_extend_imm10(1, imm)
	return "[8x16b] if " + ra + " == 0x{:X}: ".format(imm) + rt + " = 0xFFFF, else 0x0000"

def ceqi(opcode):

	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	imm    = (opcode >> 14) & 0xFF
	imm    = sign_extend_imm10(0, imm)
	return "[4x32b] if " + ra + " == 0x{:X}: ".format(imm) + rt + " = 0xFFFFFFFF, else 0x00000000"

def cgtbi(opcode):

	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	imm    = (opcode >> 14) & 0xFF
	sign   = ""
	if (imm > 0x7F):
		imm = ~imm
		imm &= 0x7F
		imm += 1
		sign = "-"
	imm   = sign + "0x{:X}".format(imm)
	return "[16x8b][signed] if " + ra + " > " + imm + ": " + rt + " = 0xFF, else 0x00"

def cgthi(opcode):

	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	imm    = (opcode >> 14) & 0x3FF
	imm    = imm10_to_signed_string(imm)
	return "[8x16b][signed] if " + ra + " > " + imm + ": " + rt + " = 0xFFFF, else 0x0000"

def cgti(opcode):

	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	imm    = (opcode >> 14) & 0x3FF
	imm    = imm10_to_signed_string(imm)
	return "[4x32b][signed] if " + ra + " > " + imm + ": " + rt + " = 0xFFFFFFFF, else 0x00000000"

def clgtbi(opcode):

	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	imm    = (opcode >> 14) & 0xFF
	return "[16x8b][unsigned] if " + ra + " > 0x{:X}: ".format(imm) + rt + " = 0xFF, else 0x00"

def clgthi(opcode):

	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	imm    = (opcode >> 14) & 0xFF
	imm    = sign_extend_imm10(1, imm)
	return "[8x16b][unsigned] if " + ra + " > 0x{:X}: ".format(imm) + rt + " = 0xFFFF, else 0x0000"

def clgti(opcode):

	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	imm    = (opcode >> 14) & 0xFF
	imm    = sign_extend_imm10(0, imm)
	return "[4x32b][unsigned] if " + ra + " > 0x{:X}: ".format(imm) + rt + " = 0xFFFFFFFF, else 0x00000000"

######################
# Channel r/w start: #
######################

def wrch(opcode):

	ca     = get_channel((opcode >> 7) & 0x7F)
	cc     = (opcode >> 7) & 0x7F
	rt     = get_preferred_reg(opcode & 0x7F)
	return ca + "[32b](ch{:d}) = ".format(cc) + rt

# fixme: 96 bits zeroed
def rdch(opcode):

	ca     = get_channel((opcode >> 7) & 0x7F)
	cc     = (opcode >> 7) & 0x7F
	rt     = get_preferred_reg(opcode & 0x7F)
	return rt + " = " + ca + "[32b](ch{:d})".format(cc)

def rchcnt(opcode):

	ca     = get_channel((opcode >> 7) & 0x7F)
	cc     = (opcode >> 7) & 0x7F
	rt     = get_preferred_reg(opcode & 0x7F)
	return rt + " = chnnel count of " + ca + " (ch{:d})".format(cc)

###############
# Misc start: #
###############

def shufb(addr, opcode):

	rt     = get_reg((opcode >> 21) & 0x7F)
	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rc     = get_reg(opcode & 0x7F)
	cmt    = ".\nfor (field = 0; field <= 15; field++)\n{\n\tx = " + rc + ".byte[field]\n\tif (x < 0x80)\n\t{\n\t    if      (x & 0x10) == 0x00) {" + rt + ".byte[field] = " + ra + ".byte[x & 0x0f];}\n\t    else if (x & 0x10) == 0x10) {" + rt + ".byte[field] = " + rb + ".byte[x & 0x0f];}\n\t}\n\telse\n\t{\n\t    if      (x >= 0x80 && x < 0xC0) {" + rt + ".byte[field] = 0x00;}\n\t    else if (x >= 0xC0 && x < 0xE0) {" + rt + ".byte[field] = 0xFF;}\n\t    else if (x >= 0xE0)             {" + rt + ".byte[field] = 0x80;}\n\t}\n}"
	cmt2   = "Are you sure you want to place this comment?\nfor (field = 0; field <= 15; field++)\n{\n    x = " + rc + ".byte[field]\n    if (x < 0x80)\n    {\n        if      (x & 0x10) == 0x00) {" + rt + ".byte[field] = " + ra + ".byte[x & 0x0f];}\n        else if (x & 0x10) == 0x10) {" + rt + ".byte[field] = " + rb + ".byte[x & 0x0f];}\n    }\n    else\n    {\n        if      (x >= 0x80 && x < 0xC0) {" + rt + ".byte[field] = 0x00;}\n        else if (x >= 0xC0 && x < 0xE0) {" + rt + ".byte[field] = 0xFF;}\n        else if (x >= 0xE0)             {" + rt + ".byte[field] = 0x80;}\n    }\n}"
	cmt3   = shufb_patterns(addr - 4, opcode)
	if cmt3 == 0:
		answer = ask_yn(0, cmt2)
		if answer < 1:
			return 1
		return cmt
	return cmt3

def selb(opcode):

	rt     = get_reg((opcode >> 21) & 0x7F)
	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rc     = get_reg(opcode & 0x7F)
	return rt + "[16x8b] = " + rc + " & " + rb + " | ~" + rc + " & " + ra + " (if bit in " + rc + " is 1 take bit from " + rb + ", else from " + ra + ")"

# FPU
def csflt(opcode):

	scale  = (opcode >> 14) & 0xFF
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	scale  = exp2(155 - scale)
	return rt +"[4xfloat] = (float)(s32)" + ra + " / " + scale

def cuflt(opcode):

	scale  = (opcode >> 14) & 0xFF
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	scale  = exp2(155 - scale)
	return rt +"[4xfloat] = (float)(u32)" + ra + " / " + scale

def cflts(opcode):

	scale  = (opcode >> 14) & 0xFF
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	scale  = exp2(173 - scale)
	return rt +"[4x32b] = (s32)((float)" + ra + " * " + scale + ")"

def cfltu(opcode):

	scale  = (opcode >> 14) & 0xFF
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	scale  = exp2(173 - scale)
	return rt +"[4x32b] = (u32)((float)" + ra + " * " + scale + ")"

def fa(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt +"[4xfloat] = " + ra + " + " + rb

def fs(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt +"[4xfloat] = " + ra + " - " + rb

def fm(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt +"[4xfloat] = " + ra + " * " + rb

def fma(opcode):

	rt     = get_reg((opcode >> 21) & 0x7F)
	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rc     = get_reg(opcode & 0x7F)
	return rt +"[4xfloat] = (" + ra + " * " + rb + ") + " + rc

def fnms(opcode):

	rt     = get_reg((opcode >> 21) & 0x7F)
	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rc     = get_reg(opcode & 0x7F)
	return rt +"[4xfloat] = " + rc + " - (" + ra + " * " + rb + ")"

def fms(opcode):

	rt     = get_reg((opcode >> 21) & 0x7F)
	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rc     = get_reg(opcode & 0x7F)
	return rt +"[4xfloat] = (" + ra + " * " + rb + ") - " + rc

def fceq(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return "[4xfloat] if " + ra + " == " + rb + ": " + rt + " = 0xFFFFFFFF, else 0x00000000"

def fcmeq(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return "[4xfloat] if abs(" + ra + ") == abs(" + rb + "): " + rt + " = 0xFFFFFFFF, else 0x00000000"

def fcgt(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return "[4xfloat] if " + ra + " > " + rb + ": " + rt + " = 0xFFFFFFFF, else 0x00000000"

def fcmgt(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return "[4xfloat] if abs(" + ra + ") > abs(" + rb + "): " + rt + " = 0xFFFFFFFF, else 0x00000000"

def fscrwr(opcode):

	ra     = get_reg((opcode >> 7) & 0x7F)
	return "[128b] FPSCR = " + ra

def fscrrd(opcode):

	rt     = get_reg(opcode & 0x7F)
	return "[128b] " + rt + " = FPSCR"

# double
def dfa(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[2xdouble] = " + ra + " + " + rb

def dfs(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[2xdouble] = " + ra + " - " + rb

def dfm(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[2xdouble] = " + ra + " * " + rb

def dfma(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[2xdouble] = (" + ra + " * " + rb + ") + " + rt

# correct?
def dfnms(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[2xdouble] = " + rt + " - (" + ra + " * " + rb + ")"

def dfms(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[2xdouble] = (" + ra + " * " + rb + ") - " + rt

# correct?
def dfnma(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return rt + "[2xdouble] = xor(" + rt + " + (" + ra + " * " + rb + "), 0x80000000:00000000)"

def dfceq(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return "[2xdouble] if " + ra + " == " + rb + ": " + rt + " = 0xFFFFFFFF:FFFFFFFF, else 0x00000000:00000000"

def dfcmeq(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return "[2xdouble] if abs(" + ra + ") == abs(" + rb + "): " + rt + " = 0xFFFFFFFF:FFFFFFFF, else 0x00000000:00000000"

def dfcgt(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return "[2xdouble] if " + ra + " > " + rb + ": " + rt + " = 0xFFFFFFFF:FFFFFFFF, else 0x00000000:00000000"

def dfcmgt(opcode):

	rb     = get_reg((opcode >> 14) & 0x7F)
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	return "[2xdouble] if abs(" + ra + ") > abs(" + rb + "): " + rt + " = 0xFFFFFFFF:FFFFFFFF, else 0x00000000:00000000"

def dftsv(opcode):

	imm    = (opcode >> 14) & 0x7F
	ra     = get_reg((opcode >> 7) & 0x7F)
	rt     = get_reg(opcode & 0x7F)
	if imm != 0:
		full_string = ".[2xdouble]\nif " + ra + " is "
		if imm & 0x40 == 0x40:
			full_string += "NaN or "
		if imm & 0x20 == 0x20:
			full_string += "+Inf or "
		if imm & 0x10 == 0x10:
			full_string += "-Inf or "
		if imm & 0x08 == 0x08:
			full_string += "+0 or "
		if imm & 0x04 == 0x04:
			full_string += "-0 or "
		if imm & 0x02 == 0x02:
			full_string += "+Denormal or "
		if imm & 0x01 == 0x01:
			full_string += "-Denormal:"
		if full_string[-2] == "r":
			full_string = full_string[0:-4] + ":"
		full_string += "\n  " + rt + " = 0xFFFFFFFF:FFFFFFFF"
		full_string += "\nelse:\n  " + rt + " = 0x00000000:00000000"
		return full_string
	else:
		return rt + " = 0x00000000:00000000" 

# Todo:
# non imm shifts, Floating Reciprocal Estimate, frsqest, fi, 
# else: simplify x by 0

def SPUAsm2C(addr):

	opcode = get_wide_dword(addr)
	opcode_name = print_insn_mnem(addr)
	if   opcode_name == "fsmbi": return fsmbi(opcode)
	elif opcode_name == "avgb": return avgb(opcode)
	elif opcode_name == "absdb": return absdb(opcode)
	elif opcode_name == "andc": return andc(opcode)
	elif opcode_name == "andbi": return andbi(opcode)
	elif opcode_name == "andhi": return andhi(opcode)
	elif opcode_name == "andi": return andi(opcode)
	elif opcode_name == "orc": return orc(opcode)
	elif opcode_name == "orbi": return orbi(opcode)
	elif opcode_name == "orhi": return orhi(opcode)
	elif opcode_name == "ori": return ori(opcode)
	elif opcode_name == "orx": return orx(opcode)
	elif opcode_name == "xorbi": return xorbi(opcode)
	elif opcode_name == "xorhi": return xorhi(opcode)
	elif opcode_name == "xori": return xori(opcode)
	elif opcode_name == "eqv": return eqv(opcode)
	elif opcode_name == "ilh": return ilh(opcode)
	elif opcode_name == "ilhu": return ilhu(opcode)
	elif opcode_name == "il": return il(opcode)
	elif opcode_name == "ila": return ila(opcode)
	elif opcode_name == "iohl": return iohl(opcode)
	elif opcode_name == "ah": return ah(opcode)
	elif opcode_name == "a": return a(opcode)
	elif opcode_name == "addx": return addx(opcode)
	elif opcode_name == "sfh": return sfh(opcode)
	elif opcode_name == "sf": return sf(opcode)
	elif opcode_name == "sfx": return sfx(opcode)
	elif opcode_name == "ahi": return ahi(opcode)
	elif opcode_name == "ai": return ai(opcode)
	elif opcode_name == "sfhi": return sfhi(opcode)
	elif opcode_name == "sfi": return sfi(opcode)
	elif opcode_name == "bg": return bg(opcode)
	elif opcode_name == "mpy": return mpy(opcode)
	elif opcode_name == "mpyu": return mpyu(opcode)
	elif opcode_name == "mpya": return mpya(opcode)
	elif opcode_name == "mpyh": return mpyh(opcode)
	elif opcode_name == "mpys": return mpys(opcode)
	elif opcode_name == "mpyhh": return mpyhh(opcode)
	elif opcode_name == "mpyhha": return mpyhha(opcode)
	elif opcode_name == "mpyhhu": return mpyhhu(opcode)
	elif opcode_name == "mpyhhau": return mpyhhau(opcode)
	elif opcode_name == "mpyi": return mpyi(opcode)
	elif opcode_name == "mpyui": return mpyui(opcode)
	elif opcode_name == "shlqbyi": return shlqbyi(opcode)
	elif opcode_name == "shlqbii": return shlqbii(opcode)
	elif opcode_name == "shli": return shli(opcode)
	elif opcode_name == "shlhi": return shlhi(opcode)
	elif opcode_name == "rotmai": return rotmai(opcode)
	elif opcode_name == "rotmahi": return rotmahi(opcode)
	elif opcode_name == "rotqmbii": return rotqmbii(opcode)
	elif opcode_name == "rotqmbyi": return rotqmbyi(opcode)
	elif opcode_name == "rotmi": return rotmi(opcode)
	elif opcode_name == "rothmi": return rothmi(opcode)
	elif opcode_name == "rotqbii": return rotqbii(opcode)
	elif opcode_name == "rotqbyi": return rotqbyi(opcode)
	elif opcode_name == "roti": return roti(opcode)
	elif opcode_name == "rothi": return rothi(opcode)
	elif opcode_name == "rotma": return rotma(opcode)
	elif opcode_name == "rotmah": return rotmah(opcode)
	elif opcode_name == "rotqmbi": return rotqmbi(opcode)
	elif opcode_name == "rotqmbybi": return rotqmbybi(opcode)
	elif opcode_name == "rotqmby": return rotqmby(opcode)
	elif opcode_name == "rotm": return rotm(opcode)
	elif opcode_name == "rothm": return rothm(opcode)
	elif opcode_name == "rotqbi": return rotqbi(opcode)
	elif opcode_name == "rotqby": return rotqby(opcode)
	elif opcode_name == "rot": return rot(opcode)
	elif opcode_name == "roth": return roth(opcode)
	elif opcode_name == "bi": return bi(opcode)
	elif opcode_name == "bisled": return bisled(addr, opcode)
	elif opcode_name == "bisl": return bisl(addr, opcode)
	elif opcode_name == "biz": return biz(addr, opcode)
	elif opcode_name == "binz": return binz(addr, opcode)
	elif opcode_name == "bihz": return bihz(addr, opcode)
	elif opcode_name == "bihnz": return bihnz(addr, opcode)
	elif opcode_name == "xsbh": return xsbh(opcode)
	elif opcode_name == "xshw": return xshw(opcode)
	elif opcode_name == "xswd": return xswd(opcode)
	elif opcode_name == "shufb": return shufb(addr, opcode)
	elif opcode_name == "selb": return selb(opcode)
	elif opcode_name == "wrch": return wrch(opcode)
	elif opcode_name == "rdch": return rdch(opcode)
	elif opcode_name == "rchcnt": return rchcnt(opcode)
	elif opcode_name == "ceqb": return ceqb(opcode)
	elif opcode_name == "ceqh": return ceqh(opcode)
	elif opcode_name == "ceq": return ceq(opcode)
	elif opcode_name == "cgtb": return cgtb(opcode)
	elif opcode_name == "cgth": return cgth(opcode)
	elif opcode_name == "cgt": return cgt(opcode)
	elif opcode_name == "clgtb": return clgtb(opcode)
	elif opcode_name == "clgth": return clgth(opcode)
	elif opcode_name == "clgt": return clgt(opcode)
	elif opcode_name == "ceqbi": return ceqbi(opcode)
	elif opcode_name == "ceqhi": return ceqhi(opcode)
	elif opcode_name == "ceqi": return ceqi(opcode)
	elif opcode_name == "cgtbi": return cgtbi(opcode)
	elif opcode_name == "cgthi": return cgthi(opcode)
	elif opcode_name == "cgti": return cgti(opcode)
	elif opcode_name == "clgtbi": return clgtbi(opcode)
	elif opcode_name == "clgthi": return clgthi(opcode)
	elif opcode_name == "clgti": return clgti(opcode)
	# fpu
	elif opcode_name == "csflt": return csflt(opcode)
	elif opcode_name == "cuflt": return cuflt(opcode)
	elif opcode_name == "cflts": return cflts(opcode)
	elif opcode_name == "cfltu": return cfltu(opcode)
	elif opcode_name == "fa": return fa(opcode)
	elif opcode_name == "fs": return fs(opcode)
	elif opcode_name == "fm": return fm(opcode)
	elif opcode_name == "fma": return fma(opcode)
	elif opcode_name == "fnms": return fnms(opcode)
	elif opcode_name == "fms": return fms(opcode)
	#elif opcode_name == "frest": return frest(opcode)
	#elif opcode_name == "frsqest": return frsqest(opcode)
	#elif opcode_name == "fi": return fi(opcode)
	#elif opcode_name == "frds": return frds(opcode)
	#elif opcode_name == "fesd": return fesd(opcode)
	elif opcode_name == "fceq": return fceq(opcode)
	elif opcode_name == "fcmeq": return fceq(opcode)
	elif opcode_name == "fcgt": return fcgt(opcode)
	elif opcode_name == "fcmgt": return fcgt(opcode)
	elif opcode_name == "fscrwr": return fscrwr(opcode)
	elif opcode_name == "fscrrd": return fscrrd(opcode)
	# fpu double
	elif opcode_name == "dfa": return dfa(opcode)
	elif opcode_name == "dfs": return dfs(opcode)
	elif opcode_name == "dfm": return dfm(opcode)
	elif opcode_name == "dfma": return dfma(opcode)
	elif opcode_name == "dfnms": return dfnms(opcode)
	elif opcode_name == "dfms": return dfms(opcode)
	elif opcode_name == "dfnma": return dfnma(opcode)
	elif opcode_name == "dfceq": return dfceq(opcode)
	elif opcode_name == "dfcmeq": return dfceq(opcode)
	elif opcode_name == "dfcgt": return dfcgt(opcode)
	elif opcode_name == "dfcmgt": return dfcgt(opcode)
	elif opcode_name == "dftsv": return dftsv(opcode)

	return 0

def run_task(start_addr, end_addr, always_insert_comment):

	# convert all instructions within the bounds
	addr = start_addr
	while(addr < end_addr):
		print_str = SPUAsm2C(addr)
		if(print_str != 0 and print_str != 1):
			set_cmt(addr, print_str, False)
		elif (print_str == 0 and always_insert_comment == True):
			msg("0x{:X}: Error converting SPU to C code\n".format(addr))
		addr += 4


def PluginMain():

	# select current line or selected lines
	always_insert_comment = False
	start_addr = read_selection_start()
	end_addr = read_selection_end()
	if(start_addr == BADADDR):
		start_addr = get_screen_ea();
		end_addr = start_addr + 4;
		always_insert_comment = True

	run_task(start_addr, end_addr, always_insert_comment)


def PluginMainF():

	# convert current function
	p_func = get_func(get_screen_ea());
	if(type(p_func) == type(None)):
		msg("Not in a function, so can't do SPU to C conversion for the current function!\n");
		return;
	start_addr = p_func.start_ea;
	end_addr = p_func.end_ea;
	always_insert_comment = False;

	run_task(start_addr, end_addr, always_insert_comment)


#/***************************************************************************************************
#*
#*	Strings required for IDA Pro's PLUGIN descriptor block
#*
#***************************************************************************************************/
#
G_PLUGIN_COMMENT = "SPU To C Conversion Assist"
G_PLUGIN_HELP = "This plugin assists in converting SPU instructions into their relevant C code.\nIt is especially useful for the tricky bit manipulation and shift instructions.\n"
G_PLUGIN_NAME = "SPU To C: Selected Lines"

#/***************************************************************************************************
#*
#*	This 'PLUGIN' data block is how IDA Pro interfaces with this plugin.
#*
#***************************************************************************************************/

class ActionHandler(idaapi.action_handler_t):

    def __init__(self, callback):

        idaapi.action_handler_t.__init__(self)
        self.callback = callback

    def activate(self, ctx):

        self.callback()
        return 1

    def update(self, ctx):

        return idaapi.AST_ENABLE_ALWAYS

def register_actions():

    actions = [
        {
            'id': 'start:plg',
            'name': G_PLUGIN_NAME,
            'hotkey': 'F10',
            'comment': G_PLUGIN_COMMENT,
            'callback': PluginMain,
            'menu_location': 'Start Plg'
        },
        {
            'id': 'start:plg1',
            'name': 'spu2c unimplemented',
            'hotkey': 'Alt-Shift-F10',
            'comment': G_PLUGIN_COMMENT,
            'callback': PluginMainF,
            'menu_location': 'Start Plg1'
        }
    ]

    for action in actions:

        if not idaapi.register_action(idaapi.action_desc_t(
            action['id'], # Must be the unique item
            action['name'], # The name the user sees
            ActionHandler(action['callback']), # The function to call
            action['hotkey'], # A shortcut, if any (optional)
            action['comment'] # A comment, if any (optional)
        )):

            print('Failed to register ' + action['id'])

        if not idaapi.attach_action_to_menu(
            action['menu_location'], # The menu location
            action['id'], # The unique function ID
            0):

            print('Failed to attach to menu '+ action['id'])

class spu_helper_t(idaapi.plugin_t):
	flags = idaapi.PLUGIN_HIDE
	comment = G_PLUGIN_COMMENT
	help = G_PLUGIN_HELP
	wanted_name = G_PLUGIN_NAME
	wanted_hotkey = "F10"

	def init(self):
		if (idaapi.ph.id == idaapi.PLFM_SPU):
			register_actions()
			idaapi.msg("spu2c: loaded\n")
			return idaapi.PLUGIN_KEEP

		return idaapi.PLUGIN_SKIP

	def run(self, arg):
		idaapi.msg("spu2c: run\n")

	def term(self):
		pass

def PLUGIN_ENTRY():
	return spu_helper_t()


repository: Shenmarukai/IdaClass
file: src/IdaClassPlugin.py
wanted_name: IdaStruct
comment: Enables the user to construct and export class definitions from the current database
created_at: 2025-03-05 21:50:12+00:00
pushed_at: 2025-03-05 21:54:25+00:00
url:
https://github.com/Shenmarukai/IdaClass/blob/ebf1a7fb939218a0cd150f23699c0ea8d5557d8e/src/IdaClassPlugin.py

content:
import ida_idaapi
import ida_funcs
import idautils
import ida_typeinf
import IdaClassImpl as IClass

class IdaClassmod(ida_idaapi.plugmod_t):
    def __del__(self):
        print(">>> IdaClassmod: destructor called.")

    def run(self, arg):
        print(f">>> IdaClassmod.run() is invoked with argument value: {arg}.")
        # Create an instance of IdaClassDatabase
        db = IClass.IdaClassDatabase()

        # Iterate over all structures in the IDB
        for ordinal in range(ida_typeinf.get_ordinal_qty(ida_typeinf.get_idati())):
            tif = ida_typeinf.tinfo_t()
            if ida_typeinf.get_numbered_type(ida_typeinf.get_idati(), ordinal, tif):
                struct_name = tif.get_name()
                struct = IClass.Struct(struct_name)
                db.root.add_struct(struct)

                # Iterate over members of the structure
                udt_data = ida_typeinf.udt_type_data_t()
                if tif.get_udt_details(udt_data):
                    for udm in udt_data:
                        member_name = udm.name
                        member_type = self.get_type_from_tinfo(udm.type)
                        member_obj = IClass.Member(member_name, type=member_type)
                        struct.add_member(member_obj)

        # Iterate over all functions in the IDB
        for func_ea in idautils.Functions():
            func_name = ida_funcs.get_func_name(func_ea)
            func = IClass.Function(func_name)
            db.root.add_child(func)

            # Get function arguments and return type
            tif = ida_typeinf.tinfo_t()
            if ida_typeinf.get_tinfo(tif, func_ea):
                func.set_return_type(self.get_type_from_tinfo(tif.get_rettype()))
                for i in range(tif.get_nargs()):
                    arg_name = tif.get_nth_arg_name(i)
                    arg_type = self.get_type_from_tinfo(tif.get_nth_arg(i))
                    arg = IClass.Variable(arg_name, type=arg_type)
                    func.add_argument(arg)

        # Print the database as C++ code
        db.print()

    def get_type_from_member(self, member):
        # Convert IDA member type to Struct
        tif = ida_typeinf.tinfo_t()
        if ida_typeinf.get_member_tinfo(tif, member):
            return self.get_type_from_tinfo(tif)
        return IClass.Int  # Default to Int if type is unknown

    def get_type_from_tinfo(self, tinfo):
        # Convert IDA tinfo_t to Struct
        type_name = tinfo.dstr()
        if type_name == "int":
            return IClass.Int
        elif type_name == "float":
            return IClass.Float
        elif type_name == "char":
            return IClass.Char
        # Add more type mappings as needed
        return IClass.Int  # Default to Int if type is unknown

class IdaClass(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_MULTI | ida_idaapi.PLUGIN_FIX
    comment = "Enables the user to construct and export class definitions from the current database"
    help = "Enables the user to construct and export class definitions from the current database"
    wanted_name = "IdaStruct"
    wanted_hotkey = "Ctrl-Shift-C"

    def init(self):
        print(">>> IdaClass: Init called.")
        return IdaClassmod()

def PLUGIN_ENTRY():
    return IdaClass()


repository: Shirasawa-D/PatternGen
file: PatternGen.py
wanted_name: PatternGen
comment: Pattern Generate tool by:chord
created_at: 2024-12-03 18:02:40+00:00
pushed_at: 2024-12-03 18:39:16+00:00
url:
https://github.com/Shirasawa-D/PatternGen/blob/24a8211bc4867ea2303a912c87a7ba73a696e38a/PatternGen.py

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.widget_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: joxeankoret/idamagicstrings
file: IDAMagicStrings.py
description: An IDA Python script to extract information from string constants.
wanted_name: Magic Strings
comment: IDA Magic Strings plugin
created_at: 2018-11-29 11:36:47+00:00
pushed_at: 2023-10-21 09:48:01+00:00
url:
https://github.com/joxeankoret/idamagicstrings/blob/40e8b332313afe64f6a83783e318e94c56a8fae8/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-2022, 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 QtCore
  from PyQt5 import QtGui
  from PyQt5 import QtWidgets
except ImportError as e:
  print(f'{os.path.basename(__file__)} importerror {e}')

try:
  import nltk
  from nltk.tokenize import word_tokenize
  from nltk.tag import pos_tag
  has_nltk = True
except ImportError as e:
  print(f'{os.path.basename(__file__)} importerror nltk {e}')
  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):
  _, 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):
  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 = 4, strtypes = [0, 1]):
  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

        for head in list(Heads(block.start_ea, block.end_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, "Symbol: %s" % full_path)

  nltk_preprocess(strings)
  if len(d) > 0 and total_files > 0:
    print("Programming languages found:\n")
    for key in src_langs:
      print("  %s %f%%" % (key.ljust(10), src_langs[key] * 100. / total_files))
  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()
    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()
    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: RioKato/dotfiles
file: .idapro/plugins/rtti.py
wanted_name: Create graph by RTTI
created_at: 2022-02-13 05:37:56+00:00
pushed_at: 2025-05-13 14:41:42+00:00
url:
https://github.com/RioKato/dotfiles/blob/157859d9d89c136a23760ec390933d73efa7f48e/.idapro/plugins/rtti.py

content:
from typing import Any
from idaapi import UI_Hooks


def binfmt_32_or_64(bit32: str, bit64: str):
    from idaapi import get_inf_structure

    inf = get_inf_structure()

    if inf.is_64bit():
        return bit64
    else:
        return bit32


def cpu_intel_or_arm(intel: int, arm: int):
    from idaapi import get_inf_structure

    inf = get_inf_structure()

    if inf.procname == 'ARM':
        return arm

    return intel


VTABLE_HEADER = binfmt_32_or_64('iI', 'qQ')
VTABLE_POINTER = binfmt_32_or_64('I', 'Q')
CLASS_TYPE_INFO = binfmt_32_or_64('II', 'QQ')
SI_CLASS_TYPE_INFO = binfmt_32_or_64('III', 'QQQ')
VMI_CLASS_TYPE_INFO = binfmt_32_or_64('IIII', 'QQII')
BASE_CLASS_TYPE_INFO = binfmt_32_or_64('II', 'QQ')
VTABLE_NAME_PATTERN = r"`vtable for'(.+)"
TYPEINFO_NAME_PATTERN = r"`typeinfo for'(.+)"
MANGLED_NAME_VTABLE_CLASS_TYPE_INFO = "_ZTVN10__cxxabiv117__class_type_infoE"
MANGLED_NAME_VTABLE_SI_CLASS_TYPE_INFO = "_ZTVN10__cxxabiv120__si_class_type_infoE"
MANGLED_NAME_VTABLE_VMI_CLASS_TYPE_INFO = "_ZTVN10__cxxabiv121__vmi_class_type_infoE"
FUNCTION_POINTER_MASK = cpu_intel_or_arm(0, 1)


def get_unpacked(ea: int, fmt: str) -> Any:
    from struct import unpack, calcsize
    from idaapi import get_bytes

    size = calcsize(fmt)
    data = get_bytes(ea, size)
    return (ea+size,)+unpack(fmt, data)


def parse_typeinfo(ea: int) -> dict[int, str]:
    from struct import calcsize
    from idc import BADADDR, get_name_ea_simple

    vtable_class_type_info = get_name_ea_simple(
        MANGLED_NAME_VTABLE_CLASS_TYPE_INFO)
    if vtable_class_type_info == BADADDR:
        print('[RTTI] vtable_class_type_info not found')

    vtable_si_class_type_info = get_name_ea_simple(
        MANGLED_NAME_VTABLE_SI_CLASS_TYPE_INFO)
    if vtable_si_class_type_info == BADADDR:
        print('[RTTI] vtable_si_class_type_info not found')

    vtable_vmi_class_type_info = get_name_ea_simple(
        MANGLED_NAME_VTABLE_VMI_CLASS_TYPE_INFO)
    if vtable_vmi_class_type_info == BADADDR:
        print('[RTTI] vtable_vmi_class_type_info not found')

    def create_name(ea: int):
        from re import fullmatch
        from idaapi import INF_SHORT_DEMNAMES, get_name, demangle_name
        if name := get_name(ea):
            if name := demangle_name(name, INF_SHORT_DEMNAMES):
                if match := fullmatch(TYPEINFO_NAME_PATTERN, name):
                    return match.group(1)

        return f'ADDR_{ea:X}'

    _, vtable, _ = get_unpacked(ea, CLASS_TYPE_INFO)
    # vtable -= calcsize(VTABLE_HEADER)
    assert (vtable in
            [vtable_class_type_info, vtable_si_class_type_info, vtable_vmi_class_type_info])

    if vtable == vtable_class_type_info:
        return {0: create_name(ea)}

    elif vtable == vtable_si_class_type_info:
        _, _, _, base_type = get_unpacked(ea, SI_CLASS_TYPE_INFO)
        return {0: create_name(base_type)}

    elif vtable == vtable_vmi_class_type_info:
        ea, _, _, _, base_count = get_unpacked(ea, VMI_CLASS_TYPE_INFO)

        typeinfo = {}
        for _ in range(base_count):
            ea, base_type, offset_flags = get_unpacked(
                ea, BASE_CLASS_TYPE_INFO)
            typeinfo[offset_flags >> 8] = create_name(base_type)

        return typeinfo

    raise ValueError


def parse_vtable(start_ea: int, end_ea: int, typeinfo: dict[int, str]) -> dict[str, tuple[int, int]]:
    from struct import calcsize
    from idaapi import get_segm_by_name

    text = get_segm_by_name('.text')
    extern = get_segm_by_name('extern')
    ea = start_ea
    vtable = {}

    for _ in typeinfo:
        head_ea = ea

        assert (ea + calcsize(VTABLE_HEADER) <= end_ea)
        ea, ofs, _ = get_unpacked(ea, VTABLE_HEADER)
        assert (-ofs in typeinfo)

        while ea < end_ea:
            next_ea, ptr = get_unpacked(ea, VTABLE_POINTER)

            if not text.start_ea <= ptr < text.end_ea and not extern.start_ea <= ptr < extern.end_ea:
                break

            ea = next_ea

        vtable[typeinfo[-ofs]] = (head_ea, ea)

    assert (ea == end_ea)

    return vtable


def create_vtable(name: str, start_ea: int, end_ea: int):
    from contextlib import suppress
    from struct import calcsize
    from idaapi import BADADDR, FF_DWORD, FF_QWORD, get_struc_id, get_struc, add_struc, get_max_offset, add_struc_member, del_struc_members, set_member_tinfo, tinfo_t, get_tinfo
    from ida_hexrays import DECOMP_NO_CACHE, decompile
    from idc import apply_type, parse_decl

    start_ea += calcsize(VTABLE_HEADER)
    assert (start_ea <= end_ea)

    sid = get_struc_id(name)

    if sid != BADADDR:
        struc = get_struc(sid)
        del_struc_members(struc, 0, get_max_offset(struc))
    else:
        sid = add_struc(-1, name)
        struc = get_struc(sid)

    size = calcsize(VTABLE_POINTER)
    flag = {
        4: FF_DWORD,
        8: FF_QWORD
    }[size]

    ea = start_ea
    i = 0

    while ea < end_ea:
        ea, ptr = get_unpacked(ea, VTABLE_POINTER)
        ptr -= ptr & FUNCTION_POINTER_MASK
        add_struc_member(
            struc, f'virtual{i}_{ptr:X}', BADADDR, flag, None, size)

        with suppress(Exception):
            decompile(ptr, flags=DECOMP_NO_CACHE)

        ti_fun = tinfo_t()
        if get_tinfo(ti_fun, ptr):
            ti_funptr = tinfo_t()
            ti_funptr.create_ptr(ti_fun)
            set_member_tinfo(struc, struc.get_last_member(), 0, ti_funptr, 0)

        i += 1

    apply_type(start_ea, parse_decl(name, 0), 0)


def popup_main(start_ea: int, end_ea: int):
    from re import fullmatch
    from struct import calcsize
    from idaapi import INF_SHORT_DEMNAMES, get_name, demangle_name

    name = get_name(start_ea)
    assert (name)
    name = demangle_name(name, INF_SHORT_DEMNAMES)
    assert (name)
    match = fullmatch(VTABLE_NAME_PATTERN, name)
    assert (match)
    name = match.group(1)

    assert (start_ea + calcsize(VTABLE_HEADER) <= end_ea)

    _, ofs, typeinfo_ea = get_unpacked(start_ea, VTABLE_HEADER)
    assert (ofs == 0)

    typeinfo = parse_typeinfo(typeinfo_ea)
    print(f'[RTTI] typeinfo = {typeinfo}')

    vtable = parse_vtable(start_ea, end_ea, typeinfo)
    print(f'[RTTI] vtable = {vtable}')

    for subname, (start_ea, end_ea) in vtable.items():
        create_vtable(f'{name}::{subname}::vtable', start_ea, end_ea)


class PopupHooks(UI_Hooks):
    def finish_populating_widget_popup(self, form, popup):
        from idaapi import BWN_DISASMS, SETMENU_INS, get_widget_type, action_handler_t, action_desc_t, attach_dynamic_action_to_popup

        if get_widget_type(form) == BWN_DISASMS:

            class handler(action_handler_t):
                def activate(self, _):
                    from idaapi import read_range_selection

                    ok, start_ea, end_ea = read_range_selection(None)
                    if ok:
                        popup_main(start_ea, end_ea)

            name = 'Create vtable from selection'
            desc = action_desc_t(None, name, handler())
            attach_dynamic_action_to_popup(
                form, popup, desc, name, SETMENU_INS)


hooks = PopupHooks()
hooks.hook()


############################################################################


def typeinfo_paths() -> set[tuple[str, str, bool]]:
    from struct import calcsize
    from re import fullmatch
    from idautils import Names
    from idaapi import INF_SHORT_DEMNAMES, get_name, demangle_name
    from idc import BADADDR, get_name_ea_simple

    vtable_si_class_type_info = get_name_ea_simple(
        MANGLED_NAME_VTABLE_SI_CLASS_TYPE_INFO)
    assert (vtable_si_class_type_info != BADADDR)

    vtable_vmi_class_type_info = get_name_ea_simple(
        MANGLED_NAME_VTABLE_VMI_CLASS_TYPE_INFO)
    assert (vtable_vmi_class_type_info != BADADDR)

    paths = set()

    for ea, child in Names():
        if child := demangle_name(child, INF_SHORT_DEMNAMES):
            if match := fullmatch(TYPEINFO_NAME_PATTERN, child):
                child = match.group(1)

                _, vtable, _ = get_unpacked(ea, CLASS_TYPE_INFO)
                vtable -= calcsize(VTABLE_HEADER)

                if vtable == vtable_si_class_type_info:
                    _, _, _, base_type = get_unpacked(ea, SI_CLASS_TYPE_INFO)

                    if paren := get_name(base_type):
                        if paren := demangle_name(paren, INF_SHORT_DEMNAMES):
                            if match := fullmatch(TYPEINFO_NAME_PATTERN, paren):
                                paren = match.group(1)
                                paths.add((paren, child, True))

                elif vtable == vtable_vmi_class_type_info:
                    ea, _, _, _, base_count = get_unpacked(
                        ea, VMI_CLASS_TYPE_INFO)

                    primary = True

                    for _ in range(base_count):
                        ea, base_type, _ = get_unpacked(
                            ea, BASE_CLASS_TYPE_INFO)

                        if paren := get_name(base_type):
                            if paren := demangle_name(paren, INF_SHORT_DEMNAMES):
                                if match := fullmatch(TYPEINFO_NAME_PATTERN, paren):
                                    paren = match.group(1)
                                    paths.add((paren, child, primary))
                                    primary = False

    return paths


def typeinfo_graphiz(file: str):
    with open(file, 'w') as fd:
        print('digraph {', file=fd)

        for (src, dst, primary) in typeinfo_paths():
            if primary:
                arrow = f'"{src}" -> "{dst}"'
            else:
                arrow = f'"{src}" -> "{dst}" [style="dotted"]'

            print(arrow, file=fd)

        print('}', file=fd)


def PLUGIN_ENTRY():
    from idaapi import PLUGIN_UNL, PLUGIN_OK, plugin_t, ask_file

    class RTTIPlugin(plugin_t):
        flags = PLUGIN_UNL
        wanted_name = 'Create graph by RTTI'

        def init(self):
            return PLUGIN_OK

        def run(self, _):
            file = ask_file(0, '~', 'typeinfo_graphiz')
            typeinfo_graphiz(file)

        def term(self):
            pass

    return RTTIPlugin()


repository: AzzOnFire/msdocsviewer-ex
file: plugins/msdocsviewer_ex.py
wanted_name: API MSDN Docs
comment: API MSDN Docs
created_at: 2024-05-29 21:26:29+00:00
pushed_at: 2024-10-19 10:43:18+00:00
description:
msdocsviewer-ex is a simple tool for viewing Microsoft's win32 API and driver technical documentation offline

url:
https://github.com/AzzOnFire/msdocsviewer-ex/blob/96cbe582cdc07622d6029ba95fb3f19db9cd2522/plugins/msdocsviewer_ex.py

content:
import pickle
import difflib
import pathlib
import zlib

import idaapi
import ida_kernwin
import ida_idaapi
import ida_name
from PyQt5 import QtWidgets


USE_CACHE = True
DOC_DB_PATH = pathlib.Path(__file__).parent / 'msdn.db'
HOTKEY = 'Ctrl-Shift-B'


class DocsDBView(object):
    def __init__(self, filepath: str, use_cache: bool = True):
        self.use_cache = use_cache
        self._data = {}
        self._keys = {}
        if not pathlib.Path(filepath).exists():
            raise FileNotFoundError('Database file not found')
        self._filepath = filepath

    def data(self) -> dict:
        if not self._data:
            with open(self._filepath, 'rb') as handle:
                self._data = pickle.load(handle)
        return self._data

    def __getitem__(self, key: str) -> str:
        value = bytes(self.data()[key])
        if not self.use_cache:
            self._data.clear()
        return zlib.decompress(value).decode()

    def keys(self) -> set:
        if not self._keys:
            self._keys = set(self.data())
            if not self.use_cache:
                self._cache.clear()
        return self._keys       


class MSDN(idaapi.PluginForm):
    widget_name = 'MSDN Docs'
    options = (
        ida_kernwin.PluginForm.WOPN_MENU
        | ida_kernwin.PluginForm.WOPN_ONTOP
        | ida_kernwin.PluginForm.WOPN_RESTORE
        | ida_kernwin.PluginForm.WOPN_PERSIST
        | ida_kernwin.PluginForm.WCLS_CLOSE_LATER
    )

    def OnCreate(self, form):
        self.closed = False
        self.parent = self.FormToPyQtWidget(form)
        self.main_layout = QtWidgets.QVBoxLayout()
        self.markdown_viewer_label = QtWidgets.QLabel()
        self.markdown_viewer = QtWidgets.QTextEdit()
        self.markdown_viewer.setReadOnly(True)
        self.main_layout.addWidget(self.markdown_viewer)
        self.parent.setLayout(self.main_layout)
        
    def OnClose(self, form):
        del form
        self.closed = True

    def Show(self):
        super().Show(self.widget_name, options=self.options)

    def Update(self, content):
        self.markdown_viewer.setMarkdown(content)


class MSDNChoose(idaapi.Choose):
    def __init__(self,):
        super().__init__(
            title="Select API",
            cols=[["API Name", 50]],
            flags=idaapi.Choose.CH_RESTORE,
            embedded=False,
        )
        self.options = []

    def OnGetSize(self) -> int:
        return len(self.options)

    def OnGetLine(self, n: int) -> list:
        return [self.options[n]]

    def Pick(self, options: list):
        self.options = options
        index = super().Show(modal=True)
        if index == -1:
            return None
        return self.options[index]


class MSDNPlugin(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_MOD
    comment = "API MSDN Docs"
    help = ""
    wanted_name = "API MSDN Docs"
    wanted_hotkey = "Ctrl-Shift-Z"

    def init(self):
        self.viewer = MSDN()
        self.choose = MSDNChoose()

        try:
            self.db = DocsDBView(filepath=str(DOC_DB_PATH), use_cache=USE_CACHE)
        except FileNotFoundError as e:
            ida_kernwin.msg(f'{e}\n')
            return idaapi.PLUGIN_SKIP

        self.cache = set(self.db.keys())
        return ida_idaapi.PLUGIN_KEEP

    def run(self, arg):
        api_name = self.get_api_name()
        if not api_name:
            ida_kernwin.msg('description not found\n')
            return

        description = self.db[api_name]
        widget = idaapi.find_widget(self.viewer.widget_name)
        if widget:
            idaapi.activate_widget(widget, True)
        else:
            self.viewer.Show()
        self.viewer.Update(content=description)
 
    def term(self):
        pass

    def get_api_name(self):
        api_name = self.get_api_name_from_selection()
        if not api_name:
            ida_kernwin.msg('invalid selection\n')
            return None

        if api_name not in self.cache:
            matches = difflib.get_close_matches(api_name, self.cache)
            if len(matches) == 0:
                return None
            elif len(matches) == 1:
                return matches[0]
            else:
                api_name = self.choose.Pick(matches)

        return api_name

    @staticmethod
    def get_api_name_from_selection():
        v = ida_kernwin.get_current_viewer()
        highlight = ida_kernwin.get_highlight(v)
        if not highlight:
            return None 

        name, _ = highlight

        # remove common prefixes
        prefixes = [ida_name.FUNC_IMPORT_PREFIX, 'cs:', 'ds:', 'j_']
        for prefix in prefixes:
            if name.startswith(prefix):
                name = name[len(prefix):]
                break
        
        # select function call in decompiler view
        pos = name.find('(')
        if pos != -1:
            name = name[:pos]

        return name


plugin = MSDNPlugin()


def PLUGIN_ENTRY():
    global plugin
    return plugin


repository: dNop90/dOffset
file: IDA Pro/plugins/dOffset.py
description: A plugin for IDA Pro and Cheat Engine to get the offset of the current module
comment: Get current address offset or jump to specific offset
created_at: 2024-05-06 18:04:12+00:00
pushed_at: 2024-05-30 00:48:38+00:00
url:
https://github.com/dNop90/dOffset/blob/c80313f14d20149f6096da3f144c8c4abd247413/IDA%20Pro/plugins/dOffset.py

content:
#------------------------------------------------------------------------------
#   IDA Plugin
#   dOffset
#   
#   Go to offset or copy offset from current line
#------------------------------------------------------------------------------

VERSION = '1.0.0'
AUTHOR = 'dNop90'
PLUGIN_NAME = "dOffset"


import idc
import idaapi
import PyQt5.QtGui as QtGui
import PyQt5.QtCore as QtCore
import PyQt5.QtWidgets as QtWidgets
from PyQt5.Qt import QApplication

class dOffset(idaapi.plugin_t):
    flags = idaapi.PLUGIN_PROC | idaapi.PLUGIN_HIDE
    comment = "Get current address offset or jump to specific offset"
    wanted_name = PLUGIN_NAME
    
    def init(self):
        self.AddActions()
        self._init_hooks()
        
        print(f"[{PLUGIN_NAME}] Loaded")
        return idaapi.PLUGIN_KEEP
        
    def run(self, arg):
        pass
        
    def term(self):
        self.RemoveAllActions()
        self._hooks.unhook()
        
    def _init_hooks(self):
        self._hooks = Hooks()
        self._hooks.ready_to_run = self._init_hexrays_hooks
        self._hooks.hook()

    def _init_hexrays_hooks(self):
        if idaapi.init_hexrays_plugin():
            idaapi.install_hexrays_callback(self._hooks.hxe_callback)
        
    ACTION_NAME_MAINCTX = "dOffset"
    ACTION_NAME_GETOFFSET = "dOffset:GetOffset"
    ACTION_NAME_GETMODULEWITHOFFSET = "dOffset:GetModuleWithOffset"
    ACTION_NAME_JUMPTOOFFSET = "dOffset:JumpToOffset"
    def AddActions(self):
        action_desc_offset = idaapi.action_desc_t(
            self.ACTION_NAME_GETOFFSET,
            "Get offset",
            GetOffsetHandler(),
            None,
            "Get current address offset"
        )
        action_desc_module_offset = idaapi.action_desc_t(
            self.ACTION_NAME_GETMODULEWITHOFFSET,
            "Get module name + offset",
            GetModuleOffsetHandler(),
            None,
            "Get current module name with current address offset"
        )
        action_desc_jump_offset = idaapi.action_desc_t(
            self.ACTION_NAME_JUMPTOOFFSET,
            "Jump to offset",
            JumpToOffsetHandler(),
            None,
            "Jump to address offset"
        )

        idaapi.register_action(action_desc_offset)
        idaapi.register_action(action_desc_module_offset)
        idaapi.register_action(action_desc_jump_offset)

        self.AttachActionToMenu()

    def AttachActionToMenu(self):
        idaapi.attach_action_to_menu(
            "Jump/dOffset/",
            self.ACTION_NAME_JUMPTOOFFSET,
            idaapi.SETMENU_APP
        )
    
    def RemoveAllActions(self):
        idaapi.unregister_action(self.ACTION_NAME_GETOFFSET)
        idaapi.unregister_action(self.ACTION_NAME_GETMODULEWITHOFFSET)
        idaapi.unregister_action(self.ACTION_NAME_JUMPTOOFFSET)


def PLUGIN_ENTRY():
    return dOffset()
    
    
class Hooks(idaapi.UI_Hooks):
    def __init__(self):
        super(Hooks, self).__init__()

    def finish_populating_widget_popup(self, widget, popup):
        if idaapi.get_widget_type(widget) == idaapi.BWN_DISASMS:
            idaapi.attach_action_to_popup(
                widget,
                popup,
                dOffset.ACTION_NAME_MAINCTX,
                "dOffset",
                idaapi.SETMENU_APP
            )
            idaapi.attach_action_to_popup(
                widget,
                popup,
                dOffset.ACTION_NAME_GETOFFSET,
                dOffset.ACTION_NAME_MAINCTX + "/"
            )
            idaapi.attach_action_to_popup(
                widget,
                popup,
                dOffset.ACTION_NAME_GETMODULEWITHOFFSET,
                dOffset.ACTION_NAME_MAINCTX + "/"
            ) 
        return 0
    
    def hxe_callback(self, event, *args):
        if event == idaapi.hxe_populating_popup:
            form, popup, vu = args
            
            if get_cursor_func_ref() == idaapi.BADADDR:
                return 0
            
            idaapi.attach_action_to_popup(
                form,
                popup,
                dOffset.ACTION_NAME_MAINCTX,
                "dOffset",
                idaapi.SETMENU_APP
            )
            idaapi.attach_action_to_popup(
                form,
                popup,
                dOffset.ACTION_NAME_GETOFFSET,
                dOffset.ACTION_NAME_MAINCTX + "/"
            )
            idaapi.attach_action_to_popup(
                form,
                popup,
                dOffset.ACTION_NAME_GETMODULEWITHOFFSET,
                dOffset.ACTION_NAME_MAINCTX + "/"
            )
        
        return 0


#Action Handlers
class GetOffsetHandler(idaapi.action_handler_t):
    def __init__(self):
        idaapi.action_handler_t.__init__(self)
        
    def activate(self, ctx):
        CurrentOffset = idaapi.get_screen_ea() - idaapi.get_imagebase()
        
        clipboard = QApplication.clipboard()
        clipboard.clear(mode=clipboard.Clipboard )
        clipboard.setText(f'{CurrentOffset:X}', mode=clipboard.Clipboard)
        
        print("[%s] %X -> 0x%X" % (PLUGIN_NAME, idaapi.get_screen_ea(), CurrentOffset))
        return 1
    
    def update(self, ctx):
        return idaapi.AST_ENABLE_ALWAYS
    
class GetModuleOffsetHandler(idaapi.action_handler_t):
    def __init__(self):
        idaapi.action_handler_t.__init__(self)
        
    def activate(self, ctx):
        CurrentOffset = idaapi.get_screen_ea() - idaapi.get_imagebase()
        file_name = idc.get_root_filename()
        
        clipboard = QApplication.clipboard()
        clipboard.clear(mode=clipboard.Clipboard )
        clipboard.setText(f'{file_name}+{CurrentOffset:X}', mode=clipboard.Clipboard)
        
        print("[%s] %X -> %s+0x%X" % (PLUGIN_NAME, idaapi.get_screen_ea(), file_name, CurrentOffset))
        return 1
    
    def update(self, ctx):
        return idaapi.AST_ENABLE_ALWAYS
    
class JumpToOffsetHandler(idaapi.action_handler_t):
    def __init__(self):
        idaapi.action_handler_t.__init__(self)
        
    def activate(self, ctx):
        InputOffset = idaapi.ask_str("", 0, "Offset:")
        if (InputOffset is None):
            return 1
        
        if "0x" in InputOffset:
            InputOffset = InputOffset.replace("0x", "")

        if(not is_hex(InputOffset)):
            print("[%s] Invalid offset value %s" % (PLUGIN_NAME, InputOffset))
            return 1
        
        InputOffsetInt = int(InputOffset.lower(), 16)
        AddressToJump = idaapi.get_imagebase() + InputOffsetInt
        if(idaapi.jumpto(AddressToJump)):
            print("[%s] 0x%X -> %X" % (PLUGIN_NAME, InputOffsetInt, AddressToJump))
        else:
            print("[%s] Failed to jump to address with offset 0x%X" % (PLUGIN_NAME, InputOffsetInt))

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

#Other
def is_hex(s):
    try:
        int(s.lower(), 16)
        return True
    except ValueError:
        return False


repository: zhouzhicun/ida-scripts
file: python/IDA/IDAScripts/plugin_zzDeObfuscate.py
description: ida相关脚本，包括字符串解密，BR/BLR混淆, 常见加密算法实现。
created_at: 2024-09-10 08:31:38+00:00
pushed_at: 2024-09-18 03:02:56+00:00
url:
https://github.com/zhouzhicun/ida-scripts/blob/fa533f6a4130921a0013133a9777847842d87c63/python/IDA/IDAScripts/plugin_zzDeObfuscate.py

content:
import idaapi
from ida_idaapi import plugin_t

import zzPluginBase.emuRun as emuRun
import zzPluginBase.emuRunFlare as falreEmuRun
import zzPluginBase.pluginUtils as pluginUtils
import zzPluginBase.decryptString as decryptString


#####################################  插件逻辑  ##############################################

#字符串解密(hikari混淆)
def decryptStringHikari():
    decryptString.HikariHandler.decryptString()

#模拟执行指令
def emuRunCode():
    (valid, start, end, regName) = emuRun.check_code()
    if valid:
        falreEmuRun.emu_run_code(start, end, regName)

#模拟执行指令并Patch
def emuRunPatchCode():
    (valid, start, end, regName) = emuRun.check_code()
    if valid:
        result = falreEmuRun.emu_run_code(start, end, regName)
        emuRun.patch_code(start, end, result)


#################################### 插件配置 ##################################################


ZZDeObfuscate_wanted_name = 'ZZDeObfuscate'
ZZDeObfuscate_comment = ''
ZZDeObfuscate_help = ''
ZZDeObfuscate_wanted_hotkey = ''
ZZDeObfuscate_flags = idaapi.PLUGIN_KEEP

ZZDeObfuscateMenuConfig = pluginUtils.PluginMenuConfig("ZZDeObfuscate/", [
    pluginUtils.PluginSubMenu('my:decrypt_string_hikari', 'Hikari混淆字符串解密', decryptStringHikari),
    pluginUtils.PluginSubMenu('my:emu_run_code', '模拟执行选中指令', emuRunCode),
    pluginUtils.PluginSubMenu('my:emu_run_code_patch', '模拟运行选中指令并Patch', emuRunPatchCode),
])


#################################### 插件框架 ##################################################

class ZZDeObfuscateUIHooks(idaapi.UI_Hooks):
    def finish_populating_widget_popup(self, widget, popup_handle, ctx):
        pluginUtils.registerUIHook(widget, popup_handle, ZZDeObfuscateMenuConfig)


class ZZDeObfuscate(plugin_t):

    wanted_name = ZZDeObfuscate_wanted_name
    comment = ZZDeObfuscate_comment
    help = ZZDeObfuscate_help
    wanted_hotkey = ZZDeObfuscate_wanted_hotkey
    flags = ZZDeObfuscate_flags


    def init(self):

        pluginUtils.registerAction(ZZDeObfuscateMenuConfig)

        global my_ui_hooks
        my_ui_hooks = ZZDeObfuscateUIHooks()
        my_ui_hooks.hook()

        return idaapi.PLUGIN_KEEP

    # 插件运行中 这里是主要逻辑
    def run(self, arg: int):
        pass

    def term(self):
        pass 


# 注册插件
def PLUGIN_ENTRY():
    return ZZDeObfuscate()


repository: z1r00/dia2ol
file: diaphora/jkutils/IDAMagicStrings.py
description: online diff
wanted_name: Magic Strings
comment: IDA Magic Strings plugin
created_at: 2024-12-11 02:37:26+00:00
pushed_at: 2024-12-31 02:47:04+00:00
url:
https://github.com/z1r00/dia2ol/blob/45e3fd5345e13ce389b51c2d6dcf21f30cb16920/diaphora/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: sfn-suitecode/Ida_copilotpro
file: ida_copilot.py
wanted_name: Copilot
comment: Copilot
created_at: 2024-03-13 04:01:13+00:00
pushed_at: 2024-03-13 04:01:22+00:00
url:
https://github.com/sfn-suitecode/Ida_copilotpro/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: tbxy09/mdpickme
file: ftplugin/markdown/mdsug.py
description: A easy use Markdown format tool in Vim
created_at: 2018-09-11 01:56:36+00:00
pushed_at: 2018-09-11 02:21:26+00:00
url:
https://github.com/tbxy09/mdpickme/blob/8dbdaa4ab0e8b62798d40d301898a34e3c9c699a/ftplugin/markdown/mdsug.py

content:
from __future__ import print_function

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

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


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

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

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

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

        print(fmt % args)


_logger = FakeLogger()
def Vim_window():
    vim.command('execute get(g:, "peekaboo_window", s:default_window)')

def plugin_entry_point():
    import vim

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

    else:
        extensions = []

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

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

    target = parse_link(cursor, lines)

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


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

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

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

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

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

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

    return VimOpen(anchor_path(target, current_file))


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

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


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

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


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


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

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

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


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


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


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

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

        else:
            os.startfile(self.target)


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

        # target = parse_path(self.target)
        target = self.target

        # TODO: make space handling more robust?
        # vim.command('e {}'.format(path.path.replace(' ', '\\ ')))
        # vim.command('execute get(g:, "peekaboo_window","vertical botright 30new")')
        mylist=[
        '# {}'.format(target),
        '## {}'.format(target),
        '### {}'.format(target),
        '#### {}'.format(target),
        '##### {}'.format(target),
        '###### {}'.format(target),
        '**{}**'.format(target),
        '``` {} ```'.format(target),
        "'{}'".format(target),
        '*{}*'.format(target),
        '* {}'.format(target),
        '- {}'.format(target),
        '+ {}'.format(target),
        '1. {}'.format(target),
        "[{0}]({1})".format(target.strip('/').split('/')[-1],target),
        '| ------ |:-----:| -----:|',
        ]

        # mylist=[each.replace('DATA',target) for each in mylist]

        vim.command('execute BPop({})'.format(mylist))
        # if path.line is not None:
        #     try:
        #         line = int(path.line)

        #     except:
        #         print('invalid line number')
        #         return

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

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


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

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

        if line is None:
            return

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

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

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

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

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

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

        return target.lower()


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

    except ImportError:
        subprocess.call(args)

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


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

    For example::

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

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

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

    return ParsedPath(path + ext)


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

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


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

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

    link_text, rel_column = select_from_start_of_link(line, column)

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

    m = link_pattern.match(link_text)

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

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

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

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

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

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

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

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

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





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

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

    # TODO: handle escapes

    if start < 0:
        return None, pos

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

    return line[start:], pos - start


if __name__ == "__main__":
    plugin_entry_point()


repository: Ph0en1x-XMU/Ph0en1x-Team
file: Reverse/ida/IDA Pro 7.5.201028 (x86,x64,ARM,ARM64,PPC,PPC64,MIPS)/python/examples/hexrays/vds19.py
wanted_name: optimize x|~x
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/vds19.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:
#        x | ~x => -1
#
#      To see this plugin in action please use be_ornot_be.idb
#

import ida_hexrays
import ida_idaapi

# recognize "x | ~x" and replace by -1
class subinsn_optimizer_t(ida_hexrays.minsn_visitor_t):
    cnt = 0
    def visit_minsn(self):      # for each instruction...
        ins = self.curins       # take a reference to the current instruction

        # THE CORE OF THE PLUGIN IS HERE:
        # check the pattern "x | ~x"
        if ins.opcode == ida_hexrays.m_or and ins.r.is_insn(ida_hexrays.m_bnot) and ins.l == ins.r.d.l:
            if not ins.l.has_side_effects(): # avoid destroying side effects
                # pattern matched, convert to "mov -1, ..."
                ins.opcode = ida_hexrays.m_mov
                ins.l.make_number(-1, ins.r.size)
                ins.r = ida_hexrays.mop_t()
                self.cnt = self.cnt + 1 # number of changes we made
        return 0 # continue traversal

# a custom instruction optimizer, boilerplate code
class sample_optimizer_t(ida_hexrays.optinsn_t):
    def func(self, blk, ins, optflags):
        opt = subinsn_optimizer_t()
        ins.for_all_insns(opt)
        if opt.cnt != 0:                # if we modified microcode,
            blk.mba.verify(True)        # run the verifier
        return opt.cnt                  # report the number of changes

# a plugin interface, boilerplate code
class my_plugin_t(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_HIDE
    wanted_name = "optimize x|~x"
    wanted_hotkey = ""
    comment = ""
    help = ""
    def init(self):
        if ida_hexrays.init_hexrays_plugin():
            self.optimizer = sample_optimizer_t()
            self.optimizer.install()
            print("Installed sample optimizer for 'x | ~x'")
            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: Thehepta/D810Learn
file: D810_Extern/D810.py
wanted_name: D-810
comment: Interface to the D-810 plugin
created_at: 2024-06-02 06:36:20+00:00
pushed_at: 2025-05-11 10:47:50+00:00
url:
https://github.com/Thehepta/D810Learn/blob/52ad29245037f279411a2ddd6b8b6e3f15e9bc9a/D810_Extern/D810.py

content:
import os
import idaapi
import ida_hexrays
import ida_kernwin


from d810.conf import D810Configuration
from d810.manager import D810State, D810_LOG_DIR_NAME
from d810.log import configure_loggers, clear_logs


D810_VERSION = "0.1"

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

    def __init__(self):
        super(D810Plugin, self).__init__()
        self.d810_config = None
        self.state = None
        self.initialized = False


    def reload_plugin(self):
        if self.initialized:
            self.term()

        self.d810_config = D810Configuration()

        #TO-DO: if [...].get raises an exception because log_dir is not found, handle exception
        real_log_dir = os.path.join(self.d810_config.get("log_dir"), D810_LOG_DIR_NAME)

        #TO-DO: if [...].get raises an exception because erase_logs_on_reload is not found, handle exception
        if self.d810_config.get("erase_logs_on_reload"):
            clear_logs(real_log_dir)

        configure_loggers(real_log_dir)
        self.state = D810State(self.d810_config)
        print("D-810 reloading...")
        self.state.start_plugin()
        self.initialized = True


    # IDA API methods: init, run, term
    def init(self):
        if not ida_hexrays.init_hexrays_plugin():
            print("D-810 need Hex-Rays decompiler. Skipping")
            return idaapi.PLUGIN_SKIP

        kv = ida_kernwin.get_kernel_version().split(".")
        if (int(kv[0]) < 7) or ((int(kv[0]) == 7) and (int(kv[1]) < 5)):
            print("D-810 need IDA version >= 7.5. Skipping")
            return idaapi.PLUGIN_SKIP
        print("D-810 initialized (version {0})".format(D810_VERSION))
        return idaapi.PLUGIN_OK


    def run(self, args):
        self.reload_plugin()


    def term(self):
        print("Terminating D-810...")
        if self.state is not None:
            self.state.stop_plugin()

        self.initialized = False


def PLUGIN_ENTRY():
    return D810Plugin()


repository: malefs/security-smell-detector-python-gist
file: gist-hash/793f7ddfc2393ffd0418/snippet.py
wanted_name: DataComment
comment: Add comment on initialisation data with \"mov dword ptr[reg + 0x10], 74657465h\
created_at: 2019-11-23 19:14:08+00:00
pushed_at: 2020-07-26 03:38:04+00:00
url:
https://github.com/malefs/security-smell-detector-python-gist/blob/e90764deb06ae4d3c45e702db7ad00351520348f/gist-hash/793f7ddfc2393ffd0418/snippet.py

content:
##########################################################################
# An IDAPython plugin that generates "comment" for data
#
# Compiler sometimes optimized initialisation of small strings
# and init them with, this example write "Hello!\n":
#     - mov dword ptr[rax + 0], 6c6c6548h
#     - mov dword ptr[rax + 4], a216fh
#
# This plugin will auto comment the following line as:
#     - mov dword ptr[rax + 0], 6c6c6548h; 'Hell'
#     - mov dword ptr[rax + 4], a216fh; 'o\n\x00\x00'
#
# In a second time, you can choose to concatenate them by select both line
# and press "Ctrl+H" or go the Edit menu and click on "Extract selected data"
#
# This will change comment to:
#     - mov dword ptr[rax + 0], 6c6c6548h; 'Hello\n\x00\x00'
#     - mov dword ptr[rax + 4], a216fh;
#
# @BestPig
##########################################################################
import idc
import idaapi
import idautils
import sys


class DataCommentPlugin(idaapi.plugin_t):
    flags = 0
    comment = "Add comment on initialisation data with \"mov dword ptr[reg + 0x10], 74657465h\""
    help = ""
    wanted_name = "DataComment"
    wanted_hotkey = ""

    NAME = "datacomment.py"

    def init(self):
        self.menu_context_extract = idaapi.add_menu_item("Edit/", "Extract selected data", "Ctrl+H", 0, self.dc_extract, (None,))
        idaapi.autoWait()
        for seg_addr in Segments():
            segment = getseg(seg_addr)
            if not segment.perm & SEGPERM_EXEC:
                continue
            self.dcu_extract_between_addr(segment.startEA, segment.endEA)
        return idaapi.PLUGIN_KEEP

    def term(self):
        idaapi.del_menu_item(self.menu_context_extract)
        return None

    def run(self, arg):
        return None

    def dcu_is_wanted_assign(self, instr, addr):
        if instr.get_canon_mnem() != 'mov':
            return False
        if GetOpType(addr, 1) != o_imm:
            return False
        if not getseg(addr).use32():
            return False
        if not 'dword ptr [' in GetOpnd(addr, 0):
            return False
        if (GetOperandValue(addr, 1) & 0xfffff000) == GetOperandValue(addr, 1):
            return False
        return True

    def dcu_extract_between_addr(self, startAddr, endAddr, toMerge=False, force=False):
        current_addr = startAddr
        merged_comment = ''
        first_mov = None
        while current_addr <= endAddr:
            instr = DecodeInstruction(current_addr)
            if not instr:
                current_addr += 1
                continue
            if self.dcu_is_wanted_assign(instr, current_addr):
                if not first_mov:
                    first_mov = current_addr
                data_str = struct.pack('I', GetOperandValue(current_addr, 1) & 0xffffffff)
                if toMerge:
                    merged_comment += data_str
                current_comment = GetCommentEx(current_addr, 1)
                if not current_comment or toMerge or force:
                    exact_comment = repr(data_str)[0 if sys.version_info.major < 3 else 1:]
                    if toMerge and current_comment == exact_comment:
                        MakeRptCmt(current_addr, ' ')
                    elif (not current_comment and not toMerge) or force:
                        MakeRptCmt(current_addr, '%s' % exact_comment)
            current_addr += instr.size
        if merged_comment and first_mov:
            MakeRptCmt(first_mov, '%s' % repr(merged_comment)[0 if sys.version_info.major < 3 else 1:])

    def dc_extract(self, arg):
        multiple, start, end = idaapi.read_selection()
        if multiple:
            self.dcu_extract_between_addr(start, end, True)
        else:
            self.dcu_extract_between_addr(here(), here(), False, True)
        return

def PLUGIN_ENTRY():
    return DataCommentPlugin()


repository: PyEGo/exp-gist
file: hard-gists-me/793f7ddfc2393ffd0418/snippet.py
description: Experiment results on hard-gists dataset
wanted_name: DataComment
comment: Add comment on initialisation data with \"mov dword ptr[reg + 0x10], 74657465h\
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/793f7ddfc2393ffd0418/snippet.py

content:
##########################################################################
# An IDAPython plugin that generates "comment" for data
#
# Compiler sometimes optimized initialisation of small strings
# and init them with, this example write "Hello!\n":
#     - mov dword ptr[rax + 0], 6c6c6548h
#     - mov dword ptr[rax + 4], a216fh
#
# This plugin will auto comment the following line as:
#     - mov dword ptr[rax + 0], 6c6c6548h; 'Hell'
#     - mov dword ptr[rax + 4], a216fh; 'o\n\x00\x00'
#
# In a second time, you can choose to concatenate them by select both line
# and press "Ctrl+H" or go the Edit menu and click on "Extract selected data"
#
# This will change comment to:
#     - mov dword ptr[rax + 0], 6c6c6548h; 'Hello\n\x00\x00'
#     - mov dword ptr[rax + 4], a216fh;
#
# @BestPig
##########################################################################
import idc
import idaapi
import idautils
import sys


class DataCommentPlugin(idaapi.plugin_t):
    flags = 0
    comment = "Add comment on initialisation data with \"mov dword ptr[reg + 0x10], 74657465h\""
    help = ""
    wanted_name = "DataComment"
    wanted_hotkey = ""

    NAME = "datacomment.py"

    def init(self):
        self.menu_context_extract = idaapi.add_menu_item("Edit/", "Extract selected data", "Ctrl+H", 0, self.dc_extract, (None,))
        idaapi.autoWait()
        for seg_addr in Segments():
            segment = getseg(seg_addr)
            if not segment.perm & SEGPERM_EXEC:
                continue
            self.dcu_extract_between_addr(segment.startEA, segment.endEA)
        return idaapi.PLUGIN_KEEP

    def term(self):
        idaapi.del_menu_item(self.menu_context_extract)
        return None

    def run(self, arg):
        return None

    def dcu_is_wanted_assign(self, instr, addr):
        if instr.get_canon_mnem() != 'mov':
            return False
        if GetOpType(addr, 1) != o_imm:
            return False
        if not getseg(addr).use32():
            return False
        if not 'dword ptr [' in GetOpnd(addr, 0):
            return False
        if (GetOperandValue(addr, 1) & 0xfffff000) == GetOperandValue(addr, 1):
            return False
        return True

    def dcu_extract_between_addr(self, startAddr, endAddr, toMerge=False, force=False):
        current_addr = startAddr
        merged_comment = ''
        first_mov = None
        while current_addr <= endAddr:
            instr = DecodeInstruction(current_addr)
            if not instr:
                current_addr += 1
                continue
            if self.dcu_is_wanted_assign(instr, current_addr):
                if not first_mov:
                    first_mov = current_addr
                data_str = struct.pack('I', GetOperandValue(current_addr, 1) & 0xffffffff)
                if toMerge:
                    merged_comment += data_str
                current_comment = GetCommentEx(current_addr, 1)
                if not current_comment or toMerge or force:
                    exact_comment = repr(data_str)[0 if sys.version_info.major < 3 else 1:]
                    if toMerge and current_comment == exact_comment:
                        MakeRptCmt(current_addr, ' ')
                    elif (not current_comment and not toMerge) or force:
                        MakeRptCmt(current_addr, '%s' % exact_comment)
            current_addr += instr.size
        if merged_comment and first_mov:
            MakeRptCmt(first_mov, '%s' % repr(merged_comment)[0 if sys.version_info.major < 3 else 1:])

    def dc_extract(self, arg):
        multiple, start, end = idaapi.read_selection()
        if multiple:
            self.dcu_extract_between_addr(start, end, True)
        else:
            self.dcu_extract_between_addr(here(), here(), False, True)
        return

def PLUGIN_ENTRY():
    return DataCommentPlugin()


repository: PyEGo/exp-gist
file: hard-gists-ego/793f7ddfc2393ffd0418/snippet.py
description: Experiment results on hard-gists dataset
wanted_name: DataComment
comment: Add comment on initialisation data with \"mov dword ptr[reg + 0x10], 74657465h\
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/793f7ddfc2393ffd0418/snippet.py


repository: PyEGo/exp-gist
file: hard-gists-reqs/793f7ddfc2393ffd0418/snippet.py
description: Experiment results on hard-gists dataset
wanted_name: DataComment
comment: Add comment on initialisation data with \"mov dword ptr[reg + 0x10], 74657465h\
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/793f7ddfc2393ffd0418/snippet.py


repository: malefs/security-smell-detector-python-gist
file: gist-src/b3143dde185cecda7c1dee7ffbce5d2c.py
comment: Decode IOCTL codes!
created_at: 2019-11-23 19:14:08+00:00
pushed_at: 2020-07-26 03:38:04+00:00
url:
https://github.com/malefs/security-smell-detector-python-gist/blob/e90764deb06ae4d3c45e702db7ad00351520348f/gist-src/b3143dde185cecda7c1dee7ffbce5d2c.py

content:
############################################################################################
##
## Quick IOCTL Decoder!
##
## All credit for actual IOCTL decode logic:
##      http://www.osronline.com/article.cfm?article=229
## 
##
## To install:
##      Copy script into plugins directory, i.e: C:\Program Files\IDA 6.8\plugins
##
## To run:
##      Highlight IOCTL then right click and select "Decode IOCTL" or press PLUGIN_HOTKEY
##      The decoded IOCTL data will be added as a comment and displayed in the IDA cli
##
############################################################################################

__VERSION__ = '0.1'
__AUTHOR__ = '@herrcore'

PLUGIN_NAME = "Quick IOCTL Decoder"
PLUGIN_HOTKEY = 'Ctrl-Alt-I'


import idaapi

class IOCTL_Decoder():
    DEVICE = [None]*55
    DEVICE[1]="BEEP"
    DEVICE[2]="CD_ROM"
    DEVICE[3]="CD_ROM_FILE_SYSTEM"
    DEVICE[4]="CONTROLLER"
    DEVICE[5]="DATALINK"
    DEVICE[6]="DFS"
    DEVICE[7]="DISK"
    DEVICE[8]="DISK_FILE_SYSTEM"
    DEVICE[9]="FILE_SYSTEM"
    DEVICE[10]="INPORT_PORT"
    DEVICE[11]="KEYBOARD"
    DEVICE[12]="MAILSLOT"
    DEVICE[13]="MIDI_IN"
    DEVICE[14]="MIDI_OUT"
    DEVICE[15]="MOUSE"
    DEVICE[16]="MULTI_UNC_PROVIDER"
    DEVICE[17]="NAMED_PIPE"
    DEVICE[18]="NETWORK"
    DEVICE[19]="NETWORK_BROWSER"
    DEVICE[20]="NETWORK_FILE_SYSTEM"
    DEVICE[21]="NULL"
    DEVICE[22]="PARALLEL_PORT"
    DEVICE[23]="PHYSICAL_NETCARD"
    DEVICE[24]="PRINTER"
    DEVICE[25]="SCANNER"
    DEVICE[26]="SERIAL_MOUSE_PORT"
    DEVICE[27]="SERIAL_PORT"
    DEVICE[28]="SCREEN"
    DEVICE[29]="SOUND"
    DEVICE[30]="STREAMS"
    DEVICE[31]="TAPE"
    DEVICE[32]="TAPE_FILE_SYSTEM"
    DEVICE[33]="TRANSPORT"
    DEVICE[34]="UNKNOWN"
    DEVICE[35]="VIDEO"
    DEVICE[36]="VIRTUAL_DISK"
    DEVICE[37]="WAVE_IN"
    DEVICE[38]="WAVE_OUT"
    DEVICE[39]="8042_PORT"
    DEVICE[40]="NETWORK_REDIRECTOR"
    DEVICE[41]="BATTERY"
    DEVICE[42]="BUS_EXTENDER"
    DEVICE[43]="MODEM"
    DEVICE[44]="VDM"
    DEVICE[45]="MASS_STORAGE"
    DEVICE[46]="SMB"
    DEVICE[47]="KS"
    DEVICE[48]="CHANGER"
    DEVICE[49]="SMARTCARD"
    DEVICE[50]="ACPI"
    DEVICE[51]="DVD"
    DEVICE[52]="FULLSCREEN_VIDEO"
    DEVICE[53]="DFS_FILE_SYSTEM"
    DEVICE[54]="DFS_VOLUME"

    def __init__(self, control_code):
        self.control_code = control_code

    def decode(self):
        out={}
        device_val = (self.control_code >> 16) & 0xFFF
        funcVal = (self.control_code >> 2) & 0xFFF
        if (device_val <= 54) and (device_val != 0):
            device_string = self.DEVICE[device_val]+ " ("+hex(device_val)+")"
        else:
            device_string = hex(device_val)
        function_string = hex(funcVal)
        out["device"] = device_string
        out["function"] = function_string
        access = (self.control_code >> 14) & 3   
        method = self.control_code & 3
        access_string = ""
        if access == 0:
            access_string = "FILE_ANY_ACCESS"
        elif access == 1:
            access_string = "FILE_READ_ACCESS"
        elif access == 2:
            access_string = "FILE_WRITE_ACCESS"
        elif access == 3:
            access_string = "Read and Write"
        method_string = ""
        if method == 0:
            method_string = "METHOD_BUFFERED"
        elif method == 1:
            method_string = "METHOD_IN_DIRECT"
        elif method == 2:
            method_string = "METHOD_OUT_DIRECT"
        elif method == 3:
            method_string = "METHOD_NEITHER"
        out["access"] = access_string
        out["method"] = method_string
        return out


class IDAIOCTLDecoder():
    @staticmethod
    def decode():
        ea = ScreenEA()
        if ea == idaapi.BADADDR:
            idaapi.msg(PLUGIN_NAME + " ERROR: Could not get get_screen_ea()")
            return

        str_id = idaapi.get_highlighted_identifier()
        if not str_id:
            idaapi.msg(PLUGIN_NAME + " ERROR: No Ioctl Code highlighted!")
            return
        try:
            if str_id[-1] == 'h':
                code = int(str_id[:-1], 16)
            elif str_id[-1] == 'o':
                code = int(str_id[:-1], 8)
            elif str_id[-1] == 'b':
                code = int(str_id[:-1], 2)
            else:
                code = int(str_id)
        except ValueError:
            idaapi.msg(PLUGIN_NAME + " ERROR: Not a valid Ioctl Code: " + str(str_id)) 
            return

        try:
            decoder = IOCTL_Decoder(code)
            ioctl_data = decoder.decode()

            #print decoded IOCTL to cli
            msg_string =  "That IOCTL decodes to: \n\tDevice: %s \n\tFunction: %s \n\tAccess: %s \n\tMethod: %s" 
            idaapi.msg(msg_string % (ioctl_data["device"], ioctl_data["function"], ioctl_data["access"], ioctl_data["method"]))

            #add decoded IOCTL as comment
            comment_string =  "dwIoControlCode: \n\t\tDevice: %s \n\t\tFunction: %s \n\t\tAccess: %s \n\t\tMethod: %s"
            idaapi.set_cmt(ea, comment_string % (ioctl_data["device"], ioctl_data["function"], ioctl_data["access"], ioctl_data["method"]), 0)
        except Exception as e:
             idaapi.msg(PLUGIN_NAME + " ERROR: " + str(e))
        return


class IOCTLDecodeHandler(idaapi.action_handler_t):
    def activate(self, ctx):
        IDAIOCTLDecoder.decode()

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

class QuickIOCTLDecoderHooks(idaapi.UI_Hooks):
    def finish_populating_tform_popup(self, form, popup):
        tft = idaapi.get_tform_type(form)
        if tft == idaapi.BWN_DISASM:
            # Note the 'None' as action name (1st parameter).
            # That's because the action will be deleted immediately
            # after the context menu is hidden anyway, so there's
            # really no need giving it a valid ID.
            desc = idaapi.action_desc_t(None, 'Decode IOCTL', IOCTLDecodeHandler())
            idaapi.attach_dynamic_action_to_popup(form, popup, desc, None)

class QuickIOCTLDecoder(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = "Decode IOCTL codes!"

    help = "Highlight IOCTL and right-click 'Decode IOCTL'"
    wanted_name = PLUGIN_NAME
    wanted_hotkey = PLUGIN_HOTKEY
    
 
    def init(self):
        idaapi.msg("Initializing: %s\n" % PLUGIN_NAME)
        global hooks
        hooks = QuickIOCTLDecoderHooks()
        re = hooks.hook()
        return idaapi.PLUGIN_OK

        
    def run(self, arg):
        IDAIOCTLDecoder.decode()
        pass

    def term(self):
        pass

def PLUGIN_ENTRY():
    return QuickIOCTLDecoder()


repository: malefs/security-smell-detector-python-gist
file: gist-hash/b3143dde185cecda7c1dee7ffbce5d2c/snippet.py
comment: Decode IOCTL codes!
created_at: 2019-11-23 19:14:08+00:00
pushed_at: 2020-07-26 03:38:04+00:00
url:
https://github.com/malefs/security-smell-detector-python-gist/blob/e90764deb06ae4d3c45e702db7ad00351520348f/gist-hash/b3143dde185cecda7c1dee7ffbce5d2c/snippet.py

content:
############################################################################################
##
## Quick IOCTL Decoder!
##
## All credit for actual IOCTL decode logic:
##      http://www.osronline.com/article.cfm?article=229
## 
##
## To install:
##      Copy script into plugins directory, i.e: C:\Program Files\IDA 6.8\plugins
##
## To run:
##      Highlight IOCTL then right click and select "Decode IOCTL" or press PLUGIN_HOTKEY
##      The decoded IOCTL data will be added as a comment and displayed in the IDA cli
##
############################################################################################

__VERSION__ = '0.1'
__AUTHOR__ = '@herrcore'

PLUGIN_NAME = "Quick IOCTL Decoder"
PLUGIN_HOTKEY = 'Ctrl-Alt-I'


import idaapi

class IOCTL_Decoder():
    DEVICE = [None]*55
    DEVICE[1]="BEEP"
    DEVICE[2]="CD_ROM"
    DEVICE[3]="CD_ROM_FILE_SYSTEM"
    DEVICE[4]="CONTROLLER"
    DEVICE[5]="DATALINK"
    DEVICE[6]="DFS"
    DEVICE[7]="DISK"
    DEVICE[8]="DISK_FILE_SYSTEM"
    DEVICE[9]="FILE_SYSTEM"
    DEVICE[10]="INPORT_PORT"
    DEVICE[11]="KEYBOARD"
    DEVICE[12]="MAILSLOT"
    DEVICE[13]="MIDI_IN"
    DEVICE[14]="MIDI_OUT"
    DEVICE[15]="MOUSE"
    DEVICE[16]="MULTI_UNC_PROVIDER"
    DEVICE[17]="NAMED_PIPE"
    DEVICE[18]="NETWORK"
    DEVICE[19]="NETWORK_BROWSER"
    DEVICE[20]="NETWORK_FILE_SYSTEM"
    DEVICE[21]="NULL"
    DEVICE[22]="PARALLEL_PORT"
    DEVICE[23]="PHYSICAL_NETCARD"
    DEVICE[24]="PRINTER"
    DEVICE[25]="SCANNER"
    DEVICE[26]="SERIAL_MOUSE_PORT"
    DEVICE[27]="SERIAL_PORT"
    DEVICE[28]="SCREEN"
    DEVICE[29]="SOUND"
    DEVICE[30]="STREAMS"
    DEVICE[31]="TAPE"
    DEVICE[32]="TAPE_FILE_SYSTEM"
    DEVICE[33]="TRANSPORT"
    DEVICE[34]="UNKNOWN"
    DEVICE[35]="VIDEO"
    DEVICE[36]="VIRTUAL_DISK"
    DEVICE[37]="WAVE_IN"
    DEVICE[38]="WAVE_OUT"
    DEVICE[39]="8042_PORT"
    DEVICE[40]="NETWORK_REDIRECTOR"
    DEVICE[41]="BATTERY"
    DEVICE[42]="BUS_EXTENDER"
    DEVICE[43]="MODEM"
    DEVICE[44]="VDM"
    DEVICE[45]="MASS_STORAGE"
    DEVICE[46]="SMB"
    DEVICE[47]="KS"
    DEVICE[48]="CHANGER"
    DEVICE[49]="SMARTCARD"
    DEVICE[50]="ACPI"
    DEVICE[51]="DVD"
    DEVICE[52]="FULLSCREEN_VIDEO"
    DEVICE[53]="DFS_FILE_SYSTEM"
    DEVICE[54]="DFS_VOLUME"

    def __init__(self, control_code):
        self.control_code = control_code

    def decode(self):
        out={}
        device_val = (self.control_code >> 16) & 0xFFF
        funcVal = (self.control_code >> 2) & 0xFFF
        if (device_val <= 54) and (device_val != 0):
            device_string = self.DEVICE[device_val]+ " ("+hex(device_val)+")"
        else:
            device_string = hex(device_val)
        function_string = hex(funcVal)
        out["device"] = device_string
        out["function"] = function_string
        access = (self.control_code >> 14) & 3   
        method = self.control_code & 3
        access_string = ""
        if access == 0:
            access_string = "FILE_ANY_ACCESS"
        elif access == 1:
            access_string = "FILE_READ_ACCESS"
        elif access == 2:
            access_string = "FILE_WRITE_ACCESS"
        elif access == 3:
            access_string = "Read and Write"
        method_string = ""
        if method == 0:
            method_string = "METHOD_BUFFERED"
        elif method == 1:
            method_string = "METHOD_IN_DIRECT"
        elif method == 2:
            method_string = "METHOD_OUT_DIRECT"
        elif method == 3:
            method_string = "METHOD_NEITHER"
        out["access"] = access_string
        out["method"] = method_string
        return out


class IDAIOCTLDecoder():
    @staticmethod
    def decode():
        ea = ScreenEA()
        if ea == idaapi.BADADDR:
            idaapi.msg(PLUGIN_NAME + " ERROR: Could not get get_screen_ea()")
            return

        str_id = idaapi.get_highlighted_identifier()
        if not str_id:
            idaapi.msg(PLUGIN_NAME + " ERROR: No Ioctl Code highlighted!")
            return
        try:
            if str_id[-1] == 'h':
                code = int(str_id[:-1], 16)
            elif str_id[-1] == 'o':
                code = int(str_id[:-1], 8)
            elif str_id[-1] == 'b':
                code = int(str_id[:-1], 2)
            else:
                code = int(str_id)
        except ValueError:
            idaapi.msg(PLUGIN_NAME + " ERROR: Not a valid Ioctl Code: " + str(str_id)) 
            return

        try:
            decoder = IOCTL_Decoder(code)
            ioctl_data = decoder.decode()

            #print decoded IOCTL to cli
            msg_string =  "That IOCTL decodes to: \n\tDevice: %s \n\tFunction: %s \n\tAccess: %s \n\tMethod: %s" 
            idaapi.msg(msg_string % (ioctl_data["device"], ioctl_data["function"], ioctl_data["access"], ioctl_data["method"]))

            #add decoded IOCTL as comment
            comment_string =  "dwIoControlCode: \n\t\tDevice: %s \n\t\tFunction: %s \n\t\tAccess: %s \n\t\tMethod: %s"
            idaapi.set_cmt(ea, comment_string % (ioctl_data["device"], ioctl_data["function"], ioctl_data["access"], ioctl_data["method"]), 0)
        except Exception as e:
             idaapi.msg(PLUGIN_NAME + " ERROR: " + str(e))
        return


class IOCTLDecodeHandler(idaapi.action_handler_t):
    def activate(self, ctx):
        IDAIOCTLDecoder.decode()

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

class QuickIOCTLDecoderHooks(idaapi.UI_Hooks):
    def finish_populating_tform_popup(self, form, popup):
        tft = idaapi.get_tform_type(form)
        if tft == idaapi.BWN_DISASM:
            # Note the 'None' as action name (1st parameter).
            # That's because the action will be deleted immediately
            # after the context menu is hidden anyway, so there's
            # really no need giving it a valid ID.
            desc = idaapi.action_desc_t(None, 'Decode IOCTL', IOCTLDecodeHandler())
            idaapi.attach_dynamic_action_to_popup(form, popup, desc, None)

class QuickIOCTLDecoder(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = "Decode IOCTL codes!"

    help = "Highlight IOCTL and right-click 'Decode IOCTL'"
    wanted_name = PLUGIN_NAME
    wanted_hotkey = PLUGIN_HOTKEY
    
 
    def init(self):
        idaapi.msg("Initializing: %s\n" % PLUGIN_NAME)
        global hooks
        hooks = QuickIOCTLDecoderHooks()
        re = hooks.hook()
        return idaapi.PLUGIN_OK

        
    def run(self, arg):
        IDAIOCTLDecoder.decode()
        pass

    def term(self):
        pass

def PLUGIN_ENTRY():
    return QuickIOCTLDecoder()


