# plugins

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

content:
import ida_idaapi
import ida_kernwin
import ida_funcs

import os
import sys
import codecs

dirpath = os.path.dirname(os.path.abspath(__file__))
cto_dir = os.path.join(dirpath, "cto")
if cto_dir not in sys.path:
    sys.path.append(cto_dir)

ida_idaapi.require("cto_base")
ida_idaapi.require("cto")
ida_idaapi.require("icon")
ida_idaapi.require("syncdata")
ida_idaapi.require("qtutils")
ida_idaapi.require("get_func_relation")


class cto_plugin_t(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_KEEP
    comment = "Call Tree Overviewer"
    toolbar_displayed_name = cto_base.cto_base.orig_title
    toolbar_name = toolbar_displayed_name + 'Toolbar'
    wanted_name = cto.CallTreeOverviewer.orig_title
    wanted_hotkey = "Alt-Shift-C"
    help = "Press '" + wanted_hotkey + "' to display the " + wanted_name + " widget. Then press 'H' to see the help after setting focus to the widget."
    
    action_name = "cto:execute"
    menu_path = "Edit/Plugins/"
    icon = icon.icon_handler(icon.g_icon_data_ascii, True)
    icon_data = icon.icon_data
    icon_data_dark = icon.icon_bg_change(icon_data, True, True)
    if icon_data_dark is None:
        icon_data_dark = icon_data
    act_icon = ida_kernwin.load_custom_icon(data=icon_data, format="png")
    act_icon_dark = ida_kernwin.load_custom_icon(data=icon_data_dark, format="png")
    
    class exec_from_toolbar(ida_kernwin.action_handler_t):
        action_name = "cto:execute_toolbar"
        def __init__(self, plugin):
            ida_kernwin.action_handler_t.__init__(self)
            import weakref
            self.v = weakref.ref(plugin)
        
        def activate(self, ctx):
            self.v().exec_cto()
            
        def update(self, ctx):
            return ida_kernwin.AST_ENABLE_ALWAYS
    
    def init(self):
        ida_kernwin.msg("############### %s (%s) ###############%s" % (self.wanted_name, self.comment, os.linesep))
        ida_kernwin.msg("%s%s" % (self.help, os.linesep))

        # attach to menu
        self.g = None
        ida_kernwin.attach_action_to_menu(
            self.menu_path,
            self.action_name,
            ida_kernwin.SETMENU_APP)

        # attach to toolbar
        ida_kernwin.register_action(
            ida_kernwin.action_desc_t(
            cto_plugin_t.exec_from_toolbar.action_name,
            cto_plugin_t.comment,
            cto_plugin_t.exec_from_toolbar(self),
            None,
            self.wanted_name,
            self.act_icon))
        
        # Insert the action in a toolbar
        ida_kernwin.create_toolbar(self.toolbar_name, self.toolbar_displayed_name)
        ida_kernwin.attach_action_to_toolbar(self.toolbar_name, cto_plugin_t.exec_from_toolbar.action_name)
        
        # install ui hook to enable toolbar later
        self.ph = qtutils.enable_toolbar_t(self.toolbar_name)
        
        # Check if IDA is darkmode or not.
        # It might fail if the mode is darcula or similar themes
        # because the window color is the same as the default theme color.
        # However, it can still distinguish the default and the dark mode.
        #if cto.CallTreeOverviewer.is_dark_mode_with_main():
        if cto_base.cto_base.is_dark_mode_with_main():
            ida_kernwin.update_action_icon(cto_plugin_t.exec_from_toolbar.action_name, self.act_icon_dark)
            ida_kernwin.update_action_icon(self.menu_path + self.wanted_name, self.act_icon_dark)
        return self.flags

    def exec_cto(self):
        global g_cto

        # for degub mode handling
        debug = False
        if "g_debug" in globals() and g_debug:
            debug = True
        elif hasattr(sys.modules["__main__"], "g_debug") and sys.modules["__main__"].g_debug:
            debug = True
            
        curr_view = None
        max_depth = 1
        # for the first message when this plugin is launched by a user by pressing shortcut key or going to menu.
        if 'g_cto' not in globals():
            ida_kernwin.msg("Launching %s (%s) ...%s" % (self.wanted_name, self.comment, os.linesep))
            ida_kernwin.msg("For the first execution, %s will analyze all functions to build the call tree. Please wait for a while.%s" % (self.wanted_name, os.linesep))
        else:
            ida_kernwin.msg("Reloading %s.%s" % (self.wanted_name, os.linesep))
            curr_view = g_cto.curr_view
            if g_cto.max_depth > 1:
                max_depth = g_cto.max_depth
            
        if self.g:
            ea = ida_kernwin.get_screen_ea()
            f = ida_funcs.get_func(ea)
            if f:
                ea = f.start_ea
                
            # convert start address if it is a vfunc
            if ea in self.g.vtbl_refs:
                ea = self.g.vtbl_refs[ea]
            
            drefs = list(get_func_relation.get_drefs_to(ea))
            if not ea in self.g.func_relations and len(drefs) == 0 and ea not in self.g.vtbl_refs:
                ida_kernwin.msg("Current ea is not a function or an address in a function.%s" % (os.linesep))
                ida_kernwin.msg("Not reloaded.%s" % (os.linesep))
                return 0
            
            # save several important data with pickle and closing hooks...
            self.g.close()
        
        # reload the main modules
        ida_idaapi.require("cto")

        # get sync data on a global variable
        sd = syncdata.sync_data()
        sync_data = sd.get()
        # execute the main function
        self.g = cto.exec_cto(cto_data=sync_data, curr_view=curr_view, max_depth=max_depth, debug=debug)
        self.g.exec_ui_action("EmptyStack")
        self.g.__dict__["sd"] = sd
        if sync_data is None:
            self.g.sd.set(self.g.cto_data)
        
        # update icon on menu
        act_icon_id = self.act_icon
        if self.g and self.g.config.dark_mode:
            act_icon_id = self.act_icon_dark
        ida_kernwin.update_action_icon(cto_plugin_t.exec_from_toolbar.action_name, act_icon_id)
        ida_kernwin.update_action_icon(self.menu_path + self.wanted_name, act_icon_id)
        
        # show the messages after launching
        if 'g_cto' not in globals():
            ida_kernwin.msg("Launched %s.%s" % (self.wanted_name, os.linesep))
            ida_kernwin.msg("%s%s" % (self.help.split(". ")[1], os.linesep))
        else:
            ida_kernwin.msg("Reloaded %s.%s" % (self.wanted_name, os.linesep))

        # put the instance in the global variables
        g_cto = self.g
        if not hasattr(sys.modules["__main__"], "g_cto"):
            setattr(sys.modules["__main__"], "g_cto", self.g)
        else:
            sys.modules["__main__"].g_cto = self.g

    def run(self, arg):
        self.exec_cto()
        
    def term(self):
        if self.g:
            self.g.close()
        if hasattr(self.g, "sd"):
            self.g.sd.close()
            
        ida_kernwin.free_custom_icon(self.act_icon)
        ida_kernwin.detach_action_from_menu(self.menu_path, self.action_name)
        
        ida_kernwin.detach_action_from_toolbar(self.wanted_name, cto_plugin_t.exec_from_toolbar.action_name)
        ida_kernwin.delete_toolbar(self.wanted_name)

        if hasattr(sys.modules["__main__"], "g_cto"):
            delattr(sys.modules["__main__"], "g_cto")

def PLUGIN_ENTRY():
    return cto_plugin_t()

def main():
    global g_cto
    g_cto = cto.exec_cto()

if __name__ == '__main__':
    main()


repository: SpaghettDev/BromaIDA
file: BromaIDA.py
description: IDA Broma support, now real
url: https://github.com/SpaghettDev/BromaIDA/blob/a1373d2a0394de1aaa89c3b01706d5b1f65a4fbd/BromaIDA.py
comment: Broma support for IDA.
created_at: 2024-04-28 12:18:16+00:00
pushed_at: 2025-04-21 05:46:30+00:00
content:
# flake8-in-file-ignores: noqa: E402
VERSION = "7.0.3"
__AUTHOR__ = "SpaghettDev"

PLUGIN_NAME = "BromaIDA"
PLUGIN_HOTKEY = "Ctrl+Shift+B"


from pathlib import Path

from idaapi import (
    msg as ida_msg, register_action, unregister_action,
    plugin_t as ida_plugin_t, action_desc_t as ida_action_desc_t,
    PLUGIN_PROC, PLUGIN_KEEP
)
from ida_kernwin import ask_file
from idautils import Names

from broma_ida.utils import stop, path_exists, IDAUtils
from broma_ida.broma.importer import BromaImporter
from broma_ida.broma.exporter import BromaExporter
from broma_ida.ida_ctx_entry import IDACtxEntry

from broma_ida.data.data_manager import DataManager

from broma_ida.ui.simple_popup import SimplePopup
from broma_ida.ui.main_form import MainForm


def on_import(form: MainForm, code: int = 0):
    form.Close(1)

    file_path: str = ask_file(False, "GeometryDash.bro", "bro")

    if not path_exists(file_path, ".bro"):
        SimplePopup("Please select a valid file!", "Cancel").show()
        stop()

    broma_importer = BromaImporter(IDAUtils.get_platform(), file_path)
    broma_importer.parse_file()
    broma_importer.import_into_idb()

    print("[+] BromaIDA: Finished importing bindings from Broma file")
    SimplePopup(
        "Finished importing "
        f"{IDAUtils.get_platform_printable()} "
        "bindings from Broma file.",
        "OK"
    ).show()

    DataManager().close()


def on_export(form: MainForm, code: int = 0):
    form.Close(1)

    platform = IDAUtils.get_platform()

    if platform.startswith("android"):
        SimplePopup(
            "Cannot export bindings from Android binary!", "Cancel"
        ).show()
        stop()

    # for_saving is not True because we need to read the file first
    # which may not even exist if the saving prompt is used
    # (since you can select files that don't exist within said prompt)
    file_path: str = ask_file(False, "GeometryDash.bro", "bro")

    if not path_exists(file_path, ".bro"):
        SimplePopup("Please select a valid file!", "Cancel").show()
        stop()

    broma_exporter = BromaExporter(platform, file_path)

    broma_exporter.import_from_idb(Names())
    broma_exporter.export()

    print(
        "[+] BromaIDA: Finished exporting "
        f"{broma_exporter.num_exports} bindings, "
        f"{broma_exporter.num_ret_exports} return types and "
        f"{broma_exporter.num_args_names_exports} argument names."
    )
    SimplePopup("Finished exporting bindings to Broma file.", "OK").show()

    DataManager().close()


def bida_main():
    """BromaIDA main entrypoint"""
    DataManager().init(
        Path.home() / "broma_ida" / "shelf"
    )

    form_code = MainForm(
        VERSION,
        IDAUtils.get_platform_printable(),
        on_import,
        on_export
    ).show()

    # cancel
    if form_code == 0:
        DataManager().close()


class BromaIDAPlugin(ida_plugin_t):
    """BromaIDA Plugin"""
    flags = PLUGIN_PROC
    comment = "Broma support for IDA."
    help = f"{PLUGIN_HOTKEY} to start the importing/exporting."
    wanted_name = PLUGIN_NAME

    ACTION_BTIDA = "bida:run_bida"
    ACTION_DESC = "Launches the BromaIDA plugin."

    def init(self):
        """Ran on plugin load"""
        self._register_action()
        ida_msg(f"{PLUGIN_NAME} v{VERSION} initialized\n")

        return PLUGIN_KEEP

    def term(self):
        """Ran on plugin unload"""
        self._unregister_action()

        ida_msg(f"{PLUGIN_NAME} v{VERSION} unloaded\n")

    def run(self, arg):
        """Ran on "File -> Script File" (shocker)
        (does not work because this plugin has multiple py files)"""
        bida_main()

    def _register_action(self):
        """Registers BromaIDA's hotkey"""
        hotkey = ida_action_desc_t(
            self.ACTION_BTIDA,
            "BromaIDA",
            IDACtxEntry(bida_main),
            PLUGIN_HOTKEY,
            self.ACTION_DESC
        )

        register_action(hotkey)

    def _unregister_action(self):
        """Unregisters BromaIDA's hotkey"""
        unregister_action(self.ACTION_BTIDA)


def PLUGIN_ENTRY():
    return BromaIDAPlugin()


repository: cra0/ida-scripts
file: plugins/cvutils-getoffset.py
description: Various IDA scripts I've created for Reverse engineering.
comment: Get the Address offset at the cursor location.
created_at: 2022-07-23 02:56:21+00:00
pushed_at: 2024-10-29 02:59:25+00:00
url:
https://github.com/cra0/ida-scripts/blob/ee4ce429b3c059d0441c63e2476a9c0c7d726ccc/plugins/cvutils-getoffset.py

content:
#------------------------------------------------------------------------------
# IDA Plugin to get an offset (rva) at the cursor position.
# Copy the 'cvutils-getoffset.py' into the plugins directory of IDA
#------------------------------------------------------------------------------

VERSION = '1.1.1'
__AUTHOR__ = 'cra0'

PLUGIN_NAME = "Get Address Offset"
PLUGIN_HOTKEY = "Ctrl+Shift+C"

import os
import sys
import idc
import idaapi
import idautils

major, minor = map(int, idaapi.get_kernel_version().split("."))
using_ida7api = (major > 6)
using_pyqt5 = using_ida7api or (major == 6 and minor >= 9)

idaver_74newer = (major == 7 and minor >= 4)
idaver_8newer = (major >= 8)

if idaver_74newer or idaver_8newer:
    newer_version_compatible = True
else:
    newer_version_compatible = False

IDA_9 = (major >= 9)

if newer_version_compatible:
    # IDA 7.4+
    # https://hex-rays.com/products/ida/support/ida74_idapython_no_bc695_porting_guide.shtml
    import ida_ida
    import ida_kernwin

if IDA_9:
    BITS = 32 if not ida_ida.inf_is_64bit() else 64
    BWN_DISASM = ida_kernwin.BWN_DISASM
    BWN_PSEUDOCODE = ida_kernwin.BWN_PSEUDOCODE
    SETMENU_APP = ida_kernwin.SETMENU_APP
else:
    BITS = 32 if not idaapi.get_inf_structure().is_64bit() else 64
    BWN_DISASM = idaapi.BWN_DISASMS
    BWN_PSEUDOCODE = idaapi.BWN_PSEUDOCODE
    SETMENU_APP = idaapi.SETMENU_APP

if using_pyqt5:
    import PyQt5.QtGui as QtGui
    import PyQt5.QtCore as QtCore
    import PyQt5.QtWidgets as QtWidgets
    from PyQt5.Qt import QApplication

else:
    import PySide.QtGui as QtGui
    import PySide.QtCore as QtCore
    QtWidgets = QtGui
    QtCore.pyqtSignal = QtCore.Signal
    QtCore.pyqtSlot = QtCore.Slot
    from PySide.QtGui import QApplication


def setClipboardText(data):
    cb = QApplication.clipboard()
    cb.clear(mode=cb.Clipboard)
    cb.setText(data, mode=cb.Clipboard)


def PLUGIN_ENTRY():
    """
    Required plugin entry point for IDAPython Plugins.
    """
    return cvutils_getoffset()


class cvutils_getoffset(idaapi.plugin_t):

    flags = idaapi.PLUGIN_PROC | idaapi.PLUGIN_HIDE
    comment = "Get the Address offset at the cursor location."
    help = "At a certain location right-click 'Get Offset'"
    wanted_name = PLUGIN_NAME
    wanted_hotkey = PLUGIN_HOTKEY

    #--------------------------------------------------------------------------
    # Plugin Overloads
    #--------------------------------------------------------------------------

    def init(self):
        """
        This is called by IDA when it is loading the plugin.
        """

        # initialize the menu actions our plugin will inject
        self._init_action_get_offset()

        # initialize plugin hooks
        self._init_hooks()

        # done
        idaapi.msg("%s %s initialized...\n" % (self.wanted_name, VERSION))
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        """
        This is called by IDA when this file is loaded as a script.
        """
        idaapi.msg("%s cannot be run as a script.\n" % self.wanted_name)

    def term(self):
        """
        This is called by IDA when it is unloading the plugin.
        """

        # unhook our plugin hooks
        self._hooks.unhook()

        # unregister our actions & free their resources
        self._del_action_get_offset()

        # done
        idaapi.msg("%s terminated...\n" % self.wanted_name)

    #--------------------------------------------------------------------------
    # Plugin Hooks
    #--------------------------------------------------------------------------

    def _init_hooks(self):
        """
        Install plugin hooks into IDA.
        """
        self._hooks = Hooks()
        self._hooks.ready_to_run = self._init_hexrays_hooks
        self._hooks.hook()

    def _init_hexrays_hooks(self):
        """
        Install Hex-Rays hooks (when available).
        NOTE: This is called when the ui_ready_to_run event fires.
        """
        if idaapi.init_hexrays_plugin():
            idaapi.install_hexrays_callback(self._hooks.hxe_callback)

    #--------------------------------------------------------------------------
    # IDA Actions
    #--------------------------------------------------------------------------

    ACTION_GET_OFFSET = "prefix:get_offset"

    def _init_action_get_offset(self):
        """
        Register the get offset action with IDA.
        """
        # If the action is already registered, unregister it first.
        if idaapi.unregister_action(self.ACTION_GET_OFFSET):
            idaapi.msg("Warning: action was already registered, unregistering it first\n")

        vaction_desc = "Get the offset from the image base of the current cursor address."
        action_desc = idaapi.action_desc_t(
            self.ACTION_GET_OFFSET,  # The action name.
            "Get Offset",  # The action text.
            IDACtxEntry(getcopy_offset),  # The action handler.
            PLUGIN_HOTKEY,  # Optional: action shortcut
            vaction_desc,  # Optional: tooltip
            31  # Copy icon
        )

        # register the action with IDA
        assert idaapi.register_action(action_desc), "Action registration failed"

    def _del_action_get_offset(self):
        """
        Delete the bulk prefix action from IDA.
        """
        idaapi.unregister_action(self.ACTION_GET_OFFSET)


#------------------------------------------------------------------------------
# Plugin Hooks
#------------------------------------------------------------------------------

class Hooks(idaapi.UI_Hooks):

    def __init__(self):
        # Call the __init__ method of the superclass
        super(Hooks, self).__init__()

        # Get the IDA version
        major, minor = map(int, idaapi.get_kernel_version().split("."))
        self.newer_version_compatible = (major == 7 and minor >= 4) or (major >= 8)

        # If the IDA version is less than 7.4, define finish_populating_tform_popup
        if not self.newer_version_compatible:
            self.finish_populating_tform_popup = self._finish_populating_tform_popup

    def finish_populating_widget_popup(self, widget, popup_handle, ctx=None):
        """
        A right click menu is about to be shown. (IDA 7.x)
        """
        widget_type = idaapi.get_widget_type(widget)
        if widget_type is not None:
            inject_address_offset_copy_actions(widget, popup_handle, widget_type)
        return 0

    def _finish_populating_tform_popup(self, form, popup):
        """
        A right click menu is about to be shown. (IDA 6.x)
        """
        inject_address_offset_copy_actions(form, popup, idaapi.get_tform_type(form))
        return 0

    def hxe_callback(self, event, *args):
        """
        Hex-Rays event callback.
        """
        if event == idaapi.hxe_populating_popup:
            form, popup, vu = args

            idaapi.attach_action_to_popup(
                form,
                popup,
                cvutils_getoffset.ACTION_GET_OFFSET,
                "Get Address Offset",
                SETMENU_APP,
            )

        # done
        return 0


#------------------------------------------------------------------------------
# Prefix Wrappers
#------------------------------------------------------------------------------

def inject_address_offset_copy_actions(widget, popup_handle, widget_type):
    try:
        if widget_type in [BWN_DISASM, BWN_PSEUDOCODE]:
            idaapi.attach_action_to_popup(
                widget,
                popup_handle,
                cvutils_getoffset.ACTION_GET_OFFSET,
                "Get Address Offset",
                SETMENU_APP
            )
    except AttributeError:
        # Handle older versions where the widget type constants might differ
        pass
    return 0


#------------------------------------------------------------------------------
# Get Screen linear address
#------------------------------------------------------------------------------

def get_screen_linear_address():
    if newer_version_compatible:
        return idc.get_screen_ea()
    else:
        return idc.ScreenEA()


#------------------------------------------------------------------------------
# Get Offset
#------------------------------------------------------------------------------

def getcopy_offset():
    """
    Gets the offset of the current cursor's address
    """
    vImagebase = idaapi.get_imagebase()
    vCurrentPos = get_screen_linear_address()
    if vCurrentPos != idaapi.BADADDR:
        vOffset = vCurrentPos - vImagebase
        print("Address [0x%x] =-> Offset [%x] Copied to Clipboard!" % (vCurrentPos, vOffset))
        setClipboardText("%x" % vOffset)
    return


#------------------------------------------------------------------------------
# IDA ctxt
#------------------------------------------------------------------------------

class IDACtxEntry(idaapi.action_handler_t):

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

    def activate(self, ctx):
        self.action_function()
        return 1

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


#------------------------------------------------------------------------------
# Utilities
#------------------------------------------------------------------------------

PLUGIN_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), "plugin"))

def plugin_resource(resource_name):
    """
    Return the full path for a given plugin resource file.
    """
    return os.path.join(PLUGIN_PATH, "resources", resource_name)


repository: handsomematt/ida_func_ptr
file: ida_func_ptr.py
description: Easily create and copy function pointers to functions in IDA.
wanted_name: funcref
comment: A plugin for easy function references
created_at: 2017-07-26 22:29:45+00:00
pushed_at: 2017-07-27 20:04:12+00:00
url:
https://github.com/handsomematt/ida_func_ptr/blob/3c196293ccfa900df0ba4c1fce42f2cc59f629fa/ida_func_ptr.py

content:
import os

import idc
import idaapi
import idautils

import PySide.QtGui as QtGui
import PySide.QtCore as QtCore
QtWidgets = QtGui
QtCore.pyqtSignal = QtCore.Signal
QtCore.pyqtSlot = QtCore.Slot

#------------------------------------------------------------------------------
# IDA Plugin
#------------------------------------------------------------------------------

VERSION = "v1.0"
AUTHORS = ['Matt Stevens']

def PLUGIN_ENTRY():
    """
    Required plugin entry point for IDAPython Plugins.
    """
    return funcref_t()

class funcref_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_PROC | idaapi.PLUGIN_HIDE
    help = ""
    comment = "A plugin for easy function references"
    wanted_name = "funcref"
    wanted_hotkey = ""

    #--------------------------------------------------------------------------
    # Plugin Overloads
    #--------------------------------------------------------------------------

    def init(self):
        # just go when we have hexrays
        if not idaapi.init_hexrays_plugin():
            return idaapi.PLUGIN_SKIP

        # initialize the menu actions our plugin will inject
        self._init_action_bulk()
        self._init_action_copy()

        # initialize plugin hooks
        self._init_hooks()

        # done
        idaapi.msg("%s %s initialized...\n" % (self.wanted_name, VERSION))
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        idaapi.msg("%s cannot be run as a script.\n" % self.wanted_name)

    def term(self):
        if not hasattr(self, '_hooks'):
            return

        # unhook our plugin hooks
        self._hooks.unhook()

        # unregister our actions & free their resources
        self._del_action_bulk()
        self._del_action_copy()

        # done
        idaapi.msg("%s terminated...\n" % self.wanted_name)

    #--------------------------------------------------------------------------
    # Plugin Hooks
    #--------------------------------------------------------------------------

    def _init_hooks(self):
        self._hooks = Hooks()
        self._hooks.hook()

        if not idaapi.init_hexrays_plugin():
            idaapi.msg("[ERROR] Failed to initialize Hex-Rays SDK")
        else:
            idaapi.install_hexrays_callback(self._hooks.hxe_callback)

    #--------------------------------------------------------------------------
    # IDA Actions
    #--------------------------------------------------------------------------

    ACTION_BULK = "funcref:bulk"
    ACTION_COPY = "funcref:copy"

    def _init_action_bulk(self):
        """
        Register the bulk prefix action with IDA.
        """

        icon_data = "".join([
                "\x89\x50\x4E\x47\x0D\x0A\x1A\x0A\x00\x00\x00\x0D\x49\x48\x44\x52\x00\x00\x00\x10\x00\x00\x00\x10\x08\x06\x00\x00\x00\x1F\xF3\xFF\x61\x00\x00\x02\xCA\x49\x44\x41\x54\x78\x5E\x65",
                "\x53\x6D\x48\x53\x6F\x14\x3F\xBA\xB5\xB7\xA0\x8D\x20\x41\xF2\xBA\x5D\xB6\x0F\x56\xF4\x41\xA2\xC0\x9C\xE9\xB4\x29\x4A\x7D\xB0\x22\x7A\x11\x02\x23\x48\x2A\xD4\x74\x53\x33\x3F\xD4",
                "\x3E\x4A\x50\x19\xE4\xB0\xD0\x22\xCD\x44\x45\x4A\x31\x8C\x92\xA2\x3E\x65\x0A\x4D\xCB\x96\x7E\xE8\xD5\x97\xCC\xFE\xFE\x37\xA7\x77\xDB\xBD\xA7\xE7\x3C\xBE\x05\x9E\xED\xB7\xB3\xF3",
                "\x7B\x39\xF7\xEE\x19\x17\xA8\xAC\x56\xDB\x54\x82\x60\x41\xB3\x59\xBC\xFF\xAC\xF9\xCA\xB5\xAE\x86\xCA\xF9\x4E\xAF\x1B\x3B\xEA\x5D\x48\x9D\x66\xE2\x49\x27\x9F\xD5\x66\x9B\xA2\x1C",
                "\x22\x02\xD0\x40\xE4\x81\x6C\x3B\x76\x37\x56\xE3\x37\x5F\x2F\x62\xE8\x0B\xD3\x66\x19\x7E\x53\xA7\x99\x78\xAE\x1F\x64\x3E\x21\x71\x69\x09\x5F\x20\x98\x2D\x58\x70\x24\x07\x07\x7B",
                "\x6F\xB0\x79\x82\x61\x81\x21\xCC\xDE\x21\x54\x16\x02\xD4\x69\x26\x9E\x74\xEE\xCB\xCF\x4D\xC7\x44\xB3\x88\x7C\x81\xC5\x22\xFE\x6C\xB9\xE9\x46\x67\x46\x1A\x8A\x16\x2B\x0A\x5B\x05",
                "\x74\x66\x65\xE1\x98\x6F\x00\x31\x32\x87\x9F\x59\x77\x66\x66\x61\x42\xBC\xC0\xF5\x6C\x47\x1A\x36\xD7\xB9\x51\x14\xC5\x1E\xBE\xA0\xC3\x5B\xD9\x98\x99\xE1\xC0\xCE\xBE\x57\x48\xD7",
                "\x9A\x63\x68\xEA\x7C\x8A\xF6\x14\x3B\x9F\xF6\xA6\xA4\x60\xEB\xE3\x3E\x9C\x5F\xD6\x5A\x7A\xFA\x71\xBF\xC3\x81\x3D\x4D\x35\x0D\x7C\xC1\xF3\x87\x57\x43\xF9\x87\x8F\x21\x95\x5E\xAB",
                "\x41\x83\x4E\x83\x54\xDB\x92\x76\x20\xCA\xBF\xD0\x99\x9D\xBB\x4E\xDB\xBD\xC7\x8E\x2F\x5A\x3D\x74\x3D\x50\x03\x80\x7E\x7A\x7A\x06\x46\x47\xFD\xA0\x33\x6C\x84\x18\x46\x0C\xBD\x1F",
                "\x86\x2D\x71\x71\x00\x52\x10\x16\x17\xE6\xC1\xE7\x1B\x61\x9A\x81\x69\x31\x30\xFC\x61\x14\xB4\x3A\x3D\x20\x82\x1E\x58\xA9\x15\x05\x41\x14\x05\xB8\x58\xEE\x82\x7D\xE9\x99\x20\xCB",
                "\x32\x94\x95\x95\xC3\xA5\xD2\x53\x00\x51\x09\xAA\x4B\x0B\xA1\xB8\xA4\x0C\x52\x53\x33\x40\xA5\x52\x81\xDB\x5D\x01\xA2\x45\x00\x45\x51\x80\x2A\x36\x12\x8D\x42\x49\x51\x01\x44\xE5",
                "\x18\x90\x22\x0A\x98\x8C\x46\xF0\x54\x14\x42\x6D\x7D\x3B\xE4\x1C\x75\x41\xAD\xB7\x1D\x3C\x55\x85\x60\x32\x19\x41\x8A\x2A\xDC\x57\x5C\x74\x12\x28\x47\xA5\x8E\x44\xE4\xF0\x76\x5B",
                "\x82\xA6\xCD\x5B\x0D\xB2\x12\xE6\xE4\x06\xB5\x1A\x66\xA7\x26\x41\x92\xC2\xA0\xD5\x6A\x60\x67\x92\x19\xAE\x7B\xCE\x70\x4D\x15\xAB\x01\xAD\xC1\x08\x3F\x46\x64\x6E\x8E\x9D\xF9\x13",
                "\xE8\x1A\xFF\xE4\x63\x8A\x0E\xE6\x02\x41\xF8\x3F\x18\x82\x40\x28\x04\xFD\xDD\x75\xF0\xB6\xFF\x2E\x75\x9A\x89\x27\x9D\xFB\xC8\x4F\x39\xBE\xE0\xB4\xAB\xCE\x35\xFE\x71\x00\x16\x17",
                "\x25\x76\x50\x26\x76\x6B\x61\x86\x08\xE4\x1D\xAF\x81\xBC\x13\x97\xA9\xD3\x4C\x3C\xE9\xDC\x47\x7E\xCA\xF1\x05\x0C\x5F\x7D\xFE\xEF\x35\x03\xAF\x9F\x00\xB0\x73\x30\x9A\xE2\x81\x0E",
                "\xF6\xC1\xED\x52\xB8\x77\xAB\x98\x3A\xCD\xC4\x73\x9D\x7C\x6F\xDE\xF9\xCF\x53\x0E\xFE\xA9\xCD\xAE\xB3\x87\xCE\x75\x35\x54\xE1\xD0\xCB\x47\x38\x39\x36\x88\xFF\x4D\xF8\x57\x41\x33",
                "\xF1\xA4\x93\x0F\x00\x36\xAD\x3E\x4C\x6B\xC5\xC9\x5D\x77\x6A\x2F\xB4\x31\xA3\xC4\x40\x4F\x21\x0F\xD1\x4C\x3C\xE9\x2B\xE1\xF5\x0B\xD6\x90\xC8\x90\x4C\xE6\x35\xD0\xCC\x79\x5E\xFF",
                "\x2E\xF8\x0B\x2F\x3D\xE5\xC3\x97\x06\xCF\xCF\x00\x00\x00\x00\x49\x45\x4E\x44\xAE\x42\x60\x82"])

        # load the icon for this action
        self._bulk_icon_id = idaapi.load_custom_icon(data=icon_data, format="png")

        # describe the action
        action_desc = idaapi.action_desc_t(
            self.ACTION_BULK,                                        # The action name.
            "Copy function pointers to selected function(s)",        # The action text.
            IDACtxEntry(bulk_function),                              # The action handler.
            None,                                                    # Optional: action shortcut
            "Copies a function pointer to the selected function(s)", # Optional: tooltip
            self._bulk_icon_id                                       # Optional: the action icon
        )

        # register the action with IDA
        assert idaapi.register_action(action_desc), "Action registration failed"


    def _init_action_copy(self):
        # load the icon for this action
        self._copy_icon_id = 199

        # describe the action
        action_desc = idaapi.action_desc_t(
            self.ACTION_COPY,                  # The action name.
            "Copy function reference",         # The action text.
            IDACtxEntry(copy_function_cursor), # The action handler.
            None,                              # Optional: action shortcut
            "Copy reference of this function", # Optional: tooltip
            self._copy_icon_id                 # Optional: the action icon
        )

        # register the action with IDA
        assert idaapi.register_action(action_desc), "Action registration failed"

    def _del_action_bulk(self):
        idaapi.unregister_action(self.ACTION_BULK)
        idaapi.free_custom_icon(self._bulk_icon_id)
        self._bulk_icon_id = idaapi.BADADDR

    def _del_action_copy(self):
        idaapi.unregister_action(self.ACTION_COPY)
        #idaapi.free_custom_icon(self._copy_icon_id)
        self._copy_icon_id = idaapi.BADADDR

#------------------------------------------------------------------------------
# Plugin Hooks
#------------------------------------------------------------------------------

class Hooks(idaapi.UI_Hooks):

    def finish_populating_tform_popup(self, form, popup):
        # disassembly window
        if idaapi.get_tform_type(form) == idaapi.BWN_DISASMS:
            if get_cursor_func_ref() == idaapi.BADADDR:
                return

            idaapi.attach_action_to_popup(
                form,
                popup,
                funcref_t.ACTION_COPY
            )

        # functions window
        elif idaapi.get_tform_type(form) == idaapi.BWN_FUNCS:
            idaapi.attach_action_to_popup(form, popup, funcref_t.ACTION_BULK, "Copy All", idaapi.SETMENU_INS)

        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,
                funcref_t.ACTION_COPY,
                "Rename global item",
                idaapi.SETMENU_APP
            )

        return 0


def copy_to_clip(data):
    QtGui.QApplication.clipboard().setText(data)

def copy_function_cursor():
    # get the function reference under the user cursor (if there is one)
    target = get_cursor_func_ref()
    if target == idaapi.BADADDR:
        return

    # execute the recursive prefix
    functionPtrDef = copy_function(target)
    copy_to_clip(functionPtrDef)
    idaapi.msg(functionPtrDef + "\n")

def copy_function(addr):

    func_addr = idc.LocByName(idaapi.get_func_name(addr))
    if func_addr == idaapi.BADADDR:
        idaapi.msg("0x%08X does not belong to a defined function\n" % addr)
        return

    callTypes = ["__cdecl", "__fastcall", "__stdcall", "__thiscall", "__usercall"]

    funcDef = str(idaapi.decompile(func_addr)).split('\n', 1)[0]

    hasCallType = any(call_type in funcDef for call_type in callTypes)
    parenthesesStart = funcDef.find('(')
    parenthesesEnd = funcDef.rfind(')')
    funcNameStart = funcDef[0 : parenthesesStart].rfind(' ')
    funcNameEnd = parenthesesStart
    returnTypeStart = 0
    returnTypeEnd = funcNameStart

    callType = ""
    if hasCallType:
        callTypeStart = funcDef[0 : funcNameStart].rfind(' ')
        callType = funcDef[callTypeStart + 1 : funcNameStart]
        returnTypeEnd = callTypeStart
        if callType == "__cdecl":
            callType = ""

    returnType = funcDef[returnTypeStart : returnTypeEnd]
    funcName = funcDef[funcNameStart + 1 : parenthesesStart]
    args = funcDef[parenthesesStart + 1 : parenthesesEnd]

    finalString = "{0} ({1}* {2})({3}) = ({0}({1}*)({3}))({4});".format(
        returnType, callType, funcName, args, "0x%08X" % func_addr
    )

    return finalString

def bulk_function():
    functionPtrs = ""
    for func_name in get_selected_funcs():
        functionPtrDef = copy_function(idc.LocByName(func_name))
        functionPtrs = functionPtrs + functionPtrDef + "\n"
    copy_to_clip(functionPtrs)
    idaapi.msg(functionPtrs)

def get_all_funcs():
    return set(idaapi.get_func_name(ea) for ea in idautils.Functions())

def get_cursor_func_ref():
    current_tform  = idaapi.get_current_tform()
    tform_type     = idaapi.get_tform_type(current_tform)

    # get the hexrays vdui (if available)
    vu = idaapi.get_tform_vdui(current_tform)
    if vu:
        cursor_addr = vu.item.get_ea()
    elif tform_type == idaapi.BWN_DISASM:
        cursor_addr = idaapi.get_screen_ea()

        op_addr = idc.GetOperandValue(cursor_addr, idaapi.get_opnum())
        op_func = idaapi.get_func(op_addr)
        if op_func and op_func.startEA == op_addr:
            return op_addr

    else:
        return idaapi.BADADDR

    cursor_func = idaapi.get_func(cursor_addr)
    if cursor_func and cursor_func.startEA == cursor_addr:
        return cursor_addr

    return idaapi.BADADDR

def get_selected_funcs():
    tform = idaapi.find_tform("Functions window")
    if not tform:
        idc.Warning("Unable to find 'Functions window'")
        return

    widget = idaapi.PluginForm.FormToPySideWidget(tform)
    table = widget.findChild(QtWidgets.QTableView)
    selected_funcs = [str(s.data()) for s in table.selectionModel().selectedRows()]

    return match_funcs(selected_funcs)

def match_funcs(qt_funcs):
    res = set()
    ida_funcs = get_all_funcs()
    for f in qt_funcs:
        for f2 in ida_funcs:
            if len(f) == len(f2):
                i = 0
                while i < len(f) and (f[i] == f2[i] or f[i] == '_'):
                    i += 1

                if i == len(f):
                    res.add(f2)
                    break

    return list(res)

def graph_down(ea, path=set()):
    """
    Recursively collect all function calls.

    Copied with minor modifications from
    http://hooked-on-mnemonics.blogspot.com/2012/07/renaming-subroutine-blocks-and.html
    """
    path.add(ea)

    #
    # iterate through all the instructions in the target function (ea) and
    # inspect all the call instructions
    #

    for x in [x for x in idautils.FuncItems(ea) if idaapi.is_call_insn(x)]:

        #  TODO
        for r in idautils.XrefsFrom(x, idaapi.XREF_FAR):
            #print "0x%08X" % h, "--calls-->", "0x%08X" % r.to
            if not r.iscode:
                    continue

            # get the function pointed at by this call
            func = idaapi.get_func(r.to)
            if not func:
                continue

            # ignore calls to imports / library calls / thunks
            if (func.flags & (idaapi.FUNC_THUNK | idaapi.FUNC_LIB)) != 0:
                continue

            #
            # if we have not traversed to the destination function that this
            # call references, recurse down to it to continue our traversal
            #

            if r.to not in path:
                graph_down(r.to, path)

    return path

class IDACtxEntry(idaapi.action_handler_t):
    """
    A basic Context Menu class to utilize IDA's action handlers.
    """

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

    def activate(self, ctx):
        """
        Execute the embedded action_function when this context menu is invoked.
        """
        self.action_function()
        return 1

    def update(self, ctx):
        """
        Ensure the context menu is always available in IDA.
        """
        return idaapi.AST_ENABLE_ALWAYS


repository: yubie-re/ida-jm-xorstr-decrypt-plugin
file: xorstr_decrypt.py
description: Attempts to decrypt JM Xorstr in some x64 binaries
wanted_name: Xorstring Decryptor
comment: Attempts to detect & decrypt JM Xorstring
created_at: 2023-02-23 01:29:46+00:00
pushed_at: 2023-03-09 18:02:08+00:00
url:
https://github.com/yubie-re/ida-jm-xorstr-decrypt-plugin/blob/d79482c69f397b81467f39459fc6738f034b560a/xorstr_decrypt.py

content:
import sys
import ida_allins
import ida_bytes
import ida_idaapi
import ida_kernwin
import ida_search
import ida_ua
import idaapi
import idc


class xor_decryption_mod(ida_idaapi.plugmod_t):
    stack_count = 0

    def __del__(self):
        ida_kernwin.msg("unloaded xor decryptor\n")

    """
    Returns the instruction at a linear address
    """
    def get_insn(self, ea: int):
        insn = idaapi.insn_t()
        idaapi.decode_insn(insn, ea)
        return insn

    """
    Returns the previous instruction
    """
    def get_previous_insn(self, ea):
        insn = idaapi.insn_t()
        idaapi.decode_prev_insn(insn, ea)
        return insn

    """
    Returns the next instruction, or None if it can't find any
    """
    def get_next_insn(self, previous_insn):
        insn = idaapi.insn_t()
        if previous_insn.size == 0:
            return None
        idaapi.decode_insn(insn, previous_insn.ea + previous_insn.size)
        return insn

    """
    Finds where the initial mov is for the key or data. This matches the stack address and checks that it is being written to with a mov.
    This moves backwards.
    returns the instruction where the first mov is for the data or key.

    movabs rax, -4762152789334367252
    >> RETURN HERE FOR DATA << mov QWORD PTR [rsp], rax
    movabs rax, -6534519754492314190
    mov QWORD PTR [rsp+8], rax
    movabs rax, -2862143164529545214
    mov QWORD PTR [rsp+16], rax
    movabs rax, -4140208776682645948
    mov QWORD PTR [rsp+24], rax
    vmovdqa ymm1, YMMWORD PTR [rsp]
    movabs rax, -2550414817236710003
    >> RETURN HERE FOR KEY << mov QWORD PTR [rsp+32], rax
    movabs rax, -4595755740016602734
    mov QWORD PTR [rsp+40], rax
    movabs rax, -5461194525092864914
    mov QWORD PTR [rsp+48], rax
    movabs rax, -4140208776682645984
    mov QWORD PTR [rsp+56], rax
    vpxor ymm0, ymm1, YMMWORD PTR [rsp+32]
    vmovdqa YMMWORD PTR [rsp], ymm0
    """
    def find_stack_push_start(self, insn, stackaddr):
        calls = 0
        while insn.itype != ida_allins.NN_mov or (insn.ops[0].addr != stackaddr) or insn.ops[0].type != ida_ua.o_displ or insn.ops[1].type != ida_ua.o_reg and calls < 5000:
            calls += 1
            if insn.ea == idaapi.SIZE_MAX:
                return None
            insn = self.get_previous_insn(insn.ea)
        if calls == 5000:
            return None
        return insn

    """
    Finds the last findable immediate value for known for a register by moving backwards until finding a mov instruction where the register is written to
    This moves backwards.
    returns the immediate value of a register

    movabs rax, >> RETURNS THIS: -6534519754492314190 <<
    mov QWORD PTR [rsp+8], rax
    movabs rax, -2862143164529545214
    mov QWORD PTR [rsp+16], rax
    movabs rax, -4140208776682645948
    mov QWORD PTR [rsp+24], rax
    vmovdqa ymm1, YMMWORD PTR [rsp]
    movabs rax, -2550414817236710003
    mov QWORD PTR [rsp+32], rax
    movabs rax, -4595755740016602734
    mov QWORD PTR [rsp+40], rax
    movabs rax, -5461194525092864914
    mov QWORD PTR [rsp+48], rax
    movabs rax, -4140208776682645984
    mov QWORD PTR [rsp+56], rax
    vpxor ymm0, ymm1, YMMWORD PTR [rsp+32]
    vmovdqa YMMWORD PTR [rsp], ymm0
    """
    def find_register_value(self, insn, reg):
        calls = 0
        while insn.itype != ida_allins.NN_mov or (insn.ops[0].type != ida_ua.o_reg) or (insn.ops[0].reg != reg) and calls < 1000:
            if insn.ea == idaapi.SIZE_MAX:
                return None
            insn = self.get_previous_insn(insn.ea)
        if calls == 1000:
            return None
        if (insn.ops[1].type != ida_ua.o_imm):
            stack_insn = self.find_stack_push_start(
                self.get_previous_insn(insn.ea), insn.ops[1].addr)
            if stack_insn == None:
                return None
            return self.find_register_value(stack_insn, stack_insn.ops[1].reg)
        return insn.ops[1].value

    """
    Used to find what stack address is moved into the xmm/ymm register later used in the pxor instructions
    This moves backwards.
    returns the movdqx instruction

    movabs rax, -6534519754492314190
    mov QWORD PTR [rsp+8], rax
    movabs rax, -2862143164529545214
    mov QWORD PTR [rsp+16], rax
    movabs rax, -4140208776682645948
    mov QWORD PTR [rsp+24], rax
    RETURN HERE >> vmovdqa ymm1, YMMWORD PTR [rsp]
    movabs rax, -2550414817236710003
    mov QWORD PTR [rsp+32], rax
    movabs rax, -4595755740016602734
    mov QWORD PTR [rsp+40], rax
    movabs rax, -5461194525092864914
    mov QWORD PTR [rsp+48], rax
    movabs rax, -4140208776682645984
    mov QWORD PTR [rsp+56], rax
    vpxor ymm0, ymm1, YMMWORD PTR [rsp+32]
    vmovdqa YMMWORD PTR [rsp], ymm0
    """
    def find_register_movdq_insn(self, insn, reg):
        calls = 0
        while (insn.itype != ida_allins.NN_vmovdqa and insn.itype != ida_allins.NN_vmovdqu and insn.itype != ida_allins.NN_movdqa and insn.itype != ida_allins.NN_movdqu) or (insn.ops[0].type != ida_ua.o_reg) or (insn.ops[0].reg != reg) and calls < 1000:
            if insn.ea == idaapi.SIZE_MAX:
                return None
            insn = self.get_previous_insn(insn.ea)
        if calls == 1000:
            return None
        return insn

    """
    Used to find where the ymm/xmm xored output is moved back onto the stack (useful to find where to place psuedocode comments)
    This moves forwards
    returns the movdqx instruction where this happens

    movabs rax, -6534519754492314190
    mov QWORD PTR [rsp+8], rax
    movabs rax, -2862143164529545214
    mov QWORD PTR [rsp+16], rax
    movabs rax, -4140208776682645948
    mov QWORD PTR [rsp+24], rax
    vmovdqa ymm1, YMMWORD PTR [rsp]
    movabs rax, -2550414817236710003
    mov QWORD PTR [rsp+32], rax
    movabs rax, -4595755740016602734
    mov QWORD PTR [rsp+40], rax
    movabs rax, -5461194525092864914
    mov QWORD PTR [rsp+48], rax
    movabs rax, -4140208776682645984
    mov QWORD PTR [rsp+56], rax
    vpxor ymm0, ymm1, YMMWORD PTR [rsp+32]
     >> RETURNS HERE << vmovdqa YMMWORD PTR [rsp], ymm0
    """
    def find_stack_movdq_insn(self, insn, reg):
        calls = 0
        while (insn.itype != ida_allins.NN_vmovdqa and insn.itype != ida_allins.NN_vmovdqu and insn.itype != ida_allins.NN_movdqa and insn.itype != ida_allins.NN_movdqu) or (insn.ops[1].type != ida_ua.o_reg) or (insn.ops[1].reg != reg) and calls < 1000:
            if insn.ea == idaapi.SIZE_MAX:
                return None
            insn = self.get_next_insn(insn)
            if insn == None:
                return None
        if calls == 1000:
            return None
        return insn

    """
    Finds the next stack push instruction which matches the stack address given.
    This moves forwards
    returns the mov instruction which accesses the address

    movabs rax, -6534519754492314190
    >> CALLED HERE << mov QWORD PTR [rsp+8], rax
    movabs rax, -2862143164529545214
    >> RETURNS HERE << mov QWORD PTR [rsp+16], rax
    movabs rax, -4140208776682645948
    mov QWORD PTR [rsp+24], rax
    vmovdqa ymm1, YMMWORD PTR [rsp]
    movabs rax, -2550414817236710003
    mov QWORD PTR [rsp+32], rax
    movabs rax, -4595755740016602734
    mov QWORD PTR [rsp+40], rax
    movabs rax, -5461194525092864914
    mov QWORD PTR [rsp+48], rax
    movabs rax, -4140208776682645984
    mov QWORD PTR [rsp+56], rax
    vpxor ymm0, ymm1, YMMWORD PTR [rsp+32]
    vmovdqa YMMWORD PTR [rsp], ymm0
    """
    def find_next_stack_push(self, insn, address):
        calls = 0
        while insn.itype != ida_allins.NN_mov or (insn.ops[0].addr != address) and calls < 5000:
            calls += 1
            if insn.ea == idaapi.SIZE_MAX:
                return None
            insn = self.get_next_insn(insn)
            if insn == None:
                return None
        if calls == 5000:
            return None
        return insn

    """
    Handles a basic xor cipher with two byte arrays
    """
    def byte_xor(self, ba1, ba2):
        return bytes([_a ^ _b for _a, _b in zip(ba1, ba2)])

    """
    Handles the string decryption.
    Steps:
    Find where it starts pushing data onto the stack
    Figure if we're dealing with xmm/ymm registers
    Find immediate values of the data and append them to a byte array
    Find where it starts pushing the key onto the stack
    Find immediate values of the key and append them to a byte array
    Xor the two arrays, then set comments as necessary
    """
    def handle_str_decryption(self, data_reg, key_address, func_addr, pxor_insn):
        previous_insn = self.find_register_movdq_insn(pxor_insn, data_reg)
        if previous_insn == None:
            return None
        data_address = previous_insn.ops[1].addr
        mov_start = self.find_stack_push_start(previous_insn, data_address)
        if mov_start == None:
            return None
        if idaapi.get_reg_name(data_reg, 16).startswith('xmm'):
            expected_pushes = 2
        elif idaapi.get_reg_name(data_reg, 16).startswith('ymm'):
            expected_pushes = 4
        else:
            return None
        mov_insn = mov_start
        xor_data = bytes()
        xor_key = bytes()
        for x in range(0, expected_pushes):
            register_val = self.find_register_value(
                mov_insn, mov_insn.ops[1].reg)
            if register_val == None:
                return None
            xor_data += (register_val.to_bytes(8, sys.byteorder))
            if x != expected_pushes - 1:
                mov_insn = self.find_next_stack_push(
                    mov_insn, data_address + (x + 1) * 8)
            if mov_insn == None:
                return None
        mov_insn = self.find_stack_push_start(previous_insn, key_address)
        if mov_insn == None:
            return None
        for x in range(0, expected_pushes):
            register_val = self.find_register_value(
                mov_insn, mov_insn.ops[1].reg)
            if register_val == None:
                return None
            xor_key += (register_val.to_bytes(8, sys.byteorder))
            if x != expected_pushes - 1:
                mov_insn = self.find_next_stack_push(
                    mov_insn, key_address + (x + 1) * 8)
            if mov_insn == None:
                return None
            
        decrypted = self.byte_xor(xor_data, xor_key)
        comment = 'Decrypted: '
        result = ''
        
        try:
            result = decrypted.decode('utf-8').rstrip('\x00')
        except:
            result = ' '.join(["{:02x}".format(b) for b in decrypted])
            comment = 'Raw: '
        
        if len(result) == 0:
            return None
        
        mov_to_stack_insn = self.find_stack_movdq_insn(pxor_insn, pxor_insn.ops[0].reg)
        idc.set_cmt(func_addr, comment + result, 0)
        cfunc = idaapi.decompile(mov_to_stack_insn.ea)
        if cfunc:
            tl = idaapi.treeloc_t()
            tl.ea = mov_to_stack_insn.ea
            tl.itp = idaapi.ITP_SEMI
            cfunc.set_user_cmt(tl, comment + result)
            cfunc.save_user_cmts()
        return result

    """
    Starts the routine for a PXOR instruction
    ex : pxor xmm0, [rbp+1F30h+var_1B90]
    """
    def handle_pxor(self, func_addr):
        insn = self.get_insn(func_addr)
        data_reg = insn.ops[0].reg
        key_address = insn.ops[1].addr
        return self.handle_str_decryption(data_reg, key_address, func_addr, insn)

    """
    Starts the routine for a VPXOR instruction
    ex : vpxor ymm0, ymm1, YMMWORD PTR [rsp+32]
    """
    def handle_vpxor(self, func_addr):
        insn = self.get_insn(func_addr)
        data_reg = insn.ops[1].reg
        key_address = insn.ops[2].addr
        return self.handle_str_decryption(data_reg, key_address, func_addr, insn)

    """
    calls the right routine depending on the instruction type
    """
    def analyze(self, func_addr):
        insn = self.get_insn(func_addr)
        if insn.itype == ida_allins.NN_vpxor:
            return self.handle_vpxor(func_addr)
        if insn.itype == ida_allins.NN_pxor:
            return self.handle_pxor(func_addr)
        return None

    """
    Analyzes all instances of an IDA Pattern with compability for IDA 7.5 (7.5 doesn't have compiled_binpat_vec_t, find_binary is deprecated in IDA 8)
    """
    def analyze_sig_75(self, sig):
        match_ea = idc.get_inf_attr(idc.INF_MIN_EA)
        while True:
            match_ea = ida_search.find_binary(
                match_ea + 1, ida_idaapi.BADADDR, sig, 16, idc.SEARCH_DOWN)
            if match_ea != idaapi.BADADDR:
                result = self.analyze(match_ea)
                if result != None:
                    print("Found match at {:08X} {}".format(match_ea, result))
            else:
                break

    """
    Analyzes all instances of an IDA Pattern
    """
    def analyze_sig(self, sig):
        match_ea = idc.get_inf_attr(idc.INF_MIN_EA)
        binpat = ida_bytes.compiled_binpat_vec_t()
        ida_bytes.parse_binpat_str(binpat, match_ea, sig, 16)
        while True:
            match_ea = ida_bytes.bin_search(
                match_ea + 1, idaapi.BADADDR, binpat, idaapi.BIN_SEARCH_FORWARD)
            if match_ea != idaapi.BADADDR:
                result = self.analyze(match_ea)
                if result != None:
                    print("Found match at {:08X} {}".format(match_ea, result))
            else:
                break

    """
    Starts plugin logic
    """
    def run(self, arg):
        if idaapi.IDA_SDK_VERSION <= 750:
            self.analyze_sig_75("C5 ? EF")  # vpxor
            self.analyze_sig_75("66 ? EF")  # pxor
        else:
            self.analyze_sig("C5 ? EF")  # vpxor
            self.analyze_sig("66 ? EF")  # pxor
        return 0

# This class is instantiated when IDA loads the plugin.
class xor_decryption_t(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_UNL | ida_idaapi.PLUGIN_MULTI
    comment = "Attempts to detect & decrypt JM Xorstring"
    help = "This is help"
    wanted_name = "Xorstring Decryptor"
    wanted_hotkey = "Alt-F8"

    # def __del__(self):
    # ida_kernwin.msg("unloaded globally\n")

    def init(self):
        ida_kernwin.msg("init() called!\n")
        return xor_decryption_mod()

    def run(self, arg):
        ida_kernwin.msg("ERROR: run() called for global object!\n")
        return 0

    def term(self):
        ida_kernwin.msg("ERROR: term() called (should never be called)\n")


def PLUGIN_ENTRY():
    return xor_decryption_t()


repository: vAlerainTech/vAlerain-Ark
file: cmake-build-debug/IDA_Pro_7.7/plugins/auto_re.py
wanted_name: Auto RE
created_at: 2024-03-25 14:26:36+00:00
pushed_at: 2024-07-27 05:07:46+00:00
description:
🔭This will be my first long-term update project, and vAlerainArk is committed to helping security personnel with testing.

url:
https://github.com/vAlerainTech/vAlerain-Ark/blob/df8bcafeda091d30ceade7bb2c7c73b5ae351a38/cmake-build-debug/IDA_Pro_7.7/plugins/auto_re.py

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

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


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


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


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


TAGS_IGNORE_LIST = {
    'OpenProcessToken',
    'DisconnectNamedPipe'
}

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

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

STRICT_TAG_NAME_CHECKING = {'file'}

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


def get_addr_width():
    return '16' if idaapi.cvar.inf.is_64bit() else '8'


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


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


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

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

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

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

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

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

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


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

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

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

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

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

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

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

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

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

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

        self.parent.setLayout(root_layout)

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

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

        self.tv.expandAll()

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

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

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

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

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

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

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

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

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

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

        return self._rename_ea_requested(addr, name_idx)

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

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

        if new_name is None:
            return

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

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

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

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

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

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

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

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

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

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

        return [ea_item, name_item, api_name]

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

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

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

            root_node.appendRow(item_tag_list)

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

                item_tag.appendRow(ref_item_list)

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

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


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

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

    _PREFIX_NAME = 'au_re_'
    _MIN_MAX_MATH_OPS_TO_ALLOW_RENAME = 10

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

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

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

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

        return idaapi.PLUGIN_OK

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

    # def _store_config(self, cfg):
    #     pass

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

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

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

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

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

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

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

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

        normalized = self.normalize_name(possible_name)

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

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

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

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

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

        return dis0.Op1.value

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

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

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

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

            if not addr:
                continue

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

            imp_stripped_name = name.lstrip('_')

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

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

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

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

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

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

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

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

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

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

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

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

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

    def term(self):
        self._data = None

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

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

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

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

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

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

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

        if not name:
            return

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

        if name in TAGS_IGNORE_LIST:
            return

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

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

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

    @classmethod
    def __get_ptr_val(cls, ea):
        if idaapi.cvar.inf.is_64bit():
            return idaapi.get_qword(ea)

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

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

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

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

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

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

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

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

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

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

        return rv

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

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


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


repository: grimdoomer/DeadRisingArcTool
file: DeadRisingArcTool/IDA/SSEHelper.py
description: Modding tools for Dead Rising .arc files
wanted_name: SSE Helper
comment: Provides high level comments for SSE instructions.
created_at: 2020-04-17 09:58:49+00:00
pushed_at: 2025-03-27 13:10:42+00:00
url:
https://github.com/grimdoomer/DeadRisingArcTool/blob/3415c31d7d06fd3b8d5f0fd592632e2d1cb13555/DeadRisingArcTool/IDA/SSEHelper.py

content:
﻿"""

"""

from idaapi import *
import idautils
from idc import *

"""
XMLoadFloat3(XMFLOAT3 *pSource)
	movss		xmm1, dword ptr [pSource.x]		// __m128 x = _mm_load_ss( &pSource->x );
	movss		xmm0, dword ptr [pSource.y]		// __m128 y = _mm_load_ss( &pSource->y );
	unpcklps	xmm1, xmm0						// __m128 xy = _mm_unpacklo_ps( x, y );
	movss		xmm0, dword ptr [pSource.z]		// __m128 z = _mm_load_ss( &pSource->z );
	movlhps		xmm1, xmm0						// return _mm_movelh_ps( xy, z );

"""

class SSEHelperPlugin(plugin_t):

	flags = idaapi.PLUGIN_PROC
	comment = "Provides high level comments for SSE instructions."
	help = "Provides high level comments for SSE instructions."
	wanted_name = "SSE Helper"
	wanted_hotkey = "Alt+F5"

	def init(self):

		#menu_item_desc = idaapi.action_desc_t("ssehelper:add_comment", "SSEHelper", self.add_comment, 
		#	"Alt+F5", "Adds a comment for the SSE instruction at the current address.", 122)

		# Add our menu item entry.
		#idaapi.register_action(menu_item_desc)
		#idaapi.attach_action_to_menu("Edit/Plugins/", "ssehelper:add_comment", idaapi.SETMENU_APP)

		return idaapi.PLUGIN_KEEP

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

	def term(self):
		pass

	def add_comment(self):

		comp = ['x', 'y', 'z', 'w']

		# Decode the instruction at the current address.
		ins = idautils.DecodeInstruction(ScreenEA())

		# Check if this instruction is supported and handle accordingly.
		print("Instruction at 0x%08x is %s" % (ScreenEA(), ins.get_canon_mnem()))
		if ins.get_canon_mnem() == "shufps":
			
			# Get register names.
			reg_dst = "xmm%s" % str(ins.Op1.reg - 64)
			reg_src = "xmm%s" % str(ins.Op2.reg - 64)

			# Get the shuffle indices.
			shuff_a = (ins.Op3.value & 0x3)
			shuff_b = (ins.Op3.value >> 2) & 0x3
			shuff_c = (ins.Op3.value >> 4) & 0x3
			shuff_d = (ins.Op3.value >> 6) & 0x3

			# Add the comment.
			MakeRptCmt(ScreenEA(), "%s = [%s.%c | %s.%c | %s.%c | %s.%c]" % 
				(reg_dst, reg_src, comp[shuff_d], reg_src, comp[shuff_c], reg_dst, comp[shuff_b], reg_dst, comp[shuff_a]))

		elif ins.get_canon_mnem() == "unpcklps":

			# Get register names.
			reg_dst = "xmm%c" % chr(ord('0') + ins.Op1.reg - 64)
			reg_src = "xmm%c" % chr(ord('0') + ins.Op2.reg - 64)

			# Add the comment.
			MakeRptCmt(ScreenEA(), "%s = [%s.y | %s.y | %s.x | %s.x]" % (reg_dst, reg_src, reg_dst, reg_src, reg_dst))

		elif ins.get_canon_mnem() == "movlhps":

			# Get register names.
			reg_dst = "xmm%c" % chr(ord('0') + ins.Op1.reg - 64)
			reg_src = "xmm%c" % chr(ord('0') + ins.Op2.reg - 64)

			# Add the comment.
			MakeRptCmt(ScreenEA(), "%s = [%s.y | %s.x | %s.y | %s.x]" % (reg_dst, reg_src, reg_src, reg_dst, reg_dst))


def PLUGIN_ENTRY():
	return SSEHelperPlugin()


repository: N3rdL0rd/EchoDecomp
file: utils/rizzo.py
description: A WIP decompilation of Echo Arena and the RAD15/14 engine
url: https://github.com/N3rdL0rd/EchoDecomp/blob/f84971cbda09872905c49ef9342027439fa024ff/utils/rizzo.py
wanted_name: Rizzo
comment: Function signature
created_at: 2024-05-30 14:42:29+00:00
pushed_at: 2024-11-10 23:30:16+00:00
content:
##########################################################################
# An IDAPython plugin that generates "fuzzy" function signatures that can
# be shared and applied amongst different IDBs.
#
# There are multiple sets of signatures that are generated:
#
#   o "Formal" signatures, where functions must match exactly
#   o "Fuzzy" signatures, where functions must only resemble each other
#     in terms of data/call references.
#   o String-based signatures, where functions are identified based on
#     unique string references.
#   o Immediate-based signatures, where functions are identified based
#     on immediate value references.
#
# These signatures are applied based on accuracy, that is, formal
# signatures are applied first, then string and immediate based
# signatures, and finally fuzzy signatures.
#
# Further, functions are identified based on call references. Consider,
# for example, two functions, one named 'foo', the other named 'bar'.
# The 'foo' function is fairly unique and a reliable signature is easily
# generated for it, but the 'bar' function is more difficult to reliably
# identify. However, 'foo' calls 'bar', and thus once 'foo' is identified,
# 'bar' can also be identified by association.
#
# Craig Heffner
# @devttys0
##########################################################################

from __future__ import print_function
import idc
import idaapi
import idautils

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

from shims import ida_shims


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

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

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

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

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

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


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

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


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

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

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


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

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


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

    DEFAULT_SIGNATURE_FILE = "rizzo.sig"

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

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

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

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

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

    def load(self):
        print(("Loading signatures from %s..." % self.sigfile), end=' ')
        fp = open(self.sigfile, "rb")
        sigs = pickle.load(fp)
        fp.close()
        print("done.")
        return sigs

    def sighash(self, value):
        h = hashlib.md5()
        h.update(value.encode("utf-8"))
        hash_value = h.hexdigest()
        del(h)
        return hash_value

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

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

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

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

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

            ea = ida_shims.next_head(ea)

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

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

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

        return blocks

    def generate(self):
        signatures = RizzoSignatures()

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

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

        # Generate formal, fuzzy, and immediate-based function signatures
        for ea in idautils.Functions():
            func = idaapi.get_func(ea)
            if func:
                # Generate a signature for each block in this function
                blocks = self.function(func)

                # Build function-wide formal and fuzzy signatures by simply
                # concatenating the individual function block signatures.
                formal = self.sighash(''.join([str(e)
                                               for (e, f, i, c) in blocks]))
                fuzzy = self.sighash(''.join([str(f)
                                              for (e, f, i, c) in blocks]))

                # Add this signature to the function dictionary.
                start_ea = ida_shims.start_ea(func)
                signatures.functions[start_ea] = (ida_shims.get_name(start_ea),
                                                  blocks)

                # Check for and remove formal duplicate signatures.
                # If no duplicates, add this to the formal signature dict.
                if formal in signatures.formal:
                    del signatures.formal[formal]
                    signatures.formaldups.add(formal)
                elif formal not in signatures.formaldups:
                    signatures.formal[formal] = ida_shims.start_ea(func)

                # Check for and remove fuzzy duplicate signatures.
                # If no duplicates, add this to the fuzzy signature dict.
                if fuzzy in signatures.fuzzy:
                    del signatures.fuzzy[fuzzy]
                    signatures.fuzzydups.add(fuzzy)
                elif fuzzy not in signatures.fuzzydups:
                    signatures.fuzzy[fuzzy] = ida_shims.start_ea(func)

                # Check for and remove immediate duplicate signatures.
                # If no duplicates, add this to the immediate signature dict.
                for (e, f, immediates, c) in blocks:
                    for immediate in immediates:
                        if immediate in signatures.immediates:
                            del signatures.immediates[immediate]
                            signatures.immediatedups.add(immediate)
                        elif immediate not in signatures.immediatedups:
                            signatures.immediates[immediate] = \
                                ida_shims.start_ea(func)

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

        # DEBUG
        signatures.show()

        return signatures

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

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

        # Match fuzzy function signatures
        start = time.time()
        for (extsig, ext_func_ea) in extsigs.fuzzy.items():
            if extsig in self.signatures.fuzzy:
                curr_fun = RizzoFunctionDescriptor(
                    self.signatures.fuzzy, self.signatures.functions, extsig)
                new_fun = RizzoFunctionDescriptor(
                    extsigs.fuzzy, extsigs.functions, extsig)
                # Only accept this as a valid match if the functions have the
                # same number of basic code blocks
                if len(curr_fun.blocks) == len(new_fun.blocks):
                    fuzzy[curr_fun] = new_fun
        end = time.time()
        print("Found %d fuzzy matches in %.2f seconds." % (len(fuzzy),
                                                           (end-start)))

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

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

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

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

    def apply(self, extsigs):
        count = 0

        start = time.time()

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

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

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

                bm = {}
                duplicates = set()

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

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

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

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

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


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


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


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


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

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

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

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


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

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

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


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

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

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

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

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

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

        return idaapi.PLUGIN_KEEP

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

    def run(self, arg):
        return None

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


def PLUGIN_ENTRY():
    return RizzoPlugin()


repository: Xunhaoz/CodeLamaDecompiler
file: 6. IDA plugin/CodeLlamaDecompilerPlugin.py
wanted_name: LLM Decompile
comment: Show decompile pseudocode and export function
created_at: 2024-08-11 06:40:30+00:00
pushed_at: 2024-08-11 09:51:15+00:00
url:
https://github.com/Xunhaoz/CodeLamaDecompiler/blob/885c725936f35525e27cd81338f7d6780f9e6283/6.%20IDA%20plugin/CodeLlamaDecompilerPlugin.py

content:
import ida_idaapi
import ida_kernwin
import ida_lines
import ida_funcs
import ida_hexrays
import ida_bytes
import ida_ua
import ida_name
import idautils
import idc
import idaapi  # Add this import
import os
import re


# Decompile view
class LLMDecompilerView(ida_kernwin.simplecustviewer_t):
    def Create(self, title):
        if not ida_kernwin.simplecustviewer_t.Create(self, title):
            return False
        # Add "Decompile pseudocode" to the viewer
        self.AddLine("Decompile pseudocode")
        return True

    # close  viewer
    def OnKeydown(self, vkey, shift):
        if vkey == ord('X'):
            self.Close()
        return True


# Plugin
class CodeLlamaDecompilerPlugin(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_UNL
    comment = "Show decompile pseudocode and export function"
    help = "This plugin shows Decompile in a custom viewer and exports the current function"
    wanted_name = "LLM Decompile"  # Plugin name
    wanted_hotkey = "Ctrl-Shift-K"  # Plugin hotkey

    def init(self):
        return ida_idaapi.PLUGIN_OK

    def run(self, arg):
        # plugin run
        # Create and display Decompile view
        viewer = LLMDecompilerView()
        if viewer.Create("Decompile Viewer"):
            viewer.Show()
        else:
            print("Failed to create Decompile viewer")

        # Get the address of the current mouse position
        current_ea = ida_kernwin.get_screen_ea()
        # Get the function where the address is located
        func = ida_funcs.get_func(current_ea)
        if func:
            # If a function is found, export the function and use the API to get function information
            self.export_function(func)

            # Use the new API interface to obtain function information
            func_info = self.api_get_function_info(current_ea)
            if func_info:
                print(f"Function Name: {func_info['name']}")
                print("Assembly (first 5 lines):")
                print("\n".join(func_info['assembly'].split("\n")[:5]))
                print("\nPseudocode (first 5 lines):")
                print("\n".join(func_info['pseudocode'].split("\n")[:5]))
        else:
            print("No function found at current address.")

    def term(self):
        # How to handle plugin termination
        pass

    def export_function(self, func):
        # Export the assembly and pseudocode of the function to the desktop
        func_name = ida_funcs.get_func_name(func.start_ea)
        desktop_path = os.path.join(os.path.expanduser("~"), "Desktop")

        # Export assembly
        asm_path = os.path.join(desktop_path, f"{func_name}_assembly.txt")
        with open(asm_path, "w") as f:
            f.write(self.get_assembly(func_name))
        print(f"Assembly exported to {asm_path}")

        # Export pseudocode
        pseudo_path = os.path.join(desktop_path, f"{func_name}_pseudocode.txt")
        with open(pseudo_path, "w") as f:
            f.write(self.get_pseudocode(func_name))
        print(f"Pseudocode exported to {pseudo_path}")

    def get_assembly(self, func_name):
        # Read the function's assembly
        func_addr = idc.get_name_ea_simple(func_name)
        if func_addr == idc.BADADDR:
            return "Function not found"
        func = idaapi.get_func(func_addr)
        if not func:
            return "Function not found"
        func_start, func_end = func.start_ea, func.end_ea
        asm_code = [f"{func_name} proc near\n"]
        for head in idautils.Heads(func_start, func_end):
            disasm_line = idc.generate_disasm_line(head, 0)
            if disasm_line:
                asm_code.append(disasm_line)
        return '\n'.join(asm_code + [f"{func_name} endp\n"])

    def get_pseudocode(self, function_name):
        # Read the function's pseudocode
        function_address = ida_name.get_name_ea(ida_idaapi.BADADDR, function_name)
        if function_address == ida_idaapi.BADADDR:
            return "Function not found"
        func = ida_funcs.get_func(function_address)
        if not func:
            return "Function not found"
        cfunc = ida_hexrays.decompile(func)
        if not cfunc:
            return "Failed to decompile function"
        return str(cfunc)

    def api_get_function_info(self, ea):
        # API interface: Get function information of the specified address
        func = ida_funcs.get_func(ea)  # 獲取指定地址的function
        if func:
            func_name = ida_funcs.get_func_name(func.start_ea)
            return {
                "name": func_name,  # 獲取function name
                "assembly": self.get_assembly(func_name),  # 獲取assembly
                "pseudocode": self.get_pseudocode(func_name)  # 獲取pseudocode
            }
        else:
            return None


# IDAplugin entry
def PLUGIN_ENTRY():
    return CodeLlamaDecompilerPlugin()


repository: gaasedelen/prefix
file: plugin/ida_prefix.py
description: Function Prefixing for IDA Pro
wanted_name: prefix
comment: A plugin for easy function prefixing
created_at: 2017-05-22 05:15:21+00:00
pushed_at: 2020-04-24 05:00:32+00:00
url:
https://github.com/gaasedelen/prefix/blob/8ad9dc63b388e947136bfcd08a509512eeffeb27/plugin/ida_prefix.py

content:
import os

import idc
import idaapi
import idautils

from PyQt5 import QtGui, QtCore, QtWidgets

#------------------------------------------------------------------------------
# IDA Plugin
#------------------------------------------------------------------------------

VERSION = "v1.2"
AUTHORS = ['Andrew Marumoto', 'Markus Gaasedelen']

def PLUGIN_ENTRY():
    """
    Required plugin entry point for IDAPython Plugins.
    """
    return prefix_t()

class prefix_t(idaapi.plugin_t):
    """
    The IDA Plugin for Prefix.
    """

    flags = idaapi.PLUGIN_PROC | idaapi.PLUGIN_HIDE
    help = ""
    comment = "A plugin for easy function prefixing"
    wanted_name = "prefix"
    wanted_hotkey = ""

    #--------------------------------------------------------------------------
    # Plugin Overloads
    #--------------------------------------------------------------------------

    def init(self):
        """
        This is called by IDA when it is loading the plugin.
        """

        # initialize the menu actions our plugin will inject
        self._init_action_bulk()
        self._init_action_clear()
        self._init_action_recursive()

        # initialize plugin hooks
        self._init_hooks()

        # done
        idaapi.msg("%s %s initialized...\n" % (self.wanted_name, VERSION))
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        """
        This is called by IDA when this file is loaded as a script.
        """
        idaapi.msg("%s cannot be run as a script.\n" % self.wanted_name)

    def term(self):
        """
        This is called by IDA when it is unloading the plugin.
        """

        # unhook our plugin hooks
        self._hooks.unhook()

        # unregister our actions & free their resources
        self._del_action_bulk()
        self._del_action_clear()
        self._del_action_recursive()

        # done
        idaapi.msg("%s terminated...\n" % self.wanted_name)

    #--------------------------------------------------------------------------
    # Plugin Hooks
    #--------------------------------------------------------------------------

    def _init_hooks(self):
        """
        Install plugin hooks into IDA.
        """
        self._hooks = Hooks()
        self._hooks.ready_to_run = self._init_hexrays_hooks
        self._hooks.hook()

    def _init_hexrays_hooks(self):
        """
        Install Hex-Rrays hooks (when available).

        NOTE: This is called when the ui_ready_to_run event fires.
        """
        if idaapi.init_hexrays_plugin():
            idaapi.install_hexrays_callback(self._hooks.hxe_callback)

    #--------------------------------------------------------------------------
    # IDA Actions
    #--------------------------------------------------------------------------

    ACTION_BULK      = "prefix:bulk"
    ACTION_CLEAR     = "prefix:clear"
    ACTION_RECURSIVE = "prefix:recursive"

    def _init_action_bulk(self):
        """
        Register the bulk prefix action with IDA.
        """

        # load the icon for this action
        self._bulk_icon_id = idaapi.load_custom_icon(plugin_resource("bulk.png"))

        # describe the action
        action_desc = idaapi.action_desc_t(
            self.ACTION_BULK,                        # The action name.
            "Prefix selected functions",             # The action text.
            IDACtxEntry(bulk_prefix),                # The action handler.
            None,                                    # Optional: action shortcut
            "Assign a user prefix to the selected functions", # Optional: tooltip
            self._bulk_icon_id                       # Optional: the action icon
        )

        # register the action with IDA
        assert idaapi.register_action(action_desc), "Action registration failed"

    def _init_action_clear(self):
        """
        Register the clear prefix action with IDA.
        """

        # load the icon for this action
        self._clear_icon_id = idaapi.load_custom_icon(plugin_resource("clear.png"))

        # describe the action
        action_desc = idaapi.action_desc_t(
            self.ACTION_CLEAR,                       # The action name.
            "Clear prefixes",                        # The action text.
            IDACtxEntry(clear_prefix),               # The action handler.
            None,                                    # Optional: action shortcut
            "Clear user prefixes from the selected functions", # Optional: tooltip
            self._clear_icon_id                      # Optional: the action icon
        )

        # register the action with IDA
        assert idaapi.register_action(action_desc), "Action registration failed"

    def _init_action_recursive(self):
        """
        Register the recursive rename action with IDA.
        """

        # load the icon for this action
        self._recursive_icon_id = idaapi.load_custom_icon(plugin_resource("recursive.png"))

        # describe the action
        action_desc = idaapi.action_desc_t(
            self.ACTION_RECURSIVE,                   # The action name.
            "Recursive function prefix",             # The action text.
            IDACtxEntry(recursive_prefix_cursor),    # The action handler.
            None,                                    # Optional: action shortcut
            "Recursively prefix callees of this function", # Optional: tooltip
            self._recursive_icon_id                  # Optional: the action icon
        )

        # register the action with IDA
        assert idaapi.register_action(action_desc), "Action registration failed"

    def _del_action_bulk(self):
        """
        Delete the bulk prefix action from IDA.
        """
        idaapi.unregister_action(self.ACTION_BULK)
        idaapi.free_custom_icon(self._bulk_icon_id)
        self._bulk_icon_id = idaapi.BADADDR

    def _del_action_clear(self):
        """
        Delete the clear prefix action from IDA.
        """
        idaapi.unregister_action(self.ACTION_CLEAR)
        idaapi.free_custom_icon(self._clear_icon_id)
        self._clear_icon_id = idaapi.BADADDR

    def _del_action_recursive(self):
        """
        Delete the recursive rename action from IDA.
        """
        idaapi.unregister_action(self.ACTION_RECURSIVE)
        idaapi.free_custom_icon(self._recursive_icon_id)
        self._recursive_icon_id = idaapi.BADADDR

#------------------------------------------------------------------------------
# Plugin Hooks
#------------------------------------------------------------------------------

class Hooks(idaapi.UI_Hooks):

    def ready_to_run(self):
        """
        UI ready to run -- an IDA event fired when everything is spunup.

        NOTE: this is a placeholder func, it gets replaced on a live instance
        but we need it defined here for IDA 7.2+ to properly hook it.
        """
        pass

    def finish_populating_widget_popup(self, widget, popup):
        """
        A right click menu is about to be shown. (IDA 7)
        """
        inject_prefix_actions(widget, popup, idaapi.get_widget_type(widget))
        return 0

    def hxe_callback(self, event, *args):
        """
        HexRays event callback.

        We lump this under the (UI) Hooks class for organizational reasons.
        """

        #
        # if the event callback indicates that this is a popup menu event
        # (in the hexrays window), we may want to install our prefix menu
        # actions depending on what the cursor right clicked.
        #

        if event == idaapi.hxe_populating_popup:
            form, popup, vu = args

            #
            # if the user cursor isn't hovering over a function ref, there
            # is nothing for us to do
            #

            if get_cursor_func_ref() == idaapi.BADADDR:
                return 0

            #
            # the user cursor is hovering over a valid target for a recursive
            # function prefix. insert the prefix action entry into the menu
            #

            idaapi.attach_action_to_popup(
                form,
                popup,
                prefix_t.ACTION_RECURSIVE,
                "Rename global item",
                idaapi.SETMENU_APP
            )

        # done
        return 0

#------------------------------------------------------------------------------
# Prefix Wrappers
#------------------------------------------------------------------------------

def recursive_prefix_cursor():
    """
    Recursive prefix under the user cursor.
    """

    # get the function reference under the user cursor (if there is one)
    target = get_cursor_func_ref()
    if target == idaapi.BADADDR:
        return

    # execute the recursive prefix
    recursive_prefix(target)

def inject_prefix_actions(form, popup, form_type):
    """
    Inject prefix actions to popup menu(s) based on context.
    """

    #
    # disassembly window
    #

    if form_type == idaapi.BWN_DISASMS:

        #
        # if the user cursor isn't hovering over a function ref, there
        # is nothing for us to do
        #

        if get_cursor_func_ref() == idaapi.BADADDR:
            return

        #
        # the user cursor is hovering over a valid target for a recursive
        # function prefix. insert the prefix action entry into the menu
        #

        idaapi.attach_action_to_popup(
            form,
            popup,
            prefix_t.ACTION_RECURSIVE,
            "Rename",
            idaapi.SETMENU_APP
        )

    #
    # functions window
    #

    elif form_type == idaapi.BWN_FUNCS:

        # inject the 'Bulk' function prefix action
        idaapi.attach_action_to_popup(
            form,
            popup,
            prefix_t.ACTION_BULK,
            "Delete function(s)...",
            idaapi.SETMENU_INS
        )

        # inject the 'Clear prefix' action
        idaapi.attach_action_to_popup(
            form,
            popup,
            prefix_t.ACTION_CLEAR,
            "Delete function(s)...",
            idaapi.SETMENU_INS
        )

        # inject a menu separator
        idaapi.attach_action_to_popup(
            form,
            popup,
            None,
            "Delete function(s)...",
            idaapi.SETMENU_INS
        )

    # done
    return 0

#------------------------------------------------------------------------------
# Prefix API
#------------------------------------------------------------------------------

PREFIX_DEFAULT = "MyPrefix"
PREFIX_SEPARATOR = '%'

def recursive_prefix(addr):
    """
    Recursively prefix a function tree with a user defined string.
    """
    func_addr = idaapi.get_name_ea(idaapi.BADADDR, idaapi.get_func_name(addr))
    if func_addr == idaapi.BADADDR:
        idaapi.msg("Prefix: 0x%08X does not belong to a defined function\n" % addr)
        return

    # prompt the user for a prefix to apply to the selected functions
    tag = idaapi.ask_str(PREFIX_DEFAULT, 0, "Function Tag")

    # the user closed the window... ignore
    if tag == None:
        return

    # the user put a blank string and hit 'okay'... notify & ignore
    elif tag == '':
        idaapi.warning("[ERROR] Tag cannot be empty [ERROR]")
        return

    # recursively collect all the functions called by this function
    nodes_xref_down = graph_down(func_addr, path=set([]))

    # graph_down returns the int address needs to be converted
    tmp  = []
    tmp1 = ''
    for func_addr in nodes_xref_down:
        tmp1 = idaapi.get_func_name(func_addr)
        if tmp1:
            tmp.append(tmp1)
    nodes_xref_down = tmp

    # prefix the tree of functions
    for rename in nodes_xref_down:
        func_addr = idaapi.get_name_ea(idaapi.BADADDR, rename)
        if tag not in rename:
            idaapi.set_name(func_addr,'%s%s%s' % (str(tag), PREFIX_SEPARATOR, rename), idaapi.SN_NOWARN)

    # refresh the IDA views
    refresh_views()

def bulk_prefix():
    """
    Prefix the Functions window selection with a user defined string.
    """

    # prompt the user for a prefix to apply to the selected functions
    tag = idaapi.ask_str(PREFIX_DEFAULT, 0, "Function Tag")

    # the user closed the window... ignore
    if tag == None:
        return

    # the user put a blank string and hit 'okay'... notify & ignore
    elif tag == '':
        idaapi.warning("[ERROR] Tag cannot be empty [ERROR]")
        return

    #
    # loop through all the functions selected in the 'Functions window' and
    # apply the user defined prefix tag to each one.
    #

    for func_name in get_selected_funcs():

        # ignore functions that already have the specified prefix applied
        if func_name.startswith(tag):
            continue

        # apply the user defined prefix to the function (rename it)
        new_name  = '%s%s%s' % (str(tag), PREFIX_SEPARATOR, func_name)
        func_addr = idaapi.get_name_ea(idaapi.BADADDR, func_name)
        idaapi.set_name(func_addr, new_name, idaapi.SN_NOWARN)

    # refresh the IDA views
    refresh_views()

def clear_prefix():
    """
    Clear user defined prefixes from the selected functions in the Functions window.
    """

    #
    # loop through all the functions selected in the 'Functions window' and
    # clear any user defined prefixes applied to them.
    #

    for func_name in get_selected_funcs():

        #
        # locate the last (rfind) prefix separator in the function name as
        # we will want to keep everything that comes after it
        #

        i = func_name.rfind(PREFIX_SEPARATOR)

        # if there is no prefix (separator), there is nothing to trim
        if i == -1:
            continue

        # trim the prefix off the original function name and discard it
        new_name  = func_name[i+1:]
        func_addr = idaapi.get_name_ea(idaapi.BADADDR, func_name)
        idaapi.set_name(func_addr, new_name, idaapi.SN_NOWARN)

    # refresh the IDA views
    refresh_views()

#------------------------------------------------------------------------------
# IDA Util
#------------------------------------------------------------------------------

def refresh_views():
    """
    Refresh the IDA views.
    """

    # refresh IDA views
    idaapi.refresh_idaview_anyway()

    # refresh hexrays
    current_widget = idaapi.get_current_widget()
    vu = idaapi.get_widget_vdui(current_widget)
    if vu:
        vu.refresh_ctext()

def get_all_funcs():
    """
    Enumerate all function names defined in the IDB.
    """
    return set(idaapi.get_func_name(ea) for ea in idautils.Functions())

def get_cursor_func_ref():
    """
    Get the function reference under the user cursor.

    Returns BADADDR or a valid function address.
    """
    current_widget = idaapi.get_current_widget()
    form_type      = idaapi.get_widget_type(current_widget)
    vu = idaapi.get_widget_vdui(current_widget)

    #
    # hexrays view is active
    #

    if vu:
        cursor_addr = vu.item.get_ea()

    #
    # disassembly view is active
    #

    elif form_type == idaapi.BWN_DISASM:
        cursor_addr = idaapi.get_screen_ea()
        opnum = idaapi.get_opnum()

        if opnum != -1:

            #
            # if the cursor is over an operand value that has a function ref,
            # use that as a valid rename target
            #

            op_addr = idc.get_operand_value(cursor_addr, opnum)
            op_func = idaapi.get_func(op_addr)

            if op_func and op_func.start_ea == op_addr:
                return op_addr

    # unsupported/unknown view is active
    else:
        return idaapi.BADADDR

    #
    # if the cursor is over a function definition or other reference, use that
    # as a valid rename target
    #

    cursor_func = idaapi.get_func(cursor_addr)
    if cursor_func and cursor_func.start_ea == cursor_addr:
        return cursor_addr

    # fail
    return idaapi.BADADDR

def get_selected_funcs():
    """
    Return the list of function names selected in the Functions window.
    """
    import sip
    twidget = idaapi.find_widget("Functions window")
    widget  = sip.wrapinstance(int(twidget), QtWidgets.QWidget)

    # TODO: test this
    if not widget:
        idaapi.warning("Unable to find 'Functions window'")
        return

    #
    # locate the table widget within the Functions window that actually holds
    # all the visible function metadata
    #

    table = widget.findChild(QtWidgets.QTableView)

    #
    # scrape the selected function names from the Functions window table
    #

    selected_funcs = [str(s.data()) for s in table.selectionModel().selectedRows()]

    #
    # re-map the scraped names as they appear in the function table, to their true
    # names as they are saved in the IDB. See the match_funcs(...) function
    # comment for more details
    #

    return match_funcs(selected_funcs)

def match_funcs(qt_funcs):
    """
    Convert function names scraped from Qt to their *actual* representation.

    The function names we scrape from the Functions window Qt table actually
    use the underscore character ('_') as a substitute for a variety of
    different characters.

    For example, a function named foo%bar in the IDB will appears as foo_bar
    in the Functions window table.

    This function takes a list of names as they appear in the Functions window
    table such as the following:

        ['foo_bar']

    And applies a best effort lookup to return a list of the 'true' function
    names as they are stored in the IDB.

       ['foo%bar']

    TODO: rewrite this to be more efficient for larger idbs
    TODO: takes first matching function, may want to change it to make the requirements more strict
    """
    res = set()
    ida_funcs = get_all_funcs()
    for f in qt_funcs:
        for f2 in ida_funcs:
            if len(f) == len(f2):
                i = 0
                while i < len(f) and (f[i] == f2[i] or f[i] == '_'):
                    i += 1

                if i == len(f):
                    res.add(f2)
                    break

    return list(res)

def graph_down(ea, path=set()):
    """
    Recursively collect all function calls.

    Copied with minor modifications from
    http://hooked-on-mnemonics.blogspot.com/2012/07/renaming-subroutine-blocks-and.html
    """
    path.add(ea)

    #
    # extract all the call instructions from the current function
    #

    call_instructions = []
    instruction_info = idaapi.insn_t()
    for address in idautils.FuncItems(ea):

        # decode the instruction
        if not idaapi.decode_insn(instruction_info, address):
            continue

        # check if this instruction is a call
        if not idaapi.is_call_insn(instruction_info):
            continue

        # save this address as a call instruction
        call_instructions.append(address)

    #
    # iterate through all the instructions in the target function (ea) and
    # inspect all the call instructions
    #

    for x in call_instructions:

        #  TODO
        for r in idautils.XrefsFrom(x, idaapi.XREF_FAR):
            #print(0x%08X" % h, "--calls-->", "0x%08X" % r.to)
            if not r.iscode:
                continue

            # get the function pointed at by this call
            func = idaapi.get_func(r.to)
            if not func:
                continue

            # ignore calls to imports / library calls / thunks
            if (func.flags & (idaapi.FUNC_THUNK | idaapi.FUNC_LIB)) != 0:
                continue

            #
            # if we have not traversed to the destination function that this
            # call references, recurse down to it to continue our traversal
            #

            if r.to not in path:
                graph_down(r.to, path)

    return path

class IDACtxEntry(idaapi.action_handler_t):
    """
    A basic Context Menu class to utilize IDA's action handlers.
    """

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

    def activate(self, ctx):
        """
        Execute the embedded action_function when this context menu is invoked.
        """
        self.action_function()
        return 1

    def update(self, ctx):
        """
        Ensure the context menu is always available in IDA.
        """
        return idaapi.AST_ENABLE_ALWAYS

#------------------------------------------------------------------------------
# Plugin Util
#------------------------------------------------------------------------------

PLUGIN_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), prefix_t.wanted_name))

def plugin_resource(resource_name):
    """
    Return the full path for a given plugin resource file.
    """
    return os.path.join(
        PLUGIN_PATH,
        "resources",
        resource_name
    )


repository: D1ag0n-Young/PwnPatcher
file: PwnPatcher.py
description: A quick fix tool for pwn
comment: PwnPatcher for awd/awdp
created_at: 2023-09-20 09:25:25+00:00
pushed_at: 2023-09-25 03:29:12+00:00
url:
https://github.com/D1ag0n-Young/PwnPatcher/blob/ce25c3e80c6e1a2b2a33efc4b6fc3c09cc976ec4/PwnPatcher.py

content:
#########################################################################################################
is_debug = False
if is_debug:
    ''' 
        Install pydevd:

        1. sudo pip install pydevd

        or

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

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

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

from PyQt5.QtWidgets import *
from PyQt5.QtCore import Qt

import idaapi
import idautils
import idc
import ida_kernwin
from pwn import *
import keystone
import ctypes
import shutil
import struct
import sys

PYTHON_VERSION = sys.version_info[0]

class AwdPwnPatcher:
    def __init__(self, path,arch_tag='',bits=0):
        try:
            self.path = path
            self.save_path = path + "_patch"
            self.binary = ELF(self.path)
            self.patch_nums = 0
            self.bits = self.binary.bits if bits == 0 else bits
            self.pie = self.binary.pie
            self.endian = self.binary.endian
            self.arch = self.binary.arch if arch_tag == '' else arch_tag
        except Exception as e:
            QMessageBox.warning(None, "Error", "Sorry, PwnPatcher not support this architecture!")
            return
        if self.bits != 32 and self.bits != 64:
            QMessageBox.warning(None, "Error", "Sorry, the architecture of program is neither 32-bit or 64-bit.")
            return False
        if self.arch == "arm":
            self.ks_arch = keystone.KS_ARCH_ARM
            self.ks_mode = keystone.KS_MODE_ARM
        elif self.arch == "aarch64":
            self.ks_arch = keystone.KS_ARCH_ARM64
            self.ks_mode = 0
        elif self.arch == "i386" or self.arch == "amd64":
            self.ks_arch = keystone.KS_ARCH_X86
            self.ks_mode = keystone.KS_MODE_32 if self.bits == 32 else keystone.KS_MODE_64
        elif self.arch == "mips" or self.arch == "mips64":
            self.ks_arch = keystone.KS_ARCH_MIPS
            self.ks_mode = keystone.KS_MODE_MIPS32 if self.bits == 32 else keystone.KS_MODE_MIPS64
        else:
            self.ks_mode = 0
            self.ks_arch = 0
        if self.endian == "little":
            self.ks_mode |= keystone.KS_MODE_LITTLE_ENDIAN
        else:
            self.ks_mode |= keystone.KS_MODE_BIG_ENDIAN
        if self.ks_arch != 0:
            self.ks = keystone.Ks(self.ks_arch, self.ks_mode)
        self.eh_frame_section = self.binary.get_section_by_name(".eh_frame")
        self.eh_frame_addr = self.eh_frame_section.header.sh_addr
        self.eh_frame_size = self.eh_frame_section.header.sh_size
        self.offset = 0
        self.constant_offset = 0
        self.adjust_eh_frame_size()

    def adjust_eh_frame_size(self):
        if self.arch == "arm" or self.arch == "aarch64" or self.arch == "mips" or self.arch == "mips64":
            PAGE_SIZE = 0x1000
            for i in range(self.binary.num_sections()):
                section = self.binary.get_section(i)
                if self.binary._get_section_name(section) == ".eh_frame":
                    break
            if self.arch == "mips64":
                self.note_section = self.binary.get_section(i+1)
                self.ctors_section = self.binary.get_section(i+2)
                self.offset = self.eh_frame_size + self.note_section.header.sh_size
                self.eh_frame_next_section = self.ctors_section
            else:
                self.eh_frame_next_section = self.binary.get_section(i+1)
            self.eh_frame_section_header_offset = self.binary._section_offset(i)
            actual_size = self.eh_frame_next_section.header.sh_offset - self.eh_frame_section.header.sh_offset
            self.eh_frame_end_addr = self.eh_frame_addr + self.eh_frame_size
            if (self.eh_frame_end_addr % PAGE_SIZE) != 0:
                self.eh_frame_end_addr_align = (self.eh_frame_end_addr + PAGE_SIZE) & ctypes.c_uint32(~PAGE_SIZE + 1).value
            self.old_eh_frame_size = self.eh_frame_size
            if self.eh_frame_addr + actual_size > self.eh_frame_end_addr_align:
                self.eh_frame_size = self.eh_frame_end_addr_align - self.eh_frame_addr
            else:
                self.eh_frame_size = actual_size
            load_segment = self.binary.get_segment_for_address(self.eh_frame_addr)
            for i in range(self.binary.num_segments()):
                segment = self.binary.get_segment(i)
                if segment.header.p_vaddr == load_segment.header.p_vaddr:
                    break
            self.load_segment_header_offset = self.binary._segment_offset(i)
            if self.endian == "little":
                endian_fmt = "<"
            else:
                endian_fmt = ">"
            new_size = self.eh_frame_size - self.old_eh_frame_size + load_segment.header.p_filesz
            shutil.copy2(self.path, self.save_path)
            self.bin_file = open(self.save_path, "rb+")
            if self.bits == 32:
                self.bin_file.seek(self.load_segment_header_offset+16)
                self.bin_file.write(struct.pack(endian_fmt+"I", new_size))
                self.bin_file.write(struct.pack(endian_fmt+"I", new_size))
            else:
                self.bin_file.seek(self.load_segment_header_offset+32)
                self.bin_file.write(struct.pack(endian_fmt+"Q", new_size))
                self.bin_file.write(struct.pack(endian_fmt+"Q", new_size))
            self.bin_file.close()
            self.binary = ELF(self.save_path)

            print("old eh_frame_size: %#x" % self.old_eh_frame_size)
        print("eh_frame_size: %#x" % self.eh_frame_size)

    def add_patch_in_ehframe(self, assembly="", machine_code=[]):
        patch_start_addr = self.eh_frame_addr + self.offset
        if len(assembly) != 0 :
            shellcode, count = self.ks.asm(assembly, addr=patch_start_addr)
            shellcode = "".join([chr(x) for x in shellcode])
        elif len(machine_code) != 0:
            shellcode = "".join([chr(x) for x in machine_code])
        else:
            shellcode = ""

        if len(shellcode) == 0:
            return 0

        self.offset += len(shellcode)
        assert(self.offset <= self.eh_frame_size)
        if PYTHON_VERSION == 3:
            shellcode = shellcode.encode("latin-1")
        self.binary.write(patch_start_addr, shellcode)
        self.patch_nums += len(shellcode)
        return patch_start_addr

    def patch_origin(self, start, end=0, assembly="", machine_code=[]):
        if len(assembly) != 0:
            shellcode, count = self.ks.asm(assembly, addr=start)
            shellcode = "".join([chr(x) for x in shellcode])
        elif len(machine_code) != 0:
            shellcode = "".join([chr(x) for x in machine_code])
        else:
            shellcode = ""
        if end != 0:
            assert(len(shellcode) <= (end-start))
            shellcode = shellcode.ljust(end - start, "\x90")
        if PYTHON_VERSION == 3:
            shellcode = shellcode.encode("latin-1")
        self.binary.write(start, shellcode)
        self.patch_nums += len(shellcode)

    def patch_by_jmp(self, jmp_from, jmp_to=0, assembly="", machine_code=[]):
        if self.arch == "i386" or self.arch == "amd64":
            jmp_ins = "jmp"
        elif self.arch == "arm" or self.arch == "aarch64":
            jmp_ins = "b"
        elif self.arch == "mips" or self.arch =="mips64":
            if self.pie:
                jmp_ins = "b"
            else:
                jmp_ins = "j"
        if jmp_to:
            payload = "{} {}".format(jmp_ins, hex(jmp_to))
            if len(assembly) != 0:
                assembly += "\n" + payload
            else:
                addr = self.get_next_patch_start_addr() + len(machine_code)
                shellcode, count = self.ks.asm(payload, addr=addr)
                machine_code += shellcode
        patch_start_addr = self.add_patch_in_ehframe(assembly=assembly, machine_code=machine_code)
        if jmp_to:
            # fix translation bug of mips jump code: when keystone translates jmp code, it treats the value of argument start as the base address,
            # rather than the address of jump code.
            # FYI: shellcode, count = self.ks.asm(assembly, addr=patch_start_addr)
            if self.arch == "mips" or self.arch == "mips64":
                next_patch_addr = self.get_next_patch_start_addr()
                payload = "{} {}".format(jmp_ins, hex(jmp_to))
                # why - 8? because a nop code will be added automatically after jmp code.
                self.patch_origin(next_patch_addr-8, assembly=payload)

        if patch_start_addr == 0:
            return 0
        payload = "{} {}".format(jmp_ins, hex(patch_start_addr))
        self.patch_origin(jmp_from, assembly=payload)
        return patch_start_addr

    def patch_by_call(self, call_from, assembly="", machine_code=[]):
        if self.arch != "i386" and self.arch != "amd64":
            QMessageBox.warning(None, "Error", "Sorry, patch_by_call only support x86 architecture!")
            return False
        patch_start_addr = self.add_patch_in_ehframe(assembly=assembly, machine_code=machine_code)
        if patch_start_addr == 0:
            return 0

        payload = "call {}".format(hex(patch_start_addr))
        self.patch_origin(call_from, assembly=payload)
        return patch_start_addr

    def add_constant_in_ehframe(self, string):
        patch_start_addr = self.eh_frame_addr + self.offset
        if PYTHON_VERSION == 3:
            string = string.encode("latin-1")
        self.binary.write(patch_start_addr, string)
        self.offset += len(string)
        self.patch_nums += len(string)
        return patch_start_addr

    def add_constant_in_addr(self, patch_start_addr, string):
        patch_start_addr = patch_start_addr + self.constant_offset
        if PYTHON_VERSION == 3:
            string = string.encode("latin-1")
        self.binary.write(patch_start_addr, string)
        self.constant_offset += len(string)
        self.patch_nums += len(string)
        return patch_start_addr,self.constant_offset
    
    def patch_fmt_by_call(self, call_from):
        if self.arch != "i386" and self.arch != "amd64":
            QMessageBox.warning(None, "Error", "Sorry, patch_fmt_by_call only support x86 architecture!")
            return False,''
        fmt_addr = self.add_constant_in_ehframe("%s\x00\x00")
        patch_start_addr = self.eh_frame_addr + self.offset
        call_first = u8(self.binary.read(call_from, 1))
        if 0xe8 == call_first:
            printf_addr = (call_from + 5 + u32(self.binary.read(call_from+1, 4))) & 0xffffffff
        else:
            QMessageBox.warning(None, "Error", "The assembly at this address is not a call instruction")
            return False,''
        if self.bits == 32 and not self.pie:
            assembly = """
            mov eax, dword ptr [esp+4]
            push eax
            lea eax, dword ptr [{0}]
            push eax
            call {1}
            add esp, 0x8
            ret
            """.format(hex(fmt_addr), hex(printf_addr))
        elif self.bits == 32 and self.pie:
            assembly = """
            call {0}
            mov eax, dword ptr [esp+8]
            push eax
            mov eax, dword ptr [esp+4]
            sub eax, {0}
            add eax, {1}
            push eax
            call {2}
            add esp, 0xc
            ret
            """.format(hex(patch_start_addr+5), fmt_addr, hex(printf_addr))
        else:
            assembly = """
            push rbp
            mov rsi, rdi
            lea rdi, qword ptr [{0}]
            call {1}
            pop rbp
            ret
            """.format(hex(fmt_addr), hex(printf_addr))
        self.patch_by_call(call_from, assembly=assembly)
        return True,assembly
            
    def save(self):
        self.fix_eh_frame_flags()
        self.binary.save(self.save_path)

    def get_next_patch_start_addr(self):
        return self.eh_frame_addr + self.offset

    def fix_eh_frame_flags(self):
        e_phnum = self.binary.header.e_phnum
        e_phoff = self.binary.header.e_phoff
        phdr_size = 32 if self.bits == 32 else 56
        p_flags_offset = 24 if self.bits == 32 else 4
        for i in range(0, e_phnum):
            phdr = self.binary.get_segment(i).header
            page_start = int((phdr.p_vaddr / 0x1000) * 0x1000)
            page_end = phdr.p_vaddr + phdr.p_memsz
            if page_end % 0x1000 != 0:
                page_end = (page_end / 0x1000) * 0x1000 + 0x1000
                page_end = int(page_end)
            if phdr.p_type == "PT_LOAD" and page_start <= self.eh_frame_addr and page_end >= self.eh_frame_addr + self.eh_frame_size:
                print("fix_eh_frame_flags:\npage_start: {} page_end: {} eh_frame_addr: {} eh_frame_size: {} origin phdr.p_flags: {}"
                      .format(hex(page_start), hex(page_end), hex(self.eh_frame_addr), hex(self.eh_frame_size), str(phdr.p_flags)))
                flags = chr(phdr.p_flags | 1)
                if PYTHON_VERSION == 3:
                    flags = flags.encode("latin-1")
                self.binary.write(e_phoff + phdr_size * i + p_flags_offset, flags)

class PatchRecord:
    def __init__(self):
        self.records = []
    
    def add_patch(self, start, end, code_len,patch_mode, asm_code):
        patch = {'start_addr': start, 'end_addr':end, 'len':code_len, 'patch_mod': patch_mode, 'asm_code': asm_code}
        self.records.append(patch)
    
    def export_table(self):
        table = "start_addr\tend_addr \tcoedlen\tpatch_mod\t\t\tasm_code\n"  # 表头
        for record in self.records:
            row = ""
            if record['patch_mod'] == 'patch_by_original':
                row = "0x{:08X}\t0x{:08X}\t0x{:02X}  \t{} \t{}\n".format(record['start_addr'],record['end_addr'], record['len'],record['patch_mod'], record['asm_code'])
            elif record['patch_mod'] == 'patch_by_call':
                row = "0x{:08X}\t0x{:08X}\t0x{:02X}  \t{} \t\t{}\n".format(record['start_addr'],record['end_addr'], record['len'],record['patch_mod'], record['asm_code'])
            else:
                row = "0x{:08X}\t0x{:08X}\t0x{:02X}  \t{} \t\t\t{}\n".format(record['start_addr'],record['end_addr'], record['len'],record['patch_mod'], record['asm_code'])
            table += row
        return table

class PatchDialog(QDialog):
    def __init__(self, parent=None):
        super(PatchDialog, self).__init__(parent)
        try:
            self.arch = ""
            self.bits = 0
            self.start_addr = 0
            self.end_addr = 0
            self.asm_code = ""
            self.constant = ""
            self.constant_addr = 0
            self.total_patch_nums = 0
            self.constant_type = "eh_frame" 
            self.is_fmt_patch = False
            self.path = self.get_file_patch()
            self.awdpwnpatcher = AwdPwnPatcher(self.path)
            self.init_ui()
            self.record = PatchRecord()
        except Exception as e:
            print(str(e))
            import traceback
            traceback.print_exc()
            return

    def get_start_addr_default(self):
        return ida_kernwin.get_screen_ea()
    
    def get_start_item_size(self,addr):
        return idc.get_item_size(addr)
    
    def get_asm_by_addr(self,addr):
        return idc.GetDisasm(addr).split(';')[0]
    
    def get_file_patch(self):
        return idaapi.get_input_file_path()
    
    def get_constant(self):
        return self.constant_edit.text().strip()

    def set_start_addr_user(self):
        self.start_addr = int(self.start_edit.text().strip(),16)
    
    def set_end_addr_user(self):
        self.end_addr = int(self.end_edit.text().strip(),16)
    
    def set_asm_code(self):
        self.asm_code = self.asm_edit.toPlainText().strip()
    
    def update_output_label(self):
        # 获取编辑框的文本内容
        asm_code = self.asm_edit.toPlainText().strip()
        #print(asm_code)
        # 判断asm_code是否是合法的汇编语句
        try:
            
            if asm_code:
                payload, _ = self.awdpwnpatcher.ks.asm(asm_code, addr=self.start_addr)
                # 更新显示窗口的内容
                self.output_label.setText("Encode code:[%d bytes]" % len(payload))
                payload = ' '.join([hex(num)[2:].upper().zfill(2) for num in payload])
                self.encode_asm.setText(payload)
            else:
                self.output_label.setText("Encode code:")
                self.encode_asm.setText("")
        except (keystone.KsError,Exception) as e:
            # 如果汇编语句不合法，显示错误信息
            self.output_label.setText("Encode code[Syntax error]:")
            self.encode_asm.setText("...")
    
    def init_ui(self):
        layout = QVBoxLayout()
        self.start_addr = self.get_start_addr_default()
        self.start_item_size = self.get_start_item_size(self.start_addr)
        self.start_asm_code = self.get_asm_by_addr(self.start_addr)
        addr_label = QLabel("Patch address range (hex):")
        layout.addWidget(addr_label)

        addr_hbox = QHBoxLayout()
        self.start_edit = QLineEdit()
        #self.start_edit.setReadOnly(True)
        self.start_edit.setText(hex(self.start_addr))
        #self.start_edit.setFixedWidth(200)
        addr_hbox.addWidget(self.start_edit)

        div_label = QLabel(" - ")
        addr_hbox.addWidget(div_label)

        self.end_edit = QLineEdit()
        self.end_edit.setText(hex(self.start_addr+self.start_item_size))
        #self.end_edit.setFixedWidth(200)
        addr_hbox.addWidget(self.end_edit)

        layout.addLayout(addr_hbox)

        arch_label = QLabel("patch constant in:")
        layout.addWidget(arch_label)

        self.arch_combobox = QComboBox()
        self.arch_combobox.addItem("eh_frame")
        self.arch_combobox.addItem("address")
        # 根据self.arch设置默认显示的Item
        self.arch_combobox.setCurrentText("eh_frame")
        self.arch_combobox.currentIndexChanged.connect(lambda index: self.on_constant_change(index))
        layout.addWidget(self.arch_combobox)
        
        constant_hbox = QHBoxLayout()
        addr_label = QLabel("constant:")
        constant_hbox.addWidget(addr_label)
        self.constant_edit = QLineEdit()
        self.constant_edit.setText('%s')
        constant_hbox.addWidget(self.constant_edit)
        layout.addLayout(constant_hbox)

        # 创建水平布局和标签
        constant_addr_hbox = QHBoxLayout()
        address_label = QLabel("constant_address:")
        constant_addr_hbox.addWidget(address_label)
        self.constant_addr_edit = QLineEdit()
        self.constant_addr_edit.setText('None')
        self.constant_addr_edit.setReadOnly(True)
        constant_addr_hbox.addWidget(self.constant_addr_edit)
        layout.addLayout(constant_addr_hbox)

        self.arch = self.awdpwnpatcher.arch
        asm_label = QLabel("Assembly code[%s]:" % self.arch)
        layout.addWidget(asm_label)
        self.asm_edit = QTextEdit()
        self.asm_edit.setText(self.start_asm_code)
        fixed_height = 150
        self.asm_edit.setFixedHeight(fixed_height)
        layout.addWidget(self.asm_edit)
        
        # 添加一个用于实时显示编辑框中更改的数据的文本显示窗口
        self.output_label = QLabel("Encode code:")
        layout.addWidget(self.output_label)
        self.encode_asm = QTextEdit()
        self.update_output_label()
        fixed_height = 100
        self.encode_asm.setFixedHeight(fixed_height)
        self.encode_asm.setReadOnly(True)
        layout.addWidget(self.encode_asm)
        
        # 将编辑框的textChanged信号与更新显示窗口内容的槽函数关联
        self.asm_edit.textChanged.connect(self.update_output_label)


        btn_hbox = QVBoxLayout()
        patch_by_ori_btn = QPushButton("patch constant")
        patch_by_ori_btn.clicked.connect(lambda: self.write_constant())
        btn_hbox.addWidget(patch_by_ori_btn)

        patch_by_jmp_btn = QPushButton("Patch by jmp")
        patch_by_jmp_btn.clicked.connect(lambda: self.patch_by_jmp())
        btn_hbox.addWidget(patch_by_jmp_btn)

        patch_by_ori_btn = QPushButton("Patch by original")
        patch_by_ori_btn.clicked.connect(lambda: self.patch_by_original())
        btn_hbox.addWidget(patch_by_ori_btn)

        patch_by_call_btn = QPushButton("Patch fmt by call")
        patch_by_call_btn.clicked.connect(lambda: self.patch_fmt_by_call())
        btn_hbox.addWidget(patch_by_call_btn)

        patch_init_btn = QPushButton("init patcher")
        patch_init_btn.clicked.connect(lambda: self.init_patcher())
        btn_hbox.addWidget(patch_init_btn)

        btn_row = QHBoxLayout()
        patch_log_btn = QPushButton("patcher log")
        patch_log_btn.clicked.connect(lambda: self.export_log())
        btn_row.addWidget(patch_log_btn)

        save_btn = QPushButton("Save")
        btn_row.addWidget(save_btn)
        save_btn.clicked.connect(lambda: self.save_elf())
        btn_hbox.addLayout(btn_row)

        layout.addLayout(btn_hbox)

        self.setLayout(layout)
        self.setWindowTitle("PwnPatcher - Version 1.5")
        self.show()


    def on_constant_change(self, index):
        if index == 0:
            self.constant_type = "eh_frame"
        else:
            self.constant_type = "address"
   

    def set_value(self):
        self.set_start_addr_user()
        self.set_end_addr_user()
        self.set_asm_code()

    def init_patcher(self):
        self.start_edit.setText(hex(self.start_addr))
        self.end_edit.setText(hex(self.start_addr+self.start_item_size))
        # 根据self.arch设置默认显示的Item
        self.arch_combobox.setCurrentText("eh_frame")
        self.constant_edit.setText('%s')
        self.constant_addr_edit.setText('None')
        self.asm_edit.setText(self.start_asm_code)
        self.output_label.setText("Encode code:")
        self.update_output_label()
        self.total_patch_nums = 0
        self.awdpwnpatcher = AwdPwnPatcher(self.path,self.awdpwnpatcher.binary.arch,self.awdpwnpatcher.binary.bits)
        self.record = PatchRecord()

    def export_log(self):
        print('-'*70)
        print(self.record.export_table())
        print("total nums: " + str(self.total_patch_nums))

    def save_elf(self):
        self.awdpwnpatcher.save()
        print('binary patch saved success!')

    def write_constant(self):
        self.awdpwnpatcher.patch_nums = 0
        self.set_value()
        constant = self.get_constant()
        if not constant:
            QMessageBox.warning(None, "Error", "constant is None!")
            return
        if self.constant_type == "eh_frame":
            self.constant_addr = self.awdpwnpatcher.add_constant_in_ehframe(constant+'\x00\x00')
            if self.constant_addr :
                ehframe_startaddr = self.awdpwnpatcher.eh_frame_addr+self.awdpwnpatcher.offset
                ehframe_endaddr = self.awdpwnpatcher.eh_frame_addr+self.awdpwnpatcher.offset+len(constant)+2
                self.record.add_patch(ehframe_startaddr,ehframe_endaddr,self.awdpwnpatcher.patch_nums,self.constant_type,constant)
                QMessageBox.information(None, "Success", "constant str add successful! Please use offset %s for access"%(hex(self.constant_addr)))
        else:
            self.constant_addr,conut= self.awdpwnpatcher.add_constant_in_addr(self.start_addr,constant+'\x00')
            if self.constant_addr :
                self.record.add_patch(self.constant_addr,self.constant_addr+conut,len(constant+'\x00'),self.constant_type+' ',constant)
                QMessageBox.information(None, "Success", "constant str add in %s"%hex(self.constant_addr))
        
        self.constant_addr_edit.setText(hex(self.constant_addr))
        self.total_patch_nums+=self.awdpwnpatcher.patch_nums

    def get_jmp_len(self):
        if self.arch == "i386" or self.arch == "amd64":
            jmp_ins = "jmp"
        elif self.arch == "arm" or self.arch == "aarch64":
            jmp_ins = "b"
        elif self.arch == "mips" or self.arch =="mips64":
            if self.awdpwnpatcher.pie:
                jmp_ins = "b"
            else:
                jmp_ins = "j"
        patch_start_addr = self.awdpwnpatcher.eh_frame_addr
        payload = "{} {}".format(jmp_ins, hex(patch_start_addr))
        payloadcode, _ = self.awdpwnpatcher.ks.asm(payload,addr=self.start_addr)
        return len(payloadcode),payload


    def patch_by_jmp(self):
        self.awdpwnpatcher.patch_nums = 0
        self.is_fmt_patch = False
        valid,patch_max = self.validate_input()
        if not valid:
            return
        if not self.path:
            QMessageBox.warning(None, "Error", "Failed to get file patch!")
            return
        jmp_len,jmp_asm = self.get_jmp_len()
        try:
            if patch_max < jmp_len:
                QMessageBox.warning(None, "Error", "The given address range is less than the encoding length of jmp_to[%d]!"%self.get_jmp_len())
                return False
            self.awdpwnpatcher.patch_by_jmp(self.start_addr,self.end_addr,self.asm_code)
        except Exception as e:
            QMessageBox.warning(None, "Error", "Please input correct assembly code! " + "info: " + str(e))
            return False
        jmp_mod_asm = jmp_asm +'; -> '+ self.asm_code + ';'+"jmp %s"%hex(self.end_addr)
        self.record.add_patch(self.start_addr,self.end_addr,self.awdpwnpatcher.patch_nums,"patch_by_jmp",jmp_mod_asm.replace('\n',';'))
        self.total_patch_nums+=self.awdpwnpatcher.patch_nums
        print("Code patch_by_jmp[%d bytes] successfully!" % self.awdpwnpatcher.patch_nums)
        QMessageBox.information(None, "Success", "Code patch_by_jmp[%d bytes] successfully!" % self.awdpwnpatcher.patch_nums)

    def patch_by_original(self):
        self.awdpwnpatcher.patch_nums = 0
        self.is_fmt_patch = False
        valid,patch_max = self.validate_input()
        if not valid:
            return
        if not self.path:
            QMessageBox.warning(None, "Error", "Failed to get file patch!")
            return
        hellcode = ""
        try:
            hellcode, _ = self.awdpwnpatcher.ks.asm(self.asm_code,addr=self.start_addr)
            if len(hellcode) > patch_max:
                QMessageBox.warning(None, "Error", "Assembly code length greater than maximum patch length!")
                return False
        except Exception as e:
            QMessageBox.warning(None, "Error", "Please input correct assembly code! " + "info: " + str(e))
            return False
        self.awdpwnpatcher.patch_origin(self.start_addr,self.end_addr,self.asm_code)
        self.record.add_patch(self.start_addr,self.end_addr,self.awdpwnpatcher.patch_nums,"patch_by_original",self.asm_code.replace('\n',';'))
        self.total_patch_nums+=self.awdpwnpatcher.patch_nums
        print("Code patch_by_original[%d bytes] successfully!" % self.awdpwnpatcher.patch_nums)
        QMessageBox.information(None, "Success", "Code patch_by_original[%d bytes] successfully!" % self.awdpwnpatcher.patch_nums)

    def patch_fmt_by_call(self):
        self.awdpwnpatcher.patch_nums = 0
        self.is_fmt_patch = True
        valid, _ = self.validate_input()
        if not valid:
            return
        if not self.path:
            QMessageBox.warning(None, "Error", "Failed to get file patch!")
            return
        call_asm = ''
        call_result = False
        try:
            call_result,call_asm = self.awdpwnpatcher.patch_fmt_by_call(self.start_addr)
        except Exception as e:
            QMessageBox.warning(None, "Error", "patch_fmt_by_call: %s" % str(e))
            return
        if call_result:
            self.record.add_patch(self.start_addr,self.end_addr,self.awdpwnpatcher.patch_nums,"patch_by_call",call_asm.replace('\n',';').replace('            ',''))
            self.total_patch_nums+=self.awdpwnpatcher.patch_nums
            print("Code patch_fmt_by_call[%d bytes] successfully!" % self.awdpwnpatcher.patch_nums)
            QMessageBox.information(None, "Success", "Code patch_fmt_by_call[%d bytes] successfully!" % self.awdpwnpatcher.patch_nums)

    def validate_input(self):
        try:
            self.set_value()
            if not self.start_addr :
                QMessageBox.warning(None, "Error", "Please input start address!")
                return False,0
            if not self.is_fmt_patch and not self.end_addr:
                QMessageBox.warning(None, "Error", "Please input end address!")
                return False,0
        except:
            QMessageBox.warning(None, "Error", "Please input correct start/end address!")
            return False,0

        try:

            if not self.is_fmt_patch :
                if self.start_addr >= self.end_addr:
                    QMessageBox.warning(None, "Error", "Invalid address range!")
                    return False,0

        except ValueError:
            QMessageBox.warning(None, "Error", "Invalid address format!")
            return False,0
        
        if not self.asm_code:
            QMessageBox.warning(None, "Error", "Please input assembly code!")
            return False,0

        path_max = 0
        if not self.is_fmt_patch:
            path_max = self.end_addr-self.start_addr
        return True,path_max


class PatchFilePlugin(idaapi.plugin_t):
    PLUGIN_NAME = "PwnPatcher"
    PLUGIN_VERSION = "1.0"
    comment = "PwnPatcher for awd/awdp"
    help = ""
    wanted_name = PLUGIN_NAME
    flags = idaapi.PLUGIN_KEEP
    wanted_hotkey = "ALT-K"

    def init(self):
        return idaapi.PLUGIN_OK


    def run(self, arg):
        try:
            dialog = PatchDialog()
        except Exception as e:
            return

    def term(self):
        print("PwnPatcher success exited!")

def PLUGIN_ENTRY():
    return PatchFilePlugin()

# pl = PatchFilePlugin()
# pl.run('d')


repository: neoz/arm_ops
file: arm_ops.py
url: https://github.com/neoz/arm_ops/blob/1efb6792cbb9eba1a5ae30dd4868c5e314906d9e/arm_ops.py
wanted_name: ARM Opcode Assembler
comment: Arm Opcode Assembler helper
created_at: 2015-02-04 05:59:51+00:00
pushed_at: 2021-03-03 04:17:25+00:00
content:
import idaapi
import idautils
import idc
import sys

class Error(Exception):
    """Base class for exceptions in this module."""
    pass

class InputError(Error):
    """Exception raised for errors in the input.

    Attributes:
        expr -- input expression in which the error occurred
        msg  -- explanation of the error
    """

    def __init__(self, expr, msg):
        self.expr = expr
        self.msg = msg

def get_arch():
	(arch, bits) = (None, None)
	for x in idaapi.ph_get_regnames():
		name = x
		if name == 'RAX':
			arch = 'amd64'
			bits = 64
			break
		elif name == 'EAX':
			arch = 'x86'
			bits = 32
			break
		elif name == 'R0':
			arch = 'arm'
			bits = 32
			break
	return (arch, bits)

class myplugin_t(idaapi.plugin_t):
	flags = idaapi.PLUGIN_UNL
	comment = "Arm Opcode Assembler helper"
	help = "This is help"
	wanted_name = "ARM Opcode Assembler"
	wanted_hotkey = "Alt-N"
	(arch, bits) = (None,None)

	def init(self):
		(self.arch, self.bits) = get_arch()
		return idaapi.PLUGIN_OK

	def run(self, arg):
		if self.arch == 'arm':
			startarm()

	def term(self):
		pass

def PLUGIN_ENTRY():
	return myplugin_t()

def tohex(val, nbits):
	return hex((val + (1 << nbits)) % (1 << nbits))

def remove_doublespace(str):
	pos = str.find("  ")
	while pos != -1:
		str = str.replace("  "," ")
		pos = str.find("  ")
	return str

def clean_part(str):
	str = str.strip()
	pos = str.find(' ')
	while pos != -1:
		str = str.replace(' ','')
		pos = str.find(' ')
	return str

def read_number(str):
	err_flag = 0
	try:
		err_flag = 0
		num = int(str)
	except ValueError:
		err_flag = 1
		num = 0

	if err_flag == 1:
		try:
			err_flag = 0
			num = int(str,16)
		except ValueError:
			err_flag = 1
			num = 0

	if err_flag == 1:
		try:
			err_flag = 0
			num = int(str,2)
		except ValueError:
			err_flag = 1
			num = 0

	if err_flag == 1:
		raise InputError(str,"Error: Could not read number"+str)

	return num

def remove_prefix(prefix,str):
	if str.startswith(prefix):
		str = str[len(prefix):]
	return str

def get_cond_bit(cond_str):
	if cond_str == "eq":
		cond = "0000"
	elif cond_str == "ne":
		cond = "0001"
	elif cond_str == "cs":
		cond = "0010"
	elif cond_str == "cc":
		cond = "0011"
	elif cond_str == "mi":
		cond = "0100"
	elif cond_str == "pl":
		cond = "0101"
	elif cond_str == "vs":
		cond = "0110"
	elif cond_str == "vc":
		cond = "0111"
	elif cond_str == "hi":
		cond = "1000"
	elif cond_str == "ls":
		cond = "1001"
	elif cond_str == "ge":
		cond = "1010"
	elif cond_str == "lt":
		cond = "1011"
	elif cond_str == "gt":
		cond = "1100"
	elif cond_str == "le":
		cond = "1101"
	elif cond_str == "al":
		cond = "1110"
	else:
		cond = "1110"
	return cond

def get_Rn_bit(Rn_str):
	RnStr_ = Rn_str
	if RnStr_ == "sp":
		RnStr_ = 'r13'
	elif RnStr_ == "lr":
		RnStr_ = 'r14'
	elif RnStr_ == "pc":
		RnStr_ = 'r15'

	n = int(remove_prefix("r",RnStr_))
	return format(n,'b').zfill(4)

def get_args(arg_str):
	parts = arg_str.split(',')
	for i in range(0,len(parts)):
		parts[i] = clean_part(parts[i])
	return parts

def bx_opcode(startea, part):
	cond_str = remove_prefix("bx",part[0])
	cond = get_cond_bit(cond_str)
	Rn = get_Rn_bit(part[1])
	ins_str = cond + '0001' + '0010' + '1111' + '1111' + '1111' + '0001' + Rn
	ins_value = int(ins_str,2)
	PatchDword(startea, ins_value)

def bandbl_opcode(startea, part):
	dest_addr =  read_number(part[1]);
	offset = dest_addr - startea - 8
	offset = offset >> 2
	if part[0] == "bl":
			linkbit = '1'
	elif part[0] in ('b','ble','bls','blt'):
		linkbit = '0'
	else:
		if part[0].startswith("bl"):
			linkbit = '1'
		else:
			linkbit = '0'

	if part[0].startswith("bl") and len(part[0]) == 4:
			cond_str = remove_prefix("bl",part[0])
	else:
		cond_str = remove_prefix("b",part[0])

	cond = get_cond_bit(cond_str)
	ins_str = cond + '101' + linkbit
	ins_value = int(ins_str,2) << 24
	delta = offset & 0x00ffffff
	ins_value = ins_value | delta
	PatchDword(startea, ins_value)

def dataprocessing_opcode(startea, part):
	if part[0].startswith("and"):
		args = get_args(part[1])
		cond_str = remove_prefix("and",part[0])
		cond = get_cond_bit(cond_str)
		#cond =  '1110'
		opcode = '0000'
		imm_operand = '0'
		setcondcode = '0'
		Rn = get_Rn_bit(args[1])
		Rd = get_Rn_bit(args[0])
		operand2 = '000000000000'
		if args[2].startswith("#"):
			imm_operand = '1'
			imm_str = remove_prefix("#",args[2])
			imm_value = read_number(imm_str)
			imm_value = imm_value & 0xFF
			operand2 = '0000' + format(imm_value,'b').zfill(8)
		else:	
			imm_operand = '0'
			operand2 = '00000000' + get_Rn_bit(args[2])
	elif part[0].startswith("eor"):
		args = get_args(part[1])
		cond_str = remove_prefix("eor",part[0])
		cond = get_cond_bit(cond_str)
		#cond =  '1110'
		opcode = '0001'
		imm_operand = '0'
		setcondcode = '0'
		Rn = get_Rn_bit(args[1])
		Rd = get_Rn_bit(args[0])
		operand2 = '000000000000'
		if args[2].startswith("#"):
			imm_operand = '1'
			imm_str = remove_prefix("#",args[2])
			imm_value = read_number(imm_str)
			imm_value = imm_value & 0xFF
			operand2 = '0000' + format(imm_value,'b').zfill(8)
		else:	
			imm_operand = '0'
			operand2 = '00000000' + get_Rn_bit(args[2])
	elif part[0].startswith("sub"):
		args = get_args(part[1])
		cond_str = remove_prefix("sub",part[0])
		cond = get_cond_bit(cond_str)
		#cond =  '1110'
		opcode = '0010'
		imm_operand = '0'
		setcondcode = '0'
		Rn = get_Rn_bit(args[1])
		Rd = get_Rn_bit(args[0])
		operand2 = '000000000000'
		if args[2].startswith("#"):
			imm_operand = '1'
			imm_str = remove_prefix("#",args[2])
			imm_value = read_number(imm_str)
			imm_value = imm_value & 0xFF
			operand2 = '0000' + format(imm_value,'b').zfill(8)
		else:	
			imm_operand = '0'
			operand2 = '00000000' + get_Rn_bit(args[2])
	elif part[0].startswith("rsb"):
		args = get_args(part[1])
		cond_str = remove_prefix("rsb",part[0])
		cond = get_cond_bit(cond_str)
		#cond =  '1110'
		opcode = '0011'
		imm_operand = '0'
		setcondcode = '0'
		Rn = get_Rn_bit(args[1])
		Rd = get_Rn_bit(args[0])
		operand2 = '000000000000'
		if args[2].startswith("#"):
			imm_operand = '1'
			imm_str = remove_prefix("#",args[2])
			imm_value = read_number(imm_str)
			imm_value = imm_value & 0xFF
			operand2 = '0000' + format(imm_value,'b').zfill(8)
		else:	
			imm_operand = '0'
			operand2 = '00000000' + get_Rn_bit(args[2])
	elif part[0].startswith("add"):
		args = get_args(part[1])
		cond_str = remove_prefix("add",part[0])
		cond = get_cond_bit(cond_str)
		#cond =  '1110'
		opcode = '0100'
		imm_operand = '0'
		setcondcode = '0'
		Rn = get_Rn_bit(args[1])
		Rd = get_Rn_bit(args[0])
		operand2 = '000000000000'
		if args[2].startswith("#"):
			imm_operand = '1'
			imm_str = remove_prefix("#",args[2])
			imm_value = read_number(imm_str)
			imm_value = imm_value & 0xFF
			operand2 = '0000' + format(imm_value,'b').zfill(8)
		else:	
			imm_operand = '0'
			operand2 = '00000000' + get_Rn_bit(args[2])
	elif part[0].startswith("adc"):
		args = get_args(part[1])
		cond_str = remove_prefix("adc",part[0])
		cond = get_cond_bit(cond_str)
		#cond =  '1110'
		opcode = '0101'
		imm_operand = '0'
		setcondcode = '0'
		Rn = get_Rn_bit(args[1])
		Rd = get_Rn_bit(args[0])
		operand2 = '000000000000'
		if args[2].startswith("#"):
			imm_operand = '1'
			imm_str = remove_prefix("#",args[2])
			imm_value = read_number(imm_str)
			imm_value = imm_value & 0xFF
			operand2 = '0000' + format(imm_value,'b').zfill(8)
		else:	
			imm_operand = '0'
			operand2 = '00000000' + get_Rn_bit(args[2])
	elif part[0].startswith("sbc"):
		args = get_args(part[1])
		cond_str = remove_prefix("sbc",part[0])
		cond = get_cond_bit(cond_str)
		#cond =  '1110'
		opcode = '0110'
		imm_operand = '0'
		setcondcode = '0'
		Rn = get_Rn_bit(args[1])
		Rd = get_Rn_bit(args[0])
		operand2 = '000000000000'
		if args[2].startswith("#"):
			imm_operand = '1'
			imm_str = remove_prefix("#",args[2])
			imm_value = read_number(imm_str)
			imm_value = imm_value & 0xFF
			operand2 = '0000' + format(imm_value,'b').zfill(8)
		else:	
			imm_operand = '0'
			operand2 = '00000000' + get_Rn_bit(args[2])
	elif part[0].startswith("rsc"):
		args = get_args(part[1])
		cond_str = remove_prefix("rsc",part[0])
		cond = get_cond_bit(cond_str)
		#cond =  '1110'
		opcode = '0111'
		imm_operand = '0'
		setcondcode = '0'
		Rn = get_Rn_bit(args[1])
		Rd = get_Rn_bit(args[0])
		operand2 = '000000000000'
		if args[2].startswith("#"):
			imm_operand = '1'
			imm_str = remove_prefix("#",args[2])
			imm_value = read_number(imm_str)
			imm_value = imm_value & 0xFF
			operand2 = '0000' + format(imm_value,'b').zfill(8)
		else:	
			imm_operand = '0'
			operand2 = '00000000' + get_Rn_bit(args[2])
	elif part[0].startswith("tst"):
		args = get_args(part[1])
		cond_str = remove_prefix("tst",part[0])
		cond = get_cond_bit(cond_str)
		#cond =  '1110'
		opcode = '1000'
		imm_operand = '0'
		setcondcode = '1'
		Rn = get_Rn_bit(args[0])
		Rd = '0000' #get_Rn_bit(args[0])
		operand2 = '000000000000'
		if args[1].startswith("#"):
			imm_operand = '1'
			imm_str = remove_prefix("#",args[1])
			imm_value = read_number(imm_str)
			imm_value = imm_value & 0xFF
			operand2 = '0000' + format(imm_value,'b').zfill(8)
		else:	
			imm_operand = '0'
			operand2 = '00000000' + get_Rn_bit(args[1])
	elif part[0].startswith("teq"):
		args = get_args(part[1])
		#cond =  '1110'
		cond_str = remove_prefix("teq",part[0])
		cond = get_cond_bit(cond_str)
		opcode = '1001'
		imm_operand = '0'
		setcondcode = '1'
		Rn = get_Rn_bit(args[0])
		Rd = '0000' #get_Rn_bit(args[0])
		operand2 = '000000000000'
		if args[1].startswith("#"):
			imm_operand = '1'
			imm_str = remove_prefix("#",args[1])
			imm_value = read_number(imm_str)
			imm_value = imm_value & 0xFF
			operand2 = '0000' + format(imm_value,'b').zfill(8)
		else:	
			imm_operand = '0'
			operand2 = '00000000' + get_Rn_bit(args[1])
	elif part[0].startswith("cmp"):
		args = get_args(part[1])
		#cond =  '1110'
		cond_str = remove_prefix("cmp",part[0])
		cond = get_cond_bit(cond_str)
		opcode = '1010'
		imm_operand = '0'
		setcondcode = '1'
		Rn = get_Rn_bit(args[0])
		Rd = get_Rn_bit(args[0])
		operand2 = '000000000000'
		if args[1].startswith("#"):
			imm_operand = '1'
			imm_str = remove_prefix("#",args[1])
			imm_value = read_number(imm_str)
			imm_value = imm_value & 0xFF
			operand2 = '0000' + format(imm_value,'b').zfill(8)
		else:	
			imm_operand = '0'
			operand2 = '00000000' + get_Rn_bit(args[1])
	elif part[0].startswith("cmn"):
		args = get_args(part[1])
		cond_str = remove_prefix("cmn",part[0])
		cond = get_cond_bit(cond_str)
		#cond =  '1110'
		opcode = '1011'
		imm_operand = '0'
		setcondcode = '1'
		Rn = get_Rn_bit(args[0])
		Rd = get_Rn_bit(args[0])
		operand2 = '000000000000'
		if args[1].startswith("#"):
			imm_operand = '1'
			imm_str = remove_prefix("#",args[1])
			imm_value = read_number(imm_str)
			imm_value = imm_value & 0xFF
			operand2 = '0000' + format(imm_value,'b').zfill(8)
		else:	
			imm_operand = '0'
			operand2 = '00000000' + get_Rn_bit(args[1])
	elif part[0].startswith("orr"):
		args = get_args(part[1])
		cond_str = remove_prefix("orr",part[0])
		cond = get_cond_bit(cond_str)
		#cond =  '1110'
		opcode = '1100'
		imm_operand = '0'
		setcondcode = '0'
		Rn = get_Rn_bit(args[1])
		Rd = get_Rn_bit(args[0])
		operand2 = '000000000000'
		if args[2].startswith("#"):
			imm_operand = '1'
			imm_str = remove_prefix("#",args[2])
			imm_value = read_number(imm_str)
			imm_value = imm_value & 0xFF
			operand2 = '0000' + format(imm_value,'b').zfill(8)
		else:	
			imm_operand = '0'
			operand2 = '00000000' + get_Rn_bit(args[2])
	elif part[0].startswith("mov"):
		args = get_args(part[1])
		cond_str = remove_prefix("mov",part[0])
		cond = get_cond_bit(cond_str)
		#cond =  '1110'
		opcode = '1101'
		imm_operand = '0'
		setcondcode = '0'
		Rn = '0000'
		Rd = get_Rn_bit(args[0])
		operand2 = '000000000000'
		if args[1].startswith("#"):
			imm_operand = '1'
			imm_str = remove_prefix("#",args[1])
			imm_value = read_number(imm_str)
			imm_value = imm_value & 0xFF
			operand2 = '0000' + format(imm_value,'b').zfill(8)
		else:	
			imm_operand = '0'
			operand2 = '00000000' + get_Rn_bit(args[1])
	elif part[0].startswith("bic"):
		args = get_args(part[1])
		#cond =  '1110'
		cond_str = remove_prefix("bic",part[0])
		cond = get_cond_bit(cond_str)
		opcode = '1110'
		imm_operand = '0'
		setcondcode = '0'
		Rn = get_Rn_bit(args[1])
		Rd = get_Rn_bit(args[0])
		operand2 = '000000000000'
		if args[2].startswith("#"):
			imm_operand = '1'
			imm_str = remove_prefix("#",args[2])
			imm_value = read_number(imm_str)
			imm_value = imm_value & 0xFF
			operand2 = '0000' + format(imm_value,'b').zfill(8)
		else:	
			imm_operand = '0'
			operand2 = '00000000' + get_Rn_bit(args[2])
	elif part[0].startswith("mvn"):
		args = get_args(part[1])
		#cond =  '1110'
		cond_str = remove_prefix("mvn",part[0])
		cond = get_cond_bit(cond_str)
		opcode = '1111'
		imm_operand = '0'
		setcondcode = '0'
		Rn = '0000'
		Rd = get_Rn_bit(args[0])
		operand2 = '000000000000'
		if args[1].startswith("#"):
			imm_operand = '1'
			imm_str = remove_prefix("#",args[1])
			imm_value = read_number(imm_str)
			imm_value = imm_value & 0xFF
			operand2 = '0000' + format(imm_value,'b').zfill(8)
		else:	
			imm_operand = '0'
			operand2 = '00000000' + get_Rn_bit(args[1])

	ins_str = cond + '00' + imm_operand + opcode + setcondcode + Rn + Rd + operand2 
	ins_value = int(ins_str,2)
	PatchDword(startea, ins_value)

def singledatatransfer_opcode(startea, part):
	if part[0] == "str":
		part1 = part[1].replace('[','')
		part1 = part1.replace(']','')
		args = get_args(part1)
		cond =  '1110'
		immediaOffset = '0'
		prepostIndexbit = '1'
		updownbit = '1'
		bytewordbit = '0'
		writebackbit = '0'
		loadstorebit = '0'
		Rn = get_Rn_bit(args[1])
		Rd = get_Rn_bit(args[0])
		offset = '000000000000'
		if args[2].startswith("#"):
			immediaOffset = '0'
			imm_str = remove_prefix("#",args[2])
			imm_value = read_number(imm_str)
			imm_value = imm_value & 0xFF
			offset = format(imm_value,'b').zfill(12)
		else:	
			immediaOffset = '1'
			offset = '00000000' + get_Rn_bit(args[2])
	elif part[0] == "strb":
		part1 = part[1].replace('[','')
		part1 = part1.replace(']','')
		args = get_args(part1)
		cond =  '1110'
		immediaOffset = '0'
		prepostIndexbit = '1'
		updownbit = '1'
		bytewordbit = '1'
		writebackbit = '0'
		loadstorebit = '0'
		Rn = get_Rn_bit(args[1])
		Rd = get_Rn_bit(args[0])
		offset = '000000000000'
		if args[2].startswith("#"):
			immediaOffset = '0'
			imm_str = remove_prefix("#",args[2])
			imm_value = read_number(imm_str)
			imm_value = imm_value & 0xFF
			offset = format(imm_value,'b').zfill(12)
		else:	
			immediaOffset = '1'
			offset = '00000000' + get_Rn_bit(args[2])
	elif part[0] == "ldr":
		part1 = part[1].replace('[','')
		part1 = part1.replace(']','')
		args = get_args(part1)
		cond =  '1110'
		immediaOffset = '0'
		prepostIndexbit = '1'
		updownbit = '1'
		bytewordbit = '0'
		writebackbit = '0'
		loadstorebit = '1'
		Rn = get_Rn_bit(args[1])
		Rd = get_Rn_bit(args[0])
		offset = '000000000000'
		if args[2].startswith("#"):
			immediaOffset = '0'
			imm_str = remove_prefix("#",args[2])
			imm_value = read_number(imm_str)
			imm_value = imm_value & 0xFF
			offset = format(imm_value,'b').zfill(12)
		else:	
			immediaOffset = '1'
			offset = '00000000' + get_Rn_bit(args[2])
	elif part[0] == "ldrb":
		part1 = part[1].replace('[','')
		part1 = part1.replace(']','')
		args = get_args(part1)
		cond =  '1110'
		immediaOffset = '0'
		prepostIndexbit = '1'
		updownbit = '1'
		bytewordbit = '1'
		writebackbit = '0'
		loadstorebit = '1'
		Rn = get_Rn_bit(args[1])
		Rd = get_Rn_bit(args[0])
		offset = '000000000000'
		if args[2].startswith("#"):
			immediaOffset = '0'
			imm_str = remove_prefix("#",args[2])
			imm_value = read_number(imm_str)
			imm_value = imm_value & 0xFF
			offset = format(imm_value,'b').zfill(12)
		else:	
			immediaOffset = '1'
			offset = '00000000' + get_Rn_bit(args[2])

	ins_str = cond + '01' + immediaOffset + prepostIndexbit + updownbit + bytewordbit + writebackbit + loadstorebit + Rn + Rd + offset
	ins_value = int(ins_str,2)
	PatchDword(startea, ins_value)

def arm(startea,instruction_string):
	#startea = idc.ScreenEA()
	instruction_string = instruction_string.lower()
	instruction_string = remove_doublespace(instruction_string)
	part = instruction_string.split(' ',1)
	for i in range(0,len(part)):
		part[i] = clean_part(part[i])


	if part[0].startswith(('and','eor','sub','rsb','add','adc','sbc','rsc','tst','teq','cmp','cmn','orr','mov','bic','mvn')):
		dataprocessing_opcode(startea,part)
	elif part[0].startswith("bx"):
		bx_opcode(startea,part)
	elif part[0].startswith("b"):
		bandbl_opcode(startea,part)
	elif part[0] in ('ldr','str','ldrb','strb'):
		singledatatransfer_opcode(startea,part)
	else:
		raise InputError(str,"Invalid instruction "+instruction_string)

def startarm():
	curEA = idc.ScreenEA()
	isCont = 1
	while isCont:
		t = idaapi.generate_disasm_line(curEA)
		if t:
			line = idaapi.tag_remove(t)
		else:
			line = ""
		str = AskStr(line,"Address :"+hex(curEA)+"\nInstruction")
		if str:
			try:
				arm(curEA,str)
				curEA = curEA + 4
			except InputError as e:
				print e.msg
		else:
			isCont = 0


repository: splishy-splash/RE-scripts
file: reference_instruction.py
description: a small curation of created/stolen scripts for reverse engineering
comment: Reference Instruction
created_at: 2022-04-09 13:27:01+00:00
pushed_at: 2024-05-03 02:41:56+00:00
url:
https://github.com/splishy-splash/RE-scripts/blob/1ed9df1bb9baaa05450f60a6fc1636224ad4a9f0/reference_instruction.py

content:
############################################################################################
##
## IDA plugin template code shamelessly stolen from @herrcore's GitHub (with permission)
## https://github.com/OALabs/hexcopy-ida/blob/main/hexcopy.py
##
## One-Click instruction search!
##
## Updated for IDA 7.xx and Python 3
##
## To install:
##      Copy script into plugins directory, i.e: C:\Program Files\<ida version>\plugins
##      Install requests for whichever version of Python3 you are running
##
## To run:
##      Highlight x86 or x64 instruction and right click and select "Reference Instruction"
##      The link to the MSDN page for that function will open in a new browser window!
##
############################################################################################
 
 
__AUTHOR__ = 'SplishSplash'
 
PLUGIN_NAME = "Instruction_Reference"
PLUGIN_HOTKEY = 'Ctrl+I'
VERSION = '1.0'
 
import sys
import ida_kernwin
import idc
import idaapi
import idautils
import webbrowser
import requests
import re
 
major, minor = map(int, idaapi.get_kernel_version().split("."))
using_ida7api = (major > 6)
using_pyqt5 = using_ida7api or (major == 6 and minor >= 9)
 
if using_pyqt5:
    import PyQt5.QtGui as QtGui
    import PyQt5.QtCore as QtCore
    import PyQt5.QtWidgets as QtWidgets
    from PyQt5.Qt import QApplication
 
else:
    import PySide.QtGui as QtGui
    import PySide.QtCore as QtCore
 
    QtWidgets = QtGui
    QtCore.pyqtSignal = QtCore.Signal
    QtCore.pyqtSlot = QtCore.Slot
    from PySide.QtGui import QApplication
 
 
def PLUGIN_ENTRY():
    """
    Required plugin entry point for IDAPython Plugins.
    """
    return reference_init()
 
 
class reference_init(idaapi.plugin_t):
    """
    The IDA Plugin for Reference_Instruction.
    """
 
    flags = idaapi.PLUGIN_PROC | idaapi.PLUGIN_HIDE
    comment = "Reference Instruction"
    help = "Highlight Windows API and right-click 'Reference Instruction'"
    wanted_name = PLUGIN_NAME
    wanted_hotkey = PLUGIN_HOTKEY
 
    # --------------------------------------------------------------------------
    # Plugin Overloads
    # --------------------------------------------------------------------------
 
    def init(self):
        """
        This is called by IDA when it is loading the plugin.
        """
 
        # initialize the menu actions our plugin will inject
        self._init_action_reference_instruction()
 
        # initialize plugin hooks
        self._init_hooks()
 
        # done
        idaapi.msg("%s %s initialized...\n" % (self.wanted_name, VERSION))
        return idaapi.PLUGIN_KEEP
 
    def run(self, arg):
        """
        This is called by IDA when this file is loaded as a script.
        """
        idaapi.msg("%s cannot be run as a script.\n" % self.wanted_name)
 
    def term(self):
        """
        This is called by IDA when it is unloading the plugin.
        """
 
        # unhook our plugin hooks
        self._hooks.unhook()
 
        # unregister our actions & free their resources
        self._del_action_reference_instruction()
 
        # done
        idaapi.msg("%s terminated...\n" % self.wanted_name)
 
    # --------------------------------------------------------------------------
    # Plugin Hooks
    # --------------------------------------------------------------------------
 
    def _init_hooks(self):
        """
        Install plugin hooks into IDA.
        """
        self._hooks = Hooks()
        self._hooks.ready_to_run = self._init_hexrays_hooks
        self._hooks.hook()
 
    def _init_hexrays_hooks(self):
        """
        Install Hex-Rrays hooks (when available).
        NOTE: This is called when the ui_ready_to_run event fires.
        """
        if idaapi.init_hexrays_plugin():
            idaapi.install_hexrays_callback(self._hooks.hxe_callback)
 
    # --------------------------------------------------------------------------
    # IDA Actions
    # --------------------------------------------------------------------------
 
    REF_INSTR = "prefix:Reference_Instruction"
 
    def _init_action_reference_instruction(self):
        """
        Register the copy bytes action with IDA.
        """
        if (sys.version_info > (3, 0)):
            # Describe the action using python3 copy
            action_desc = idaapi.action_desc_t(
                self.REF_INSTR,  # The action name.
                "Reference Instruction",  # The action text.
                IDACtxEntry(reference_instr),  # The action handler.
                PLUGIN_HOTKEY,  # Optional: action shortcut
                "Reference x86/x64 Instruction",  # Optional: tooltip
                31  # Copy icon
            )
 
        # register the action with IDA
        assert idaapi.register_action(action_desc), "Action registration failed"
 
    def _del_action_reference_instruction(self):
        """
        Delete the bulk prefix action from IDA.
        """
        idaapi.unregister_action(self.REF_INSTR)
 
 
# ------------------------------------------------------------------------------
# Plugin Hooks
# ------------------------------------------------------------------------------
 
class Hooks(idaapi.UI_Hooks):
 
    def finish_populating_widget_popup(self, widget, popup):
        """
        A right click menu is about to be shown. (IDA 7)
        """
        inject_reference_actions(widget, popup, idaapi.get_widget_type(widget))
        return 0
 
    def finish_populating_tform_popup(self, form, popup):
        """
        A right click menu is about to be shown. (IDA 6.x)
        """
        inject_reference_actions(form, popup, idaapi.get_tform_type(form))
        return 0
 
    def hxe_callback(self, event, *args):
        """
        HexRays event callback.
        We lump this under the (UI) Hooks class for organizational reasons.
        """
 
        #
        # if the event callback indicates that this is a popup menu event
        # (in the hexrays window), we may want to install our prefix menu
        # actions depending on what the cursor right clicked.
        #
 
        if event == idaapi.hxe_populating_popup:
            form, popup, vu = args
 
            idaapi.attach_action_to_popup(
                form,
                popup,
                reference_init.REF_INSTR,
                "Reference Instruction",
                idaapi.SETMENU_APP,
            )
 
        # done
        return 0
 
 
# ------------------------------------------------------------------------------
# Prefix Wrappers
# ------------------------------------------------------------------------------
 
 
def inject_reference_actions(form, popup, form_type):
    """
    Inject prefix actions to popup menu(s) based on context.
    """
 
    #
    # disassembly window
    #
 
    if (form_type == idaapi.BWN_DISASMS):
        # insert the prefix action entry into the menu
        #
 
        idaapi.attach_action_to_popup(
            form,
            popup,
            reference_init.REF_INSTR,
            "Reference Instruction",
            idaapi.SETMENU_APP
        )
 
    # done
    return 0
 
 
# ------------------------------------------------------------------------------
# DO PLUGIN FUNCTIONALITY
# ------------------------------------------------------------------------------
 
def reference_instr():
    """
    Actual plugin functionality
    """
    if using_ida7api:
        data = ida_kernwin.get_highlight(ida_kernwin.get_current_viewer())
        print("%s" % data[0])
        page = 'https://www.felixcloutier.com/x86/' + data[0]
        webbrowser.open(page)
    else:
        data = idc.get_highlighted_identifier()
        print("%s" % data)
        page = 'https://www.felixcloutier.com/x86/' + data[0]
        webbrowser.open(page)
    return
 
 
# ------------------------------------------------------------------------------
# IDA ctxt
# ------------------------------------------------------------------------------
 
class IDACtxEntry(idaapi.action_handler_t):
    """
    A basic Context Menu class to utilize IDA's action handlers.
    """
 
    def __init__(self, action_function):
        idaapi.action_handler_t.__init__(self)
        self.action_function = action_function
 
    def activate(self, ctx):
        """
        Execute the embedded action_function when this context menu is invoked.
        """
        self.action_function()
        return 1
 
    def update(self, ctx):
        """
        Ensure the context menu is always available in IDA.
        """
        return idaapi.AST_ENABLE_ALWAYS


repository: HyperionGray/DriverBuddy-7.4-plus
file: DriverBuddy.py
description: A port of NCCGroupd DriverBuddy project to IDA 7.4
wanted_name: Driver Buddy
created_at: 2020-06-22 20:00:41+00:00
pushed_at: 2020-06-22 20:07:39+00:00
url:
https://github.com/HyperionGray/DriverBuddy-7.4-plus/blob/5a9a86d112d07e36797f73a3961271035cfd1b2b/DriverBuddy.py

content:
from idaapi import *
from idautils import *
from idc import *
from ida_auto import auto_wait
from DriverBuddy import data
from DriverBuddy import ioctl


'''#######################################################################################

DriverBuddy.py: Entry point for IDA python plugin used in Windows driver
                vulnerability research.

Written by Braden Hollembaek and Adam Pond of NCC Group
#######################################################################################'''
def PLUGIN_ENTRY():
    return DriverBuddyPlugin()

class DriverBuddyPlugin(plugin_t):
    flags = PLUGIN_UNL
    comment = ('Plugin to aid in Windows driver vulnerability research. ' +
               'Automatically tries to find IOCTL handlers, decode IOCTLS, '+
               'flag dangerous C/C++ functions, find Windows imports for privesc, '+
               'and identify the type of Windows driver.')
    help = ''
    wanted_name = 'Driver Buddy'
    wanted_hotkey = 'Ctrl-Alt-D'

    def init(self):
        self.hotkeys = []
        self.hotkeys.append(add_hotkey("Ctrl+Alt+I", self.decode))
        return PLUGIN_KEEP

    def run(self, args):
        print "[+] Welcome to Driver Buddy"
        auto_wait() # Wait for IDA autoanalysis to complete
        driver_entry = data.is_driver()
	if driver_entry == "":
            print "[-] No DriverEntry stub found"
            print "[-] Exiting..."
            return
        print "[+] DriverEntry found"
        if data.populate_data_structures() == False:
            print "[-] Unable to load functions"
            print "[-] Exiting..."
            return
	driver_type = data.get_driver_id(driver_entry)
        if driver_type == "":
            print "[-] Unable to determine driver type assuming wdm"
        else:
            print "[+] Driver type detected: " + driver_type
        if ioctl.find_ioctls() == False:
            print "[-] Unable to automatically find any IOCTLs"
        return

    def decode(self, _=0):
        if idc.get_operand_type(idc.get_screen_ea(), 1) != 5:   # Immediate
            return
        value = idc.get_operand_value(idc.get_screen_ea(), 1) & 0xffffffff
        ioctl.get_ioctl_code(value)

    def term(self):
        pass


repository: move47/Advanced-Topics-in-Malware-Analysis
file: Lab #3/def_use_plugin.py
description: Course material for CS 6747: Advanced Topics in Malware Analysis
wanted_name: My Python plugin
comment: This is a comment
created_at: 2023-08-22 14:40:01+00:00
pushed_at: 2023-08-23 14:51:56+00:00
url:
https://github.com/move47/Advanced-Topics-in-Malware-Analysis/blob/c3eefbd5a5a6427c7438bf8ab1fcc7840fb02c5f/Lab%20%233/def_use_plugin.py

content:
import idaapi
from cProfile import label
from re import I
import idautils
import idc
import idaapi
import pickle
from typing import Dict, List, Tuple

def check_call(ins):
    return idaapi.is_call_insn(ins)

def check_branches(ins):
    # reference: https://github.com/lifting-bits/mcsema-legacy/blob/master/tools/mcsema_disass/ida/get_cfg.py
    COND_BRANCHES = [\
    idaapi.NN_ja,\
    idaapi.NN_jae,\
    idaapi.NN_jb,\
    idaapi.NN_jbe,\
    idaapi.NN_jc,\
    idaapi.NN_jcxz,\
    idaapi.NN_je,\
    idaapi.NN_jecxz,\
    idaapi.NN_jg,\
    idaapi.NN_jge,\
    idaapi.NN_jl,\
    idaapi.NN_jle,\
    idaapi.NN_jna,\
    idaapi.NN_jnae,\
    idaapi.NN_jnb,\
    idaapi.NN_jnbe,\
    idaapi.NN_jnc,\
    idaapi.NN_jne,\
    idaapi.NN_jng,\
    idaapi.NN_jnge,\
    idaapi.NN_jnl,\
    idaapi.NN_jnle,\
    idaapi.NN_jno,\
    idaapi.NN_jnp,\
    idaapi.NN_jns,\
    idaapi.NN_jnz,\
    idaapi.NN_jo,\
    idaapi.NN_jp,\
    idaapi.NN_jpe,\
    idaapi.NN_jpo,\
    idaapi.NN_jrcxz,\
    idaapi.NN_js,\
    idaapi.NN_jz,]

    UCOND_BRANCHES = [\
        idaapi.NN_jmp,\
        idaapi.NN_jmpfi,\
        idaapi.NN_jmpni,\
        idaapi.NN_jmpshort]

    op = idautils.DecodeInstruction(ins).itype
    return op in COND_BRANCHES or op in UCOND_BRANCHES

def check_push(ins):
    insn_t = idautils.DecodeInstruction(ins)
    return insn_t.itype == idaapi.NN_push


def check_pop(ins):
    op = idautils.DecodeInstruction(ins).itype
    return op == idaapi.NN_pop

def check_ret(ins):
    op = idautils.DecodeInstruction(ins).itype
    return op == idaapi.NN_retn

def check_eflag_use(ins):
    if check_branches(ins):
        return True
    # TODO: add to this list
    FLAG_USE_CMDS = [
    ]
    op = idautils.DecodeInstruction(ins).itype
    return op in FLAG_USE_CMDS

def check_eflag_def(ins):
    # TODO: add to this list
    FLAG_DEF_CMDS = [
        idaapi.NN_add,
        idaapi.NN_sub,
        idaapi.NN_mul,
        idaapi.NN_div,
        idaapi.NN_xor,
        idaapi.NN_or,
        idaapi.NN_and,
        idaapi.NN_not,
        idaapi.NN_cmp,
        idaapi.NN_inc,
        idaapi.NN_dec,
    ]
    op = idautils.DecodeInstruction(ins).itype
    return op in FLAG_DEF_CMDS

class inst_info():
    def __init__(self, node_id, D_list, U_list, inst):
        self.D_list = D_list
        self.U_list = U_list
        self.inst = inst
        self.node_id = node_id
        self.parent_list = list()

def get_Defined_Used(ins):

    """
    for every operand in the instruction check if write or read flag is set.
    mark operand as used if read is set, mark it as defined if write flag is set
    reference: https://reverseengineering.stackexchange.com/questions/12053/ida-generic-approach-to-determine-if-an-instruction-reads-from-or-writes-to-m
    """

    D = list()
    U = list()

    OPND_WRITE_FLAGS = {
        0: idaapi.CF_CHG1,
        1: idaapi.CF_CHG2,
        2: idaapi.CF_CHG3,
        3: idaapi.CF_CHG4,
        4: idaapi.CF_CHG5,
        5: idaapi.CF_CHG6,
    }

    OPND_READ_FLAGS = {
        0: idaapi.CF_USE1,
        1: idaapi.CF_USE2,
        2: idaapi.CF_USE3,
        3: idaapi.CF_USE4,
        4: idaapi.CF_USE5,
        5: idaapi.CF_USE6,
    }

    insn = idautils.DecodeInstruction(ins)
    feature = insn.get_canon_feature()

    for op in insn.ops:
        if op.type == idaapi.o_void:
            break

        is_write = feature & OPND_WRITE_FLAGS[op.n]
        is_read = feature & OPND_READ_FLAGS[op.n]

        if(is_write):
            D.append(idc.print_operand(ins, op.n))

        if(is_read):
            U.append(idc.print_operand(ins, op.n))

    if check_push(ins):
        D += ['esp', '[esp]']
        U += ['esp']

    if check_pop(ins):
        D += ["esp"]
        U += ["esp", "[esp]"]

    if check_call(ins):
        D += ["eax", "esp"]
        U += ["esp"]

    if check_ret(ins):
        pass
        #nothing doing, ref does nothing

    if check_eflag_def(ins):
        D += ['eFlags']

    if check_eflag_use(ins):
        U += ['eFlags']

    return D,U

# Logic is mostly similar to DFS
def find_graph(block, parent, visited, parent_adj: Dict[str, inst_info], count,output_file):
    visited[block.start_ea]=1

    #First parsing current block instructions
    insts = idautils.Heads(block.start_ea, block.end_ea)

    for ins in insts:
        node_id = "n{}".format(count[0])
        parent_adj[hex(ins)] = {}
        D_list, U_list = get_Defined_Used(ins)
        parent_adj[hex(ins)] = inst_info(node_id, D_list, U_list, idc.GetDisasm(ins))
        parent_adj[hex(ins)].parent_list.append(parent)


        print("{0} [label=\"{1}; D: {2}; U: {3}\"]".format(node_id, hex(ins), D_list, U_list))
        output_file.write('{0} [label=\"{1}; D: {2}; U: {3}\"]\n'.format(node_id, hex(ins), D_list, U_list))
        parent = node_id
        count[0] = count[0] + 1

    # Now call for all the successive blocks
    for succ_block in block.succs():
        #print("outside",block)
        if succ_block.start_ea not in visited:
            find_graph(succ_block, parent, visited, parent_adj, count,output_file) # recursive call
        else:
            # gets_added_in_calle_first_ins_as_parent
            ins = succ_block.start_ea
            parent_adj[hex(ins)].parent_list.append(parent)
    return

def print_graph(adj_parent: Dict[str, inst_info],output_file):
    output_file.write("\n\n\n")
    for inst_addr in adj_parent:
        curr_node = adj_parent[inst_addr].node_id

        for parent in adj_parent[inst_addr].parent_list:
            print("{0} -> {1}".format(parent, curr_node))
            output_file.write("{0} -> {1}\n".format(parent, curr_node))
    output_file.write("\n}\n")

def main():
    print ("Start of the functions")
    for curr_function in idautils.Functions():
        #output_dir = "DOT_FILES"
        FUNCTION_NAME = idc.get_func_name(curr_function)
        #FUNCTION_NAME = "start"
        dot_file_name = "{}.dot".format(FUNCTION_NAME)
        #with open(dot_file_name,'w') as output_file
        output_file = open(dot_file_name,'w')
        output_file.write('digraph G {\n')
        print("\nCURRENT FUNCTION IS {0}".format(FUNCTION_NAME))
        for _ in range(0,2):
            print("**********************************************")
        #FUNCTION_NAME = 'start' #Change this to whatever func you want to look for. May have to change to _start in some cases or code your own logic

        cursor = 0
        real_names = []
        addresses = []

        adj_parent : Dict[str, inst_info]= {} # hex_address-->{hex addr: inst_info}
        count = [0]

        names = idautils.Names()

        # Populate the two above lists with their respective data from the list of tuples returned by idautils.Names()
        for n in names:
            real_names.append(n[1]) # real_names[30] ------> addresses[30]
            addresses.append(n[0])

        # search for a 'start' name in the names list and if it is there, set cursor to the address of beginning of it
        for rn in real_names:
            if rn == FUNCTION_NAME:
                cursor = addresses[int(real_names.index(FUNCTION_NAME))]
                break

        parent="start"
        visited={}

        if cursor: #if a start was found, assign it to the cursor and then
            myFunc = idaapi.get_func(cursor)
            f = idaapi.FlowChart(myFunc)
            for block in f:
                if block.start_ea not in visited:
                    find_graph(block, parent, visited, adj_parent, count,output_file)
                    print("Disconnected blocks")

            print_graph(adj_parent,output_file)

class myplugin_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = "This is a comment"
    help = "This is help"
    wanted_name = "My Python plugin"
    wanted_hotkey = "Alt-F8"

    def init(self):
        return idaapi.PLUGIN_OK

    def run(self, arg):
        main()
        print("Hello world!")

    def term(self):
        pass

def PLUGIN_ENTRY():
    return myplugin_t()


repository: thebabush/reversing-stuff
file: ida_plugins/hyperv_gs_proc.py
description: Reversing stuff.
wanted_name: hyperv_gs_proc
comment: Stolen from Rolf Rolles
created_at: 2019-02-12 01:59:55+00:00
pushed_at: 2019-02-18 13:16:22+00:00
url:
https://github.com/thebabush/reversing-stuff/blob/afb37ecd958e7c302518d4b83afe7465d9146419/ida_plugins/hyperv_gs_proc.py

content:
"""
@thebabush

# What is this all about?

Well, Hyper-V uses gs:XXXX to quickly access some important structures.
If you look at the end of `start`, you'll notice something like this:

```C
[...]
wrmsr(0xC0000101, some_addr);
[...]

That's MSR_GS_BASE and it's used to, guess what, tell the CPU where to look for
when you have something like `gs:0`.

IDA wouldn't let me force `gs`'s value so...

This hacky IDA Pro processor module is a quick way to turn
`mov reg, gs:XXXX` into `mov reg, qword ptr [rip+$(MSR_GS_BASE + XXXX - rip)]`.

I'm using capstone as a disassembler because I couldn't find a quick
way to get IDA to disasm some things for me.
This is probably my fault but whatever.

# How to use

1. Change the global `HVIX_GSBASE` to the value of your current Hyper-V binary.
2. Copy this script inside your ida `plugins/`.
3. Start ida64 from a virtualenv that has capstone in it.

# Warning

Once you put this script in `plugins/`, it's going to run for every binary you
open in IDA.
It's a quick thing to fix, but I haven't done it.

# Credits

All praise our Lord Rolf Rolles.
This is pretty much stolen from him :)

See: https://www.msreverseengineering.com/blog/2015/6/29/transparent-deobfuscation-with-ida-processor-module-extensions
"""

from __future__ import print_function

import os
import re

"""
RANT TIME:

    WHY THE FLYING FUCK DOES MOTHERFUCKING CAPSTONE COMES WITH THE GODDAMN
    DIET ENGINE WHEN YOU JFC INSTALL IT USING PIP FOR FUCK'S SAKE.

As you might have guessed... HUGE HACKS AHEAD :D
"""
import capstone

import ida_allins
import ida_offset
import ida_ua
import idaapi


# NOTE: CHANGE THIS
HVIX_GSBASE = 0xFFFFF80000603000

# Fuck capstone, really
X86_REGS = {
    "RAX": 0,
    "RCX": 1,
    "RDX": 2,
    "RBX": 3,
    "RSP": 4,
    "RBP": 5,
    "RSI": 6,
    "RDI": 7,
    "R8" : 8,
    "R9" : 9,
    "R10": 10,
    "R11": 11,
    "R12": 12,
    "R13": 13,
    "R14": 14,
    "R15": 15,
}


class hypervGsProcHook(idaapi.IDP_Hooks):
    def __init__(self):
        idaapi.IDP_Hooks.__init__(self)
        self.md = capstone.Cs(capstone.CS_ARCH_X86, capstone.CS_MODE_64)
        # What could possibly go wrong...
        # Thanks, capstone diet and my ignorance of IDA scripting
        self.re = re.compile(r'([re][0-9a-dpxsi]+), ([a-z]+) ptr gs:\[([0-9a-fx]+)\]')

    def print_instr(self, instr):
        print('WTF: 0x{:08X} {} {}'.format(
            idaapi.cmd.ea,
            instr.mnemonic,
            instr.op_str
        ))

    def ev_ana_insn(self, insn):
        addr = insn.ea

        opcodes = idaapi.get_bytes(addr, 16)
        instrs = self.md.disasm(opcodes, addr)

        for instr in instrs:
            if instr.mnemonic != 'mov':
                if instr.op_str.find('gs') != -1:
                    self.print_instr(instr)
                return False

            m = self.re.search(instr.op_str)
            if m is None:
                return False

            reg, size, offset = m.groups()
            offset = int(offset, 16)
            # What could possibly go wrong #2...
            size = getattr(ida_ua, 'dt_' + size)

            reg = reg.upper()
            if reg[0] == 'E':
                reg = 'R' + reg[1:]
            if reg[-1] == 'D':
                reg = reg[:-1]
            if reg not in X86_REGS:
                print('Unknown reg: {}'.format(reg))
                return False
            reg = X86_REGS[reg]

            insn.itype = ida_allins.NN_mov

            # Seriously...
            # Seems to make the stuff RIP-relative
            insn.auxpref = 0x1810   # addressing mode not overwritten by prefix,
                                    # operand size not overwritten by prefix,
                                    # segment type is 64-bit
            # Seems to use cs instead of ds
            insn.Op2.specval = 0x1e0000

            insn.Op1.type = idaapi.o_reg
            insn.Op1.dtype = size
            insn.Op1.reg = reg
            insn.Op2.dtype = size
            insn.Op2.type = idaapi.o_mem
            insn.Op2.addr = HVIX_GSBASE + offset
            insn.size = instr.size
            return True

        return False


class hyperv_gs_proc_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_PROC | idaapi.PLUGIN_HIDE
    comment = "Stolen from Rolf Rolles"
    wanted_hotkey = ""
    help = "Itsame, a stolen script"
    wanted_name = "hyperv_gs_proc"
    hook = None

    def init(self):
        self.hook = None
        hv_name = os.path.split(idaapi.get_input_file_path())[1].lower()
        if hv_name not in ('hvix64.exe', 'hvax64.exe') or idaapi.ph_get_id() != idaapi.PLFM_386:
            return idaapi.PLUGIN_SKIP
        self.hook = hypervGsProcHook()
        self.hook.hook()
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        pass

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


def PLUGIN_ENTRY():
    return hyperv_gs_proc_t()


repository: hellow554/RustyBlazer
file: tools/ida_plugins/sb_helper.py
description: This repo will contain a complete decompilation of the Soul Blazer game by Square Enix
wanted_name: SbHelper
comment: Soulblazer Analyzer (IDAPython)
created_at: 2022-09-21 06:44:32+00:00
pushed_at: 2025-04-01 06:05:03+00:00
url:
https://github.com/hellow554/RustyBlazer/blob/699824e0c418bba874535c807855ee9ceb303c72/tools/ida_plugins/sb_helper.py

content:
"""
summary: provide two helpers for the soulblazer game

description:
    This script implements two things:
        1. a helper to disassemble text pointer
        2. a handy plugin that inserts the COP data directly after a COP instruction has been decoded
"""

import ida_idaapi
import ida_auto
import ida_bytes
import ida_hexrays
import idc
import ida_idp
import ida_kernwin
import ida_struct
import ida_ua
import ida_segment
import ida_segregs
import ida_typeinf
import ida_xref
import idautils
import ida_nalt

from typing import TypeVar

R = TypeVar("R")


def dbg(x: R) -> R:
    print(x)
    return x


ASM_COP = 0x02
U64_MAX = 18446744073709551615


class MyHook(ida_idp.IDP_Hooks):
    def __init__(self):
        ida_idp.IDP_Hooks.__init__(self)

    def ev_emu_insn(self, insn: ida_ua.insn_t):
        byte = ida_bytes.get_byte(insn.ea)
        if byte == ASM_COP:
            next_byte = ida_bytes.get_byte(insn.ea + 1)
            struc_name = f"Cop{next_byte:02X}"
            struc_tid = ida_struct.get_struc_id(struc_name)
            if struc_tid != 0 and struc_tid != U64_MAX:
                struc_len = ida_struct.get_struc_size(struc_tid)
                ida_bytes.create_struct(insn.ea + 2, struc_len, struc_tid, True)
                ida_auto.auto_mark(insn.ea + 2 + struc_len, ida_auto.AU_CODE)
            else:
                print(f"Can't find a struct with the name {struc_name}")

        return 0


class Xref24(ida_kernwin.action_handler_t):
    def __init__(self):
        ida_kernwin.action_handler_t.__init__(self)

    def update(self, _: ida_kernwin.action_ctx_base_t):
        return ida_kernwin.AST_ENABLE

    def activate(self, ctx: ida_kernwin.action_ctx_base_t):
        cs = self._get_ptr_struct()
        self._recurse(cs)

    def _get_ptr_struct(self) -> ida_struct.struc_t | None:
        cs_id = ida_struct.get_struc_id("ptr24_t")
        if cs_id == -1 or cs_id == U64_MAX:
            print("Can't find `ptr24_t`")
            return None

        cs: ida_struct.struc_t = ida_struct.get_struc(cs_id)
        if (
            cs.memqty == 2
            and ida_struct.get_member_size(cs.get_member(0)) == 2
            and ida_struct.get_member_size(cs.get_member(1)) == 1
        ):
            return cs

        print("Error: struct `ptr24_t` is not valid!")
        return None

    def _recurse(self, struct, struct_offset=0):
        """
        either it's a xref to a struct (another struct containing this struct)
        or it's a ref to some data
        if struct, we call ourself recursivly
        if data, we create an xref
        """

        struct_size = ida_struct.get_struc_size(struct.id)
        for x_to in idautils.XrefsTo(struct.id):
            ea = x_to.frm
            flags = ida_bytes.get_full_flags(ea)
            elem = ida_bytes.get_item_size(x_to.frm) // struct_size

            if ida_bytes.is_data(flags):
                # we are in data section
                for x in range(elem):
                    frm = ea + x * struct_size  # + struct_offset
                    real_from = ea + x * struct_size + struct_offset
                    bank = ida_bytes.get_byte(real_from + 2)
                    if bank <= 0x20:
                        bank += 0x80
                    to = (bank << 16) + ida_bytes.get_word(real_from)

                    ida_xref.add_dref(frm, to, ida_xref.dr_R)

            elif flags == 0:
                (member, _name, struct) = ida_struct.get_member_by_id(ea)
                self._recurse(struct, member.get_soff())

            else:
                print("NOT SURE WHAT HAPPENS!", hex(flags))

class PrintSegmentRegister(ida_kernwin.action_handler_t):
    def __init__(self):
        ida_kernwin.action_handler_t.__init__(self)

    def update(self, ctx: ida_kernwin.action_ctx_base_t):
        return ida_kernwin.AST_ENABLE_FOR_WIDGET \
            if ctx.widget_type == ida_kernwin.BWN_DISASM \
                else ida_kernwin.AST_DISABLE_FOR_WIDGET

    def activate(self, ctx: ida_kernwin.action_ctx_base_t):
        # 8 = M
        # 9 = X
        m = ida_segregs.get_sreg(ctx.cur_ea, 8)
        x = ida_segregs.get_sreg(ctx.cur_ea, 9)
        print(".A{:<2} .I{:<2}".format(8 if m == 1 else 16, 8 if x == 1 else 16))

class SbHelperPlugin(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_HIDE
    wanted_name = "SbHelper"
    wanted_hotkey = ""
    comment = "Soulblazer Analyzer (IDAPython)"
    help = ""

    def init(self):
        self.sb_hook = MyHook()
        self.sb_hook.hook()

        ida_kernwin.register_action(
            ida_kernwin.action_desc_t(
                "xref24:addXref",
                "Xref24 Add X-References",
                Xref24(),
                "Shift+Ctrl+Z",
            )
        )

        ida_kernwin.register_action(
            ida_kernwin.action_desc_t(
                "printseg:printSegmentRegister",
                "Print current Segment Register",
                PrintSegmentRegister(),
                "Shift+Ctrl+A",
            )
        )

        return ida_idaapi.PLUGIN_KEEP  # keep us in memory

    def run(self, arg):
        pass

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


def PLUGIN_ENTRY():
    return SbHelperPlugin()


repository: Russinovich/IDA7.5_SDK
file: python/examples/hexrays/vds3.py
description: SDK and IDAPython for IDA Pro 7.5
wanted_name: Hex-Rays if-inverter (IDAPython)
comment: Sample plugin3 for Hex-Rays decompiler
created_at: 2021-09-01 02:18:48+00:00
pushed_at: 2021-07-23 08:19:09+00:00
url:
https://github.com/Russinovich/IDA7.5_SDK/blob/5244a2b5059113e421b4da4ecf771013a6fde2f8/python/examples/hexrays/vds3.py

content:
""" Invert the then and else blocks of a cif_t.

Author: EiNSTeiN_ <einstein@g3nius.org>

This is a rewrite in Python of the vds3 example that comes with hexrays sdk.
"""
from __future__ import print_function

import idautils

import ida_kernwin
import ida_hexrays
import ida_netnode
import ida_idaapi

import traceback

NETNODE_NAME = '$ hexrays-inverted-if'

inverter_actname = "vds3:invert"

class invert_action_handler_t(ida_kernwin.action_handler_t):
    def __init__(self, inverter):
        ida_kernwin.action_handler_t.__init__(self)
        self.inverter = inverter

    def activate(self, ctx):
        vdui = ida_hexrays.get_widget_vdui(ctx.widget)
        self.inverter.invert_if_event(vdui)
        return 1

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


class hexrays_callback_info(object):

    def __init__(self):
        self.vu = None

        self.node = ida_netnode.netnode()
        if not self.node.create(NETNODE_NAME):
            # node exists
            self.load()
        else:
            self.stored = []

        return

    def load(self):

        self.stored = []

        try:
            data = self.node.getblob(0, 'I')
            if data:
                self.stored = eval(data)
                print('Invert-if: Loaded %s' % (repr(self.stored), ))
        except:
            print('Failed to load invert-if locations')
            traceback.print_exc()
            return

        return

    def save(self):

        try:
            self.node.setblob(repr(self.stored), 0, 'I')
        except:
            print('Failed to save invert-if locations')
            traceback.print_exc()
            return

        return

    def invert_if(self, cfunc, insn):

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

        cif = insn.details

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

        ida_hexrays.qswap(cif.ithen, cif.ielse)
        # Make a copy of 'cif.expr': 'lnot' might destroy its toplevel
        # cexpr_t and return a pointer to its direct child (but we'll want to
        # 'swap' it later, the 'cif.expr' cexpr_t object must remain valid.)
        cond = ida_hexrays.cexpr_t(cif.expr)
        notcond = ida_hexrays.lnot(cond)

        cif.expr.swap(notcond)

        return True

    def add_location(self, ea):
        if ea in self.stored:
            self.stored.remove(ea)
        else:
            self.stored.append(ea)
        self.save()
        return

    def find_if_statement(self, vu):

        vu.get_current_item(ida_hexrays.USE_KEYBOARD)
        item = vu.item

        if item.is_citem() and item.it.op == ida_hexrays.cit_if and item.it.to_specific_type.cif.ielse is not None:
            return item.it.to_specific_type

        if vu.tail.citype == ida_hexrays.VDI_TAIL and vu.tail.loc.itp == ida_hexrays.ITP_ELSE:
            # for tail marks, we know only the corresponding ea,
            # not the pointer to if-statement
            # find it by walking the whole ctree
            class if_finder_t(ida_hexrays.ctree_visitor_t):
                def __init__(self, ea):
                    ida_hexrays.ctree_visitor_t.__init__(self, ida_hexrays.CV_FAST | ida_hexrays.CV_INSNS)

                    self.ea = ea
                    self.found = None
                    return

                def visit_insn(self, i):
                    if i.op == ida_hexrays.cit_if and i.ea == self.ea:
                        self.found = i
                        return 1 # stop enumeration
                    return 0

            iff = if_finder_t(vu.tail.loc.ea)
            if iff.apply_to(vu.cfunc.body, None):
                return iff.found

        return

    def invert_if_event(self, vu):

        cfunc = vu.cfunc.__deref__()
        i = self.find_if_statement(vu)
        if not i:
            return False

        if self.invert_if(cfunc, i):
            vu.refresh_ctext()
            self.add_location(i.ea)

        return True

    def restore(self, cfunc):

        class visitor(ida_hexrays.ctree_visitor_t):

            def __init__(self, inverter, cfunc):
                ida_hexrays.ctree_visitor_t.__init__(self, ida_hexrays.CV_FAST | ida_hexrays.CV_INSNS)
                self.inverter = inverter
                self.cfunc = cfunc
                return

            def visit_insn(self, i):
                try:
                    if i.op == ida_hexrays.cit_if and i.ea in self.inverter.stored:
                        self.inverter.invert_if(self.cfunc, i)
                except:
                    traceback.print_exc()
                return 0 # continue enumeration

        visitor(self, cfunc).apply_to(cfunc.body, None)

        return


class vds3_hooks_t(ida_hexrays.Hexrays_Hooks):
    def __init__(self, i):
        ida_hexrays.Hexrays_Hooks.__init__(self)
        self.i = i

    def populating_popup(self, widget, phandle, vu):
        ida_kernwin.attach_action_to_popup(vu.ct, None, inverter_actname)
        return 0

    def maturity(self, cfunc, maturity):
        if maturity == ida_hexrays.CMAT_FINAL:
            self.i.restore(cfunc)
        return 0

# a plugin interface, boilerplate code
class my_plugin_t(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_HIDE
    wanted_name = "Hex-Rays if-inverter (IDAPython)"
    wanted_hotkey = ""
    comment = "Sample plugin3 for Hex-Rays decompiler"
    help = ""
    def init(self):
        if ida_hexrays.init_hexrays_plugin():
            i = hexrays_callback_info()
            ida_kernwin.register_action(
                ida_kernwin.action_desc_t(
                    inverter_actname,
                    "Invert then/else",
                    invert_action_handler_t(i),
                    "I"))
            self.vds3_hooks = vds3_hooks_t(i)
            self.vds3_hooks.hook()
            return ida_idaapi.PLUGIN_KEEP # keep us in the memory
    def term(self):
        self.vds3_hooks.unhook()
    def run(self, arg):
        pass

def PLUGIN_ENTRY():
    return my_plugin_t()


repository: techbliss/ADB_Helper_QT_Super_version
file: plugins/ADB_Helper_QT_Super.py
description: All You Need For Ida Pro And Android Debugging
wanted_name: ADB PyQT ADB connector for android server & gdbserver
comment: This is a comment
created_at: 2015-01-17 22:33:12+00:00
pushed_at: 2015-04-05 07:30:50+00:00
url:
https://github.com/techbliss/ADB_Helper_QT_Super_version/blob/3944b1a2dc1707bb1bbb53d017d44096af35afb8/plugins/ADB_Helper_QT_Super.py

content:
__author__ = 'zadow aka Storm Shadow'
import re
import idaapi
import idc
from idc import *
from idaapi import *
import PyQt4
from PyQt4 import QtCore, QtGui
import idautils

class Adb(idaapi.plugin_t):
    flags = idaapi.PLUGIN_FIX
    comment = "This is a comment"

    help = "Plugin "
    wanted_name = "ADB PyQT ADB connector for android server & gdbserver"
    wanted_hotkey = "Alt-F7"



    def init(self):
        idaapi.msg(" ADB Plugin is found. \n")
        return idaapi.PLUGIN_OK

    def run(self, arg):
        idaapi.msg("run() called with %d!\n" % arg)

    def term(self):
        idaapi.msg("")

    def AddMenuElements(self):
        '''Menus are better than no GUI at all *sigh*'''

        idaapi.add_menu_item("Edit/Plugins/", "ADB Super Connector", "", 0, self.AdbCall, ())




    def run(self, arg = 0):
        idaapi.msg("Looks like its working")

        self.AddMenuElements()

    def AdbCall(self):
        g = globals()
        idahome = idaapi.idadir("QTApps\\SuperADB")
        IDAPython_ExecScript(idahome +  "\\ADB_QT.py", g)



def PLUGIN_ENTRY():
    return Adb()


repository: NoneShell/IDADelogger
file: idadelog_plugin.py
wanted_name: Log Analyzer
comment: Analyze log calls and rename functions
created_at: 2025-04-29 10:10:32+00:00
pushed_at: 2025-04-29 10:19:03+00:00
description:
A GUI plugin for IDA Pro that automatically analyzes logging function calls and renames functions based on extracted log message parameters.

url:
https://github.com/NoneShell/IDADelogger/blob/41cb0c453a08c28f2269ac631ba155b76f894973/idadelog_plugin.py

content:
import idaapi
import idautils
import idc
import re
from PyQt5 import QtWidgets, QtCore
from collections import defaultdict

class LogAnalyzerDialog(QtWidgets.QDialog):
    def __init__(self, parent=None):
        super(LogAnalyzerDialog, self).__init__(parent)
        self.setWindowTitle("Log Analyzer")
        self.setWindowFlags(self.windowFlags() & ~QtCore.Qt.WindowContextHelpButtonHint)
        self.resize(600, 400)
        
        # save analysis results
        self.analysis_results = defaultdict(list)
        self.current_matches = []
        
        self.setup_ui()
        
    def setup_ui(self):
        layout = QtWidgets.QVBoxLayout()
        
        # input area
        input_group = QtWidgets.QGroupBox("Analysis Parameters")
        input_layout = QtWidgets.QFormLayout()
        
        self.log_func_edit = QtWidgets.QLineEdit("log_function_name")
        self.arg_index_spin = QtWidgets.QSpinBox()
        self.arg_index_spin.setRange(0, 20)
        self.arg_index_spin.setValue(4)
        
        self.regex_check = QtWidgets.QCheckBox("Enable Regex Transformation")
        self.pattern_edit = QtWidgets.QLineEdit()
        self.replace_edit = QtWidgets.QLineEdit()
        
        input_layout.addRow("Log Function Name:", self.log_func_edit)
        input_layout.addRow("Argument Index (0-based):", self.arg_index_spin)
        input_layout.addRow(self.regex_check)
        input_layout.addRow("Pattern:", self.pattern_edit)
        input_layout.addRow("Replacement:", self.replace_edit)
        
        input_group.setLayout(input_layout)
        
        # preview area
        result_group = QtWidgets.QGroupBox("Preview Results")
        result_layout = QtWidgets.QVBoxLayout()
        
        self.result_table = QtWidgets.QTableWidget()
        self.result_table.setColumnCount(4)
        self.result_table.setHorizontalHeaderLabels(["Address", "Old Name", "New Name", "Log Message"])
        self.result_table.setSelectionBehavior(QtWidgets.QAbstractItemView.SelectRows)
        self.result_table.setEditTriggers(QtWidgets.QAbstractItemView.NoEditTriggers)
        self.result_table.horizontalHeader().setStretchLastSection(True)
        self.result_table.setSortingEnabled(True)
        
        result_layout.addWidget(self.result_table)
        result_group.setLayout(result_layout)
        
        # button area
        button_layout = QtWidgets.QHBoxLayout()
        self.test_button = QtWidgets.QPushButton("Test")
        self.apply_button = QtWidgets.QPushButton("Apply")
        self.close_button = QtWidgets.QPushButton("Close")
        
        self.apply_button.setEnabled(False)
        
        button_layout.addWidget(self.test_button)
        button_layout.addWidget(self.apply_button)
        button_layout.addWidget(self.close_button)
        
        # main layout
        layout.addWidget(input_group)
        layout.addWidget(result_group)
        layout.addLayout(button_layout)
        
        self.setLayout(layout)
        
        # connect buttons to functions
        self.test_button.clicked.connect(self.run_test)
        self.apply_button.clicked.connect(self.apply_changes)
        self.close_button.clicked.connect(self.close)
        self.regex_check.stateChanged.connect(self.toggle_regex_fields)
        
        # initially disable regex fields
        self.toggle_regex_fields()
    
    def toggle_regex_fields(self):
        enabled = self.regex_check.isChecked()
        self.pattern_edit.setEnabled(enabled)
        self.replace_edit.setEnabled(enabled)
    
    def run_test(self, apply=False):
        # run the analysis, use apply=False to just test
        log_func = self.log_func_edit.text().strip()
        arg_index = self.arg_index_spin.value()
        
        if not log_func:
            QtWidgets.QMessageBox.warning(self, "Error", "Please specify a log function name")
            return
        
        # get regex pattern and replacement
        rename_regex = None
        if self.regex_check.isChecked():
            pattern = self.pattern_edit.text().strip()
            replacement = self.replace_edit.text().strip()
            if pattern:
                rename_regex = (pattern, replacement)
        
        # run the analysis
        self.analysis_results.clear()
        self.current_matches = []
        
        try:
            analyzer = LogAnalyzer(log_func, arg_index, rename_regex, apply)
            analyzer.analyze()
            self.current_matches = analyzer.match_results
        except Exception as e:
            print(f"[EXCEPTION] Error : {e}")
            self.apply_button.setEnabled(False)
            
        
        # update the results table
        self.update_results_table()
        self.apply_button.setEnabled(len(self.current_matches) > 0)
    
    def update_results_table(self):
        # update the results table with the current matches
        self.result_table.setRowCount(len(self.current_matches))
        self.result_table.clearContents()
        
        for row, match in enumerate(self.current_matches):
            self.result_table.setItem(row, 0, QtWidgets.QTableWidgetItem(hex(match['address'])))
            self.result_table.setItem(row, 1, QtWidgets.QTableWidgetItem(match['old_name']))
            self.result_table.setItem(row, 2, QtWidgets.QTableWidgetItem(match['new_name']))
            self.result_table.setItem(row, 3, QtWidgets.QTableWidgetItem(match['log_message']))
        
        self.result_table.resizeColumnsToContents()
    
    def apply_changes(self):
        # apply the changes to rename functions
        print("[INFO] Applying changes...")
        if not self.current_matches:
            return
            
        renamed = 0
        for match in self.current_matches:
            if idc.set_name(match['address'], match['new_name']):
                renamed += 1
        
        QtWidgets.QMessageBox.information(
            self, 
            "Complete", 
            f"Successfully renamed {renamed}/{len(self.current_matches)} functions"
        )
        
        # update the results table
        self.run_test(apply=True)


class LogAnalyzer:
    def __init__(self, log_func_name, arg_index=1, rename_regex=None, apply=False):
        self.log_func_name = log_func_name
        self.arg_index = arg_index
        self.rename_regex = rename_regex
        self.log_func_addr = idc.get_name_ea_simple(log_func_name)
        self.result = {}
        self.match_results = []
        self.apply = apply

        if self.log_func_addr == idc.BADADDR:
            raise ValueError(f"[ERROR] Could not find function: {log_func_name}")
        print(f"[INFO] Found log function '{log_func_name}' at {hex(self.log_func_addr)}")

    def analyze(self):
        """
        Main analysis routine: find call sites, extract args, rename functions.
        """
        callers = self.find_callers()
        for func_addr, old_function_name in callers:
            if old_function_name not in self.result:
                self.result[old_function_name] = []
            extracted_args = self.extract_log_arguments(func_addr)
            # print(f"[INFO] Extracted args from {old_function_name}: {extracted_args}")
            if extracted_args:
                # Use the specified argument index to get the relevant argument
                if self.arg_index < len(extracted_args):
                    new_function_name = extracted_args[self.arg_index]
                    self.result[old_function_name].append(new_function_name)
                    # print(f"[INFO] Raw name from {old_function_name}: {new_function_name}")
            else:
                print(f"[WARNING] No arguments found for {old_function_name}")
            
        
        # rename functions by results
        for old_function_name, args in self.result.items():
            most_common_function_name = max(set(args), key=args.count)
            sanitized_name = self.sanitize_name(most_common_function_name)
            
            if old_function_name not in self.match_results:
                self.match_results.append({
                    'address': idc.get_name_ea_simple(old_function_name),
                    'old_name': old_function_name,
                    'new_name': sanitized_name,
                    'log_message': str(args)
                })
            if self.apply:
                self.rename_function(old_function_name, args[0])
        return self.match_results


        

    def find_callers(self):
        """
        Find all functions calling the log function.
        """
        callers = []
        for ref in idautils.CodeRefsTo(self.log_func_addr, 0):
            func_start = idc.get_func_attr(ref, idc.FUNCATTR_START)
            if func_start == idc.BADADDR:
                continue
            func_name = idc.get_func_name(func_start)
            callers.append((func_start, func_name))
            # print(f"[CALLER] Found call to {self.log_func_name} in {func_name} at {hex(ref)}")
        return callers

    def extract_log_arguments(self, func_addr):
        """
        Enhanced version that better handles multi-line log calls and complex arguments
        """
        try:
            cfunc = idaapi.decompile(func_addr)
            if not cfunc:
                # print(f"[ERROR] Could not decompile function at {hex(func_addr)}")
                return []

            decompiled_code = str(cfunc)
            log_call_start = self.log_func_name + '('
            log_call_len = len(log_call_start)
            
            start_pos = 0
            while True:
                # Find next log call
                call_pos = decompiled_code.find(log_call_start, start_pos)
                if call_pos == -1:
                    break
                    
                # Initialize parsing state
                in_string = False
                escape = False
                paren_depth = 1
                current_arg = []
                args = []
                i = call_pos + log_call_len  # Position after '('
                
                # Parse until matching closing paren
                while i < len(decompiled_code) and paren_depth > 0:
                    char = decompiled_code[i]
                    
                    # Handle strings and escapes
                    if char == '"' and not escape:
                        in_string = not in_string
                        current_arg.append(char)
                    elif in_string and char == '\\':
                        escape = not escape
                        current_arg.append(char)
                    elif in_string:
                        escape = False
                        current_arg.append(char)
                    # Handle parentheses
                    elif char == '(':
                        paren_depth += 1
                        current_arg.append(char)
                    elif char == ')':
                        paren_depth -= 1
                        if paren_depth > 0:
                            current_arg.append(char)
                    # Handle argument separators
                    elif char == ',' and paren_depth == 1:
                        args.append(''.join(current_arg).strip())
                        current_arg = []
                        i += 1  # Skip whitespace after comma
                        while i < len(decompiled_code) and decompiled_code[i].isspace():
                            i += 1
                        continue  # Don't add comma to current_arg
                    else:
                        current_arg.append(char)
                    
                    i += 1
                
                # Add the last argument
                if current_arg:
                    # print(f"[INFO] Found args: {args}")
                    args.append(''.join(current_arg).strip())
                
                start_pos = i
            # print(f"[INFO] Found args: {args}")
            return args

        except Exception as e:
            # print(f"[EXCEPTION] Error in function {hex(func_addr)}: {e}")
            return []
        
    def sanitize_name(self, name):
        """Apply regex substitution and sanitize to IDA-valid function name"""
        original = name
        name = name.replace("'", "").replace('"', "")

        if self.rename_regex:
            try:
                pattern, repl = self.rename_regex
                name = re.sub(pattern, repl, name)
                print(f"[TRANSFORM] Regex applied: '{original}' -> '{name}'")
            except Exception as e:
                print(f"[ERROR] Regex error: {e}")
        
        return re.sub(r"[^a-zA-Z0-9_]", "_", name)
    def rename_function(self, func_addr, new_function_name):
        """
        Rename the function using extracted argument strings and optional regex transformation.
        """
        old_function_name = idc.get_func_name(func_addr)
        if not new_function_name:
            return

        if new_function_name and old_function_name != new_function_name:
            success = idc.set_name(func_addr, new_function_name)
            if success:
                print(f"[INFO] Renamed {old_function_name} -> {new_function_name}")
            else:
                print(f"[ERROR] Rename failed for {old_function_name} -> {new_function_name}")
        else:
            print(f"[SKIP] No rename needed for {old_function_name}")

class LogAnalyzerPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    wanted_name = "Log Analyzer"
    wanted_hotkey = ""
    comment = "Analyze log calls and rename functions"
    help = "This plugin analyzes log function calls and renames functions based on log messages"

    def init(self):
        # create action
        self.action_desc = idaapi.action_desc_t(
            'loganalyzer:action',
            'Log Analyzer',
            self.run,
            None,
            'Analyze log calls and rename functions',
            199
        )
        
        # register action
        if not idaapi.register_action(self.action_desc):
            return idaapi.PLUGIN_SKIP
        
        # add action to menu
        if not idaapi.attach_action_to_menu(
            "Edit/",
            'loganalyzer:action',
            idaapi.SETMENU_APP
        ):
            return idaapi.PLUGIN_SKIP
        
        return idaapi.PLUGIN_OK

    def run(self, arg):
        # create and show the dialog
        dialog = LogAnalyzerDialog()
        dialog.exec_()

    def term(self):
        # unregister action
        idaapi.unregister_action('loganalyzer:action')

def PLUGIN_ENTRY():
    return LogAnalyzerPlugin()


repository: idiom/stackstack
file: src/stackstack.py
url: https://github.com/idiom/stackstack/blob/4567acaf21812a48c32b11f4da3cb0329a73d526/src/stackstack.py
wanted_name: StackStack
comment: StackStack simple emulation and tracing
created_at: 2021-04-15 23:35:38+00:00
pushed_at: 2022-04-24 20:55:05+00:00
content:
import ida_idaapi
import ida_kernwin
import idaapi
import ida_ua
import ida_diskio
import idautils

import idc
import os
import json
import logging

from stackstack.scan import YaraScanner
from stackstack.sue import SUE
from stackstack.utils import IdaHelpers, Update

BAD = [0xffffffff, 0xffffffffffffffff]


class StackStack(object):

    def __init__(self, logger):

        self.logger = logger

        self.last_bookmark = 0
        self.arch = IdaHelpers.get_arch()

    def find_end(self, offset):
        function_end = idc.get_func_attr(offset, idc.FUNCATTR_END)
        self.logger.debug('Function End:   %x' % function_end)

        found_compare = False

        while offset <= function_end:
            ins = ida_ua.insn_t()
            idaapi.decode_insn(ins, offset)

            self.logger.debug('%x: %s' % (offset, idc.generate_disasm_line(offset, 0)))

            if ins.itype == idaapi.NN_jmp:
                return idc.next_head(offset, function_end)
            elif ins.itype == idaapi.NN_call:
                return idc.next_head(offset, function_end)
            elif ins.itype in [idaapi.NN_jmp, idaapi.NN_jnz, idaapi.NN_jb]:
                if found_compare:
                    return idc.next_head(offset, function_end)
            elif ins.itype == idaapi.NN_cmp:
                if idc.get_operand_type(offset, 0) == idc.o_reg \
                        and idc.get_operand_type(offset, 1) == idc.o_imm:
                    try:
                        self.logger.debug("Found compare..")
                        found_compare = True
                        string_length = int(idc.print_operand(offset, 1)) + 1
                        self.logger.debug("Setting String length: %d" % string_length)
                    except ValueError:
                        string_length = int(idc.print_operand(offset, 1)[:-1], 16) + 1
                        self.logger.debug("Setting String length: %d" % string_length)
            elif ins.itype == idaapi.NN_dec:
                if idc.get_operand_type(offset, 0) == idc.o_reg:
                    found_compare = True
            offset = idc.next_head(offset, function_end)
        return 0

    def get_string_length(self, offset):
        function_end = idc.get_func_attr(offset, idc.FUNCATTR_END)
        self.logger.debug('Function End:   %x' % function_end)

        while offset <= function_end:
            ins = ida_ua.insn_t()
            idaapi.decode_insn(ins, offset)

            if ins.itype == idaapi.NN_jmp:
                return 0
            elif ins.itype == idaapi.NN_call:
                return 0
            elif ins.itype == idaapi.NN_cmp:
                if idc.get_operand_type(offset, 0) == idc.o_reg \
                        and idc.get_operand_type(offset, 1) == idc.o_imm:
                    try:
                        self.logger.debug("Found compare..")
                        return int(idc.print_operand(offset, 1)) + 1
                    except ValueError:
                        return int(idc.print_operand(offset, 1)[:-1], 16) + 1
            offset = idc.next_head(offset, function_end)

        return 0

    def detect_blocks(self, offset, trace_end=False):
        """
        Use ida basic blocks to detect last block etc.

        :param offset:
        :param trace_end:
        :return:
        """
        pass

    def _has_call(self, start, end):

        while start < end:
            ins = ida_ua.insn_t()
            idaapi.decode_insn(ins, start)
            if ins.itype == idaapi.NN_call:
                self.logger.debug("Found call at 0x%x..skipping using function start." % start)
                return True
            start = idc.next_head(start, end)
        return False

    def scan_reg_increment(self, blob_start, offset):
        """

        Check if the block uses a register to increment the counter. Some samples will have something similar to

        <snip>
        xor     r15d, r15d
        <snip>
        lea     r12d, [r15+1]
        <snip>
        add     rax, r12

        :return: None if not found or identified register
        """

        function_start = idc.get_func_attr(offset, idc.FUNCATTR_START)
        cursor = 0

        found_compare = False
        found_operand = False
        cmp_reg = None
        inc_reg = None
        while offset >= blob_start:

            if offset <= function_start:
                return None

            cursor += 1
            ins = ida_ua.insn_t()
            idaapi.decode_insn(ins, offset)

            self.logger.debug('%x: %s' % (offset, idc.generate_disasm_line(offset, 0)))

            if not found_compare:
                if ins.itype == idaapi.NN_cmp:
                    if idc.get_operand_type(offset, 0) == idc.o_reg:
                        cmp_reg = idc.get_operand_value(offset, 0)
                        found_compare = True

            if found_compare:
                if ins.itype == idaapi.NN_add:
                    if idc.get_operand_type(offset, 0) == idc.o_reg \
                            and idc.get_operand_value(offset, 0) == cmp_reg \
                            and idc.get_operand_type(offset, 1) == idc.o_reg:
                                    #inc_reg = idc.get_operand_value(offset, 1)
                                    inc_reg = idc.print_operand(offset, 1)
                                    return inc_reg
                                    #cmp_reg = idc.get_operand_value(offset, 0)

                elif ins.itype == idaapi.NN_add:
                    if found_operand:
                        if idc.get_operand_type(offset, 0) == idc.o_reg and idc.get_operand_value(offset, 0) == cmp_reg:
                            # first op is a reg and matches the expected
                            if idc.get_operand_type(offset, 1) == idc.o_reg:
                                inc_reg = idc.print_operand(offset, 1)
                                self.logger.debug("Found inc reg :: %s" % inc_reg)
                                return inc_reg
                                # second op is a reg

            offset = idc.prev_head(offset, function_start)
        return None

    def backtrace_start(self, offset, max_instructions=1024):
        """
        From the current offset - backtrace and find the best instruction to start emulation.

        :param offset:
        :param max_instructions: The maximum number of instructions to backtrace
        :return:
        """
        function_start = idc.get_func_attr(offset, idc.FUNCATTR_START)
        blob_start = 0
        last_mov_or_alt = True

        trace_instruction_types = [idaapi.NN_mov,
                                   idaapi.NN_sub,
                                   idaapi.NN_xor,
                                   idaapi.NN_lea,
                                   idaapi.NN_add,
                                   idaapi.NN_inc,
                                   idaapi.NN_movupd,
                                   idaapi.NN_movups,
                                   idaapi.NN_movaps,
                                   idaapi.NN_movapd,]

        # Back trace
        if offset <= function_start + 64:
            """
            Note this can cause issues. If there is a call or other that is expected to have
            initialized data. Do a quick check first before returning the function start                         
            """

            if not self._has_call(function_start, offset):
                return function_start
        icount = 0
        while offset >= function_start:
            icount += 1

            if icount > max_instructions:
                return 0

            ins = ida_ua.insn_t()

            idaapi.decode_insn(ins, offset)
            self.logger.debug("0x%x %s" % (offset, idc.generate_disasm_line(idc.prev_head(offset), 0)))

            if ins.itype in trace_instruction_types:
                if ins.itype in [idaapi.NN_mov]:
                    last_mov_or_alt = True
                elif ins.itype == idaapi.NN_xor:
                    if idc.print_operand(offset, 0) != idc.print_operand(offset, 1):
                        if idc.get_operand_type(offset, 1) != idaapi.o_imm:
                            blob_start = idc.next_head(offset)
                            break
                        else:
                            last_mov_or_alt = True
                    else:
                        last_mov_or_alt = False
                elif ins.itype == idaapi.NN_sub:
                    if not last_mov_or_alt:
                        blob_start = idc.next_head(offset)
                        break
                    last_mov_or_alt = False
                elif ins.itype in [idaapi.NN_lea, idaapi.NN_inc, idaapi.NN_add]:
                    last_mov_or_alt = True

                blob_start = offset

                if offset <= function_start:
                    self.logger.debug("Error back-tracing ADVBLOB...Using function start")
                    blob_start = function_start
                    break

            else:
                blob_start = idc.next_head(offset)
                break

            offset = idc.prev_head(offset)

        if blob_start <= function_start + 64:
            if not self._has_call(function_start, blob_start):
                blob_start = function_start

        self.logger.debug("BLOB Start: %x" % blob_start)
        self.logger.debug(idc.print_insn_mnem(blob_start))
        return blob_start


class DecodeHandler(ida_kernwin.action_handler_t):
    """
        Handle hot key and mouse actions.

    """

    def __init__(self, logger, set_bookmarks=True, has_hexrays=False):
        ida_kernwin.action_handler_t.__init__(self)
        self.scanner = YaraScanner(logger)
        self.logger = logger
        self.has_hexrays = has_hexrays

        self.set_bookmarks = set_bookmarks
        self.mode = int(IdaHelpers.get_arch()/8)
        self.stacks = StackStack(logger)

    def trace_bytes(self):
        start = idc.read_selection_start()
        end = idc.read_selection_end()

        if start in BAD:
            self.logger.error("Nothing selected")
            idc.warning("No instructions selected!")
            return

        self.logger.debug("Selection Start: 0x%x" % start)
        self.logger.debug("Selection End:   0x%x" % end)

        if not start or not end:
            idc.warning("Error: Range not selected")
            return

    def _identify_impl_type(self, start, end):
        """
        Identify the implementation type.

        0 - func - string is decrypted in a call to a function. Offset is returned in eax
        1 - inline1 -

        TODO: This should be expanded to better detect different implementations

        :param start:
        :param end:
        :return:
        """

        self.logger.debug('End Param: %x' % end)
        ins = ida_ua.insn_t()
        idaapi.decode_insn(ins, idc.prev_head(end))
        idc.generate_disasm_line(idc.prev_head(end), 0)
        self.logger.debug("0x%x %s" % (end, idc.generate_disasm_line(idc.prev_head(end), 0)))
        if ins.itype == idaapi.NN_call:
            self.logger.debug("Ends in a call")
            return 0

    def _process(self, start, end, string_length=0):
        self.logger.debug("_process->enter")

        if start >= end:
            self.logger.error("End block before start!")
            self.logger.debug("Start: %x" % start)
            self.logger.debug("End:   %x" % end)
            return

        if end:
            ireg = self.stacks.scan_reg_increment(start, end)
            iregs = []
            if ireg:
                self.logger.debug("[*] Initializing [%s] to 1" % ireg)
                iregs.append((ireg, 1))

            self.logger.debug("[*] Using ImageBase: %x" % idaapi.get_imagebase())

            impl_type = self._identify_impl_type(start, end)

            semu = SUE(code_base=idaapi.get_imagebase(), mode=self.mode, logger=self.logger)

            sresult = semu.deobfuscate_stack(start, end, string_length=string_length, impl_type=impl_type, init_regs=iregs)

            self.logger.debug("...Complete....")

            if not sresult:
                self.logger.debug("No result data!")
                return

            self.logger.debug("-" * 16)
            for rk in sresult.keys():
                try:
                    self.logger.debug("%s: 0x%x" % (rk, sresult[rk]))
                except TypeError:
                    self.logger.debug("%s: %s" % (rk, sresult[rk]))
            self.logger.debug("-" * 16)

            decoded = sresult['data']

            if decoded:
                if self.set_bookmarks:
                    IdaHelpers.add_bookmark(start, decoded)
                IdaHelpers.add_comment(start, decoded, hexrays=self.has_hexrays)
                return decoded

    def process_matches(self, matches, function_start):
        """

        :param matches:
        :return:
        """

        if not matches:
            return

        # stacks = StackStack(self.logger)

        last_start = 0
        last_end = 0
        decoded_strings = []
        deocded_offsets = []
        for match in matches:
            try:
                match_offset = function_start + match
                self.logger.debug("Processing match offset: %x" % match_offset)
                block_start = self.stacks.backtrace_start(match_offset)
                if not block_start:
                    self.logger.error("Could not find block start for %x .. skipping" % match_offset)
                    continue

                block_end = self.stacks.find_end(block_start)
                if not block_end:
                    self.logger.error("Could not find block end for %x .. skipping" % match_offset)
                    continue

                if last_end == 0 == last_start:
                    last_end = block_end
                    last_start = block_start
                else:
                    if last_start < match_offset < last_end:
                        self.logger.debug("Skipping offset [%x]" % match_offset)
                        continue


                string_length = self.stacks.get_string_length(block_start)
                self.logger.debug("Processing: start: %x, end: %x, string_length: %d" % (block_start, block_end, string_length))
                decoded = self._process(block_start, block_end, string_length=string_length)
                if decoded:
                    if not block_start in deocded_offsets:
                        decoded_strings.append((block_start, decoded))
                        deocded_offsets.append(block_start)

                self.logger.debug("Using start of: %x" % block_start)
            except Exception as ex:
                self.logger.error("Error processing block: %s" % ex)

        return decoded_strings

    def update(self, ctx):
        """
            Required by action handler
        """
        return ida_kernwin.AST_ENABLE_ALWAYS

    def activate(self, ctx):
        if ctx.action == 'ssp_decode_selected':
            self.decode_selected()
        elif ctx.action == 'ssp_decode_current':
            self.decode_current()
        elif ctx.action == 'ssp_decode_func':
            self.decode_function()
        elif ctx.action == 'ssp_trace_selected':
            self.trace_bytes()
        elif ctx.action == 'ssp_decode_all':
            self.decode_all()
        else:
            self.logger.debug(ctx.cur_func)
            self.logger.debug("Not Supported")
            idc.warning("Not Implemented")
        return True

    def decode_function(self):

        # advs = StackStack(self.logger)
        offset = idaapi.get_item_head(idc.here())

        function_name = idc.get_func_name(offset)

        self.logger.debug("Processing: %s" % function_name)

        function_start = idc.get_func_attr(offset, idc.FUNCATTR_START)
        function_end = idc.get_func_attr(offset, idc.FUNCATTR_END)

        fdata = idc.get_bytes(function_start, function_end - function_start)

        refs = self.scanner.scan_function(fdata)

        if refs:
            self.logger.debug("Found [%d] Obfuscated Strings" % len(refs))

            decoded_strings = self.process_matches(refs, function_start)
            if decoded_strings:
                # TODO: Add this to a UI Pop-up
                print("--- Decoded Function Strings ---")
                for o, s in decoded_strings:
                    print(" %x: %s" % (o, s))

        else:
            self.logger.debug("No code blocks found..")

    def decode_all(self):
        try:
            # ss = StackStack(self.logger)
            decoded_strings = []
            for func_entry in idautils.Functions():
                if idc.get_func_attr(func_entry, idc.FUNCATTR_FLAGS) & idc.FUNC_LIB:
                    self.logger.debug("Skipping lib function %s" % idc.get_func_name(func_entry))
                    continue

                fdata = idc.get_bytes(func_entry, idc.get_func_attr(func_entry, idc.FUNCATTR_END) - func_entry)
                refs = self.scanner.scan_function(fdata)
                if refs:
                    self.logger.debug("Found [%d] Obfuscated Strings" % len(refs))
                    decoded_strings.extend(self.process_matches(refs, func_entry))

            if decoded_strings:
                # TODO: Add this to a UI Pop-up
                print("--- Decoded Function Strings ---")
                for o, s in decoded_strings:
                    print(" %x: %s" % (o, s))
        except Exception as ex:
            self.logger.error("Error processing file: %s" % ex)

    def decode_current(self):
        # stacks = StackStack(self.logger)

        offset = idaapi.get_item_head(idc.here())
        self.logger.debug("Starting scan at offset: %x" % offset)

        # TODO: Make this a configuration option to either use back trace or blocks
        start = self.stacks.backtrace_start(offset)

        self.logger.debug("Snippet Start: %x" % start)

        if start:
            end = self.stacks.find_end(start)
            if end:
                string_length = self.stacks.get_string_length(start)
                self.logger.debug("Processing: start: %x, end: %x, string_length: %d" % (start, end, string_length))
                self._process(start, end, string_length=string_length)
            else:
                self.logger.error("Did not find end. Manually select instructions!")

        else:
            idc.warning("Decode Failed! Unable to get ADVBlock")

    def decode_selected(self):

        start = idc.read_selection_start()
        end = idc.read_selection_end()

        if start in BAD:
            self.logger.debug("Nothing selected")
            idc.warning("Nothing Selected!")
            return

        self.logger.debug("Selection Start: 0x%x" % start)
        self.logger.debug("Selection End:   0x%x" % end)

        if start:
            self._process(start, end)
        else:
            idc.warning("Selected Block ")


class ScanHandler(ida_kernwin.action_handler_t):

    def __init__(self, logger):
        self.logger = logger
        ida_kernwin.action_handler_t.__init__(self)
        self.scanner = YaraScanner(logger)

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

    def _scan_bin(self):
        self.logger.info("Scanning file...")
        scan_result = self.scanner.scan_functions()
        self.logger.info("Scan complete...")
        if not scan_result:
            self.logger.error("Found no suspect code blocks!")
            return

        self.logger.info("-" * 16)
        self.logger.info("Found %d suspect functions" % len(scan_result.keys()))
        for ref in scan_result:
            self.logger.info('%s\t%d' % (ref, scan_result[ref]))
        self.logger.info("-" * 16)

    def activate(self, ctx):
        if ctx.action == 'ssp_scan':
            self._scan_bin()
        else:
            self.logger.error("Unsupported scan option :: %s" % ctx.action)
            idc.warning("Unsupported Scan Option")
        return True

    def term(self):
        pass


class StackStackConfig(object):
    pass


class StackStackPlugin(ida_idaapi.plugin_t):

    flags = ida_idaapi.PLUGIN_KEEP

    comment = "StackStack simple emulation and tracing"
    help = "StackStack - simple emulation and tracing"
    wanted_name = "StackStack"
    wanted_hotkey = ""

    _version = 1.07

    def init(self):
        try:
            self.config = self.load_configuration()
            try:
                self.logger = self._init_logger(logging._checkLevel(self.config['loglevel'].upper()))
            except ValueError:
                self.logger = self._init_logger(self.logger.setLevel(logging.DEBUG))

            # TODO: Remove
            self.logger.setLevel(logging.DEBUG)

            self.logger.info("StackStack version: %s" % StackStackPlugin._version)

            if self.config['check_update']:
                version_check = Update.check_version(StackStackPlugin._version)
                if version_check > 0:
                    idc.warning("StackStack version %s is now available for download." % version_check)

            self.has_hexrays = IdaHelpers.has_decompiler()

            self.actions = []
            self.define_actions()
            self.menus = Menus()
            self.menus.hook()
        except Exception as ex:
            self.logger.error('Error initializing StackStack %s' % ex)
            idc.warning('Error initializing StackStack %s' % ex)

        return ida_idaapi.PLUGIN_KEEP

    def _init_logger(self, loglevel, name='stackstack'):
        """
        Initialize Logger

        :param loglevel: Log Level to use
        :param name: Log name
        :return:
        """
        # logging.basicConfig()
        # logging.root.setLevel(logging.NOTSET)

        logger = logging.getLogger(name)
        # logger.setLevel(loglevel)
        logger.setLevel(loglevel)

        log_stream = logging.StreamHandler()
        formatter = logging.Formatter('stackstack:%(levelname)s:%(message)s')
        log_stream.setFormatter(formatter)
        logger.addHandler(log_stream)
        return logger

    def load_configuration(self, config_name='stackstack.cfg', generate_default_config=True):
        path = ida_diskio.get_user_idadir()
        config_path = os.path.join(path, config_name)
        config_data = None

        if os.path.exists(config_path):
            with open(config_path, 'r') as inf:
                config_data = json.loads(inf.read())

        if config_data:
            def_config = self._generate_default_configuration()
            missing_options = False
            for key in def_config.keys():
                # iterate through the default config keys and add any missing config entries.
                try:
                    config_data[key]
                except KeyError:
                    config_data[key] = def_config[key]
                    missing_options = True
            if missing_options:
                with open(config_path, 'w') as out:
                    out.write(json.dumps(config_data))
            return config_data

        if generate_default_config:
            config_data = self._generate_default_configuration()
            with open(config_path, 'w') as out:
                out.write(json.dumps(config_data))
        return config_data

    def _generate_default_configuration(self):
        return {
            'loglevel': 'DEBUG',
            'ext_yara_file': 'stackstack.yara',
            'bookmarks': True,
            'rename_func': False,
            'check_update': False
        }

    def _get_scan_actions(self):
        scanner = ScanHandler(self.logger)
        return [
            ida_kernwin.action_desc_t(
                "ssp_scan",
                "Scan",
                scanner,
                "Shift-s",
                "Scan binary for functions with encrypted strings"
            )]

    def _get_decode_actions(self):
        decode = DecodeHandler(self.logger, has_hexrays=self.has_hexrays)
        return [
            ida_kernwin.action_desc_t(
                "ssp_trace_selected",
                "Trace Selected",
                decode,
                "Trace the selected bytes."
            ),
            ida_kernwin.action_desc_t(
                "ssp_decode_selected",
                "Decode Selected",
                decode,
                "Decode the selected bytes."
            ),
            ida_kernwin.action_desc_t(
                "ssp_decode_current",
                "Decode Current",
                decode,
                "Shift-x",
                "Detect and decode the current obfuscated bytes."
            ),
            ida_kernwin.action_desc_t(
                "ssp_decode_all",
                "Decode All",
                decode,
                "Scan and decode all instances."
            ),
            ida_kernwin.action_desc_t(
                "ssp_decode_func",
                "Decode Function",
                decode,
                "Decode all instances in the current function."
            )
        ]

    def _get_util_actions(self):
        return [
            ida_kernwin.action_desc_t(
                "ssp_decode_func",
                "Decode Function",
                DecodeHandler(has_hexrays=self.has_hexrays),
                "Decode all instances in the current function."
            )
        ]

    def define_actions(self):
        actions = []
        actions.extend(self._get_scan_actions())
        actions.extend(self._get_decode_actions())
        self.actions = actions
        for action_desc in actions:
            ida_kernwin.register_action(action_desc)

    def run(self, arg):
        print('-------------------------------------')
        print('running\n\n\n\n\n\n\n\n\n')
        print('-------------------------------------')
        pass

    def term(self):
        if self.actions:
            for action_desc in self.actions:
                ida_kernwin.unregister_action(action_desc.name)


class Menus(ida_kernwin.UI_Hooks):

    def finish_populating_widget_popup(self, form, popup):

        if ida_kernwin.get_widget_type(form) in [ida_kernwin.BWN_PSEUDOCODE]:
            ida_kernwin.attach_action_to_popup(form, popup, "ssp_decode_current", "StackStack/Decode/")

        if ida_kernwin.get_widget_type(form) == ida_kernwin.BWN_DISASM:
            ida_kernwin.attach_action_to_popup(form, popup, "ssp_decode_selected", "StackStack/Decode/")
            ida_kernwin.attach_action_to_popup(form, popup, "ssp_trace_selected", "StackStack/Trace/")

        if ida_kernwin.get_widget_type(form) in [ida_kernwin.BWN_PSEUDOCODE, ida_kernwin.BWN_DISASM]:
            ida_kernwin.attach_action_to_popup(form, popup, "ssp_scan", "StackStack/")
            ida_kernwin.attach_action_to_popup(form, popup, "ssp_decode_all", "StackStack/Decode/")
            ida_kernwin.attach_action_to_popup(form, popup, "ssp_decode_func", "StackStack/Decode/")

def PLUGIN_ENTRY():
    return StackStackPlugin()


repository: idapython/src
file: examples/misc/merge/py_mex1.py
description: IDAPython project for Hex-Ray's IDA Pro
wanted_name: IDAPython: Merge example 1
comment: IDAPython: An example 1 how to implement IDA merge 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/merge/py_mex1.py

content:
# pylint: disable=line-too-long,invalid-name,import-error

"""
summary: add merge functionality to a simple plugin

description:
  This is a primitive plugin which asks user for some info and saves it for
  some addresses.

  We will add a merge functionality to plugin.

  An IDA plugin may have two kinds of data with permanent storage:
    1. Data common for entire database (e.g. the options).
       To describe them we will use the idbattr_info_t type.
    2. Data specific to a particular address.
       To describe them we will use the merge_node_info_t type.

  Also, see SDK/plugins/mex1 example

level: advanced
"""


import ida_idaapi
import ida_ida
import ida_kernwin
import ida_netnode
import ida_funcs
import ida_merge
import ida_mergemod
import ida_idp

# --------------------------------------------------------------------------
# netnode to store plugin data
MEX_NODE_NAME = "$ idapython mex1"
# user input
MEX_OPTION_FLAGS_IDX = ida_idaapi.BADADDR & -1   # atag
MEX_OPTION_IDENT_IDX = ida_idaapi.BADADDR & -2   # stag
# EA marks
MEX_EA_TAG = 'm'

# mex_ctx_t::flags bits
MEX_FLAGS_0 = 0x01
MEX_FLAGS_1 = 0x02

# --------------------------------------------------------------------------
class idp_listener_t(ida_idp.IDP_Hooks):
    """
    we need an event listener to catch processor_t::ev_create_merge_handlers
    """

    def __init__(self, ctx):
        # by default IDP_Hooks uses ida_idp.HKCB_GLOBAL hkcb_flags,
        # in that case IDP events are sent to all DB instances.
        # Such behaviour does not matter for IDA Pro
        # but IDA Teams needs to have only one IDP event sent to plugin.
        # We set hkcb_flags to 0
        ida_idp.IDP_Hooks.__init__(self, 0, 0)
        self.ctx = ctx

    def ev_ending_undo(self):
        """
        A well behaving plugin should restore its state from the database
        upon ev_ending_undo. Otherwise its state may be conflicting with the
        database.
        """
        self.ctx.restore_from_idb()
        return 0

    def ev_create_merge_handlers(self, md):
        """
        This event occurs when IDA is performing a 3-way merge (for IDA Teams)
        Our plugins should create and register merge handler(s) for its data.
        """
        self.ctx.create_merge_handlers(md)
        return 0


    def ev_cvt64_supval(self, node, tag, idx, data):
        """
        Converter to i64 database
        """
        nn = ida_netnode.netnode(MEX_NODE_NAME)
        if nn == node:
            if tag == ida_netnode.stag:
                nn.supset(MEX_OPTION_IDENT_IDX, data)
                return 1
            if tag == ida_netnode.atag and len(data):
                val = int.from_bytes(data, 'little')
                nn.altset(MEX_OPTION_FLAGS_IDX, val)
                return 1
            if chr(tag) == MEX_EA_TAG:
                nn.supset(idx, data, MEX_EA_TAG)
                return 1
        return 0


# --------------------------------------------------------------------------
class mex_ctx_t(ida_idaapi.plugmod_t):
    """
    Regular plugin implementation below.
    For example, in our case the plugin asks for 2 bit values and a string value.
    Then the plugin stores this data in the database.
    And mark the start address of the current function.
    These data will be merged later.
    """

    def __init__(self):
        # bit flags, see above MEX_FLAGS_0/MEX_FLAGS_1
        self.flags = 0
        # unique database ident
        self.ident = ""
        # Restore the plugin data from the database into the memory.
        self.restore_from_idb()
        # Hook an event listener, to catch the merge-related event(s).
        self.idp_listener = idp_listener_t(self)
        self.idp_listener.hook()
        # MERGE: the following data must exist during plugin lifetime
        self.modmerger_helper = None
        self.idpopts_info = None
        self.merge_node_info = None

    def save_to_idb(self):
        """ Save the plugin state to the idb. """
        nn = ida_netnode.netnode(MEX_NODE_NAME, 0, True)
        nn.altset(MEX_OPTION_FLAGS_IDX, self.flags)
        nn.supset(MEX_OPTION_IDENT_IDX, self.ident)

    def restore_from_idb(self):
        """ Restore plugin variables from the idb. """
        nn = ida_netnode.netnode(MEX_NODE_NAME)
        if nn != ida_netnode.BADNODE:
            self.flags = nn.altval(MEX_OPTION_FLAGS_IDX)
            self.ident = nn.supstr(MEX_OPTION_IDENT_IDX)

    def run(self, _):
        """
        Ask user for the data and save them to database.
        Add mark for current EA.
        """
        if self._ask_form():
            self.save_to_idb()
        # Our plugin stores a string for the current function.
        # Just for illustration purposes of how plugins should merge address-specific info
        # stored in a netnode.
        ea = ida_kernwin.get_screen_ea()
        pfn = ida_funcs.get_func(ea)
        if pfn:
            one = ""
            if (self.flags & MEX_FLAGS_0) != 0:
                one = " one"
            two = ""
            if (self.flags & MEX_FLAGS_1) != 0:
                one = " two"
            mark = "IPMEX1" + one + two
            nn = ida_netnode.netnode(MEX_NODE_NAME, 0, True)
            nn.supset_ea(pfn.start_ea, mark, MEX_EA_TAG)

    def _ask_form(self):
        class MexForm(ida_kernwin.Form):
            def __init__(self):
                ida_kernwin.Form.__init__(self,
                r"""
IDAPython: merge example 1

<Flag 0:{optFlag0}>
<Flag 1:{optFlag1}>{grpFlags}>
<Ident prefix:{ident}>
                """,
                {
                    "grpFlags": ida_kernwin.Form.ChkGroupControl(("optFlag0", "optFlag1")),
                    "ident"   : ida_kernwin.Form.StringInput(swidth=10),
                })

        form = MexForm()
        form, _ = form.Compile()
        form.grpFlags.value = self.flags
        form.ident.value = self.ident
        ok = form.Execute()
        if ok == 1:
            self.flags = form.grpFlags.value
            self.ident = form.ident.value
        form.Free()
        return ok == 1

    def create_merge_handlers(self, md):
        """
        Create merge handlers for plugin
        """

        # 1. Data common for entire database (e.g. the options).
        #
        # This example shows how to merge the data from database.
        # We will describe the items to merge and pass the description
        # to create_std_modmerge_handlers(), which will do all the work for us.
        sizeof_flags = ida_netnode.SIZEOF_nodeidx_t
        self.idpopts_info = [
            # Describe both flags
            ida_ida.idbattr_info_t("MEX flag 0", MEX_OPTION_FLAGS_IDX, sizeof_flags, MEX_FLAGS_0, ida_netnode.atag, ida_ida.IDI_ALTVAL|ida_ida.IDI_SCALAR),
            ida_ida.idbattr_info_t("MEX flag 1", MEX_OPTION_FLAGS_IDX, sizeof_flags, MEX_FLAGS_1, ida_netnode.atag, ida_ida.IDI_ALTVAL|ida_ida.IDI_SCALAR),
            # Describe ident
            ida_ida.idbattr_info_t("MEX ident",  MEX_OPTION_IDENT_IDX, 0,            0,           ida_netnode.stag, ida_ida.IDI_SUPVAL|ida_ida.IDI_CSTR),
        ]

        # The descriptions are ready. Now create an instance of the standard helper
        # class to be passed to the kernel, and the kernel will take care of organizing
        # the merge process for them.

        # helper instance name
        self.modmerger_helper = ida_merge.moddata_diff_helper_t(
            "Sample merge data", # label: prefix for the attribute names, e.g. "Sample merge data.MEX flag 0"
            MEX_NODE_NAME,       # netnode name for idpopts_info and merge_node_info
            self.idpopts_info)   # field descriptions

        # 2. Data specific to a particular address.
        # We describe how the data is stored in a netnode.
        self.merge_node_info = [
            ida_merge.merge_node_info_t(
                "Function marks",       # label of the merge handler, e.g. "Plugins/Merge example 1/Function marks"
                MEX_EA_TAG,             # netnode tag
                ida_merge.NDS_MAP_IDX|ida_merge.NDS_IS_STR
                                        # netnode value descriptors and modificators, see \ref nds_flags_t
            ),
        ]

        # Now we should combine together in function create_std_modmerge_handlers.
        # This function will be called on processor_t::ev_create_merge_handlers event.
        # As a result, two merge handlers with labels
        #   "Plugins/Merge example 1/Database attributes"
        #   "Plugins/Merge example 1/Function marks"
        # will be created.
        mhp = ida_merge.merge_handler_params_t(
            md,                                  # merge handler data
            "Plugins/IDAPython merge example 1", # Label of the merge handler
            ida_merge.MERGE_KIND_NONE,           # allocate a merge kind
            ida_merge.MERGE_KIND_END,            # insert to the end of handler list
            0)                                   # merge handler flags
        ida_mergemod.create_std_modmerge_handlers(
             mhp,
             self.modmerger_helper,
             self.merge_node_info)


# --------------------------------------------------------------------------
class mex1_plugin_t(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_MULTI | ida_idaapi.PLUGIN_MOD
    wanted_name = "IDAPython: Merge example 1"
    comment = "IDAPython: An example 1 how to implement IDA merge functionality"
    wanted_hotkey = ""
    help = ""
    def init(self):
        return mex_ctx_t()
    def term(self):
        pass
    def run(self, arg):
        pass
def PLUGIN_ENTRY():
    return mex1_plugin_t()


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

content:
from __future__ import annotations

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

import idaapi
import idc

listen = "0.0.0.0"
listen_port = 31337

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

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

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

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


        DEBUG_MARSHALLING = False


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

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

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

                appender(marshalled)

            return wrapper


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

        host = listen
        port = listen_port

        mutex = threading.Condition()


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

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

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

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

                return rv[0]

            return wrapper


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


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


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

            return y


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


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

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


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

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

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

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


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

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

    def run(self, arg):
        # Code à exécuter lors du lancement via l'interface
        return

    def term(self):
        # Code à exécuter lors du déchargement du plugin
        # par exemple, lors de la fermeture de l'idb.
        pass

def PLUGIN_ENTRY():
    return autoLoadXMLRPC()


repository: ElvisBlue/emotet-deobfuscator
file: emotet_deobfuscator.py
description: IDA plugin to deobfuscate emotet CFF
wanted_name: Emotet Deobfuscator (POC)
comment: Emotet Deobfuscator (POC)
created_at: 2021-11-25 03:22:57+00:00
pushed_at: 2022-04-26 05:05:19+00:00
url:
https://github.com/ElvisBlue/emotet-deobfuscator/blob/1bba44eecf7002cea7c22ea49df1b5618e752a05/emotet_deobfuscator.py

content:
# pylint: disable=C0103,C0111,W0614

from __future__ import print_function

import ida_idaapi
from ida_hexrays import *
from collections import defaultdict

class CEmotetCFF(optblock_t):
    def __init__(self):
        optblock_t.__init__(self)
        self.isOptimize = False

    def FindDispatchReg(self, mba):
        infoCollectorObj = CDispatchFinder()
        mba.for_all_topinsns(infoCollectorObj)
        return infoCollectorObj.GetDispatchList()

    def FindBlockStatus(self, mba, dispatchReg):
        blockCnt = mba.qty
        statusResult = defaultdict()

        for i in range(1, blockCnt - 1):
            currentBlock = mba.get_mblock(i)
            if currentBlock.tail != None and currentBlock.tail.opcode in [m_jz, m_jnz]:
                mins = currentBlock.tail
                if mins.l.is_reg() and mins.r.is_constant():
                    if get_mreg_name(mins.l.r, mins.l.size) == dispatchReg:
                        status = mins.r.value(False)
                        if mins.opcode == m_jz:
                            #print("status 0x%X match with block %d" % (status, mins.d.b))
                            statusResult[status] = mins.d.b
                        elif mins.opcode == m_jnz:
                            #print("status 0x%X match with block %d" % (status, currentBlock.nextb.serial))
                            statusResult[status] = currentBlock.nextb.serial

        return statusResult

    def CorrectBlock(self, mba, dispatchReg, statusList):
        blockCnt = mba.qty
        changed = 0
        
        fixStatusList = []
        wontFixStatusList = []
        
        for i in range(1, blockCnt - 1):
            currentBlock = mba.get_mblock(i)
            
            #We find micro instruction mov dispatchReg, const
            mins = currentBlock.head
            while mins != currentBlock.tail:
            
                if mins == None:
                    break
                    
                if mins.opcode == m_mov and mins.l.t == mop_n and mins.d.is_reg():
                    if get_mreg_name(mins.d.r, mins.d.size) == dispatchReg:
                        #print(mins.dstr())
                        status = mins.l.value(False)
                        if status in statusList:
                        
                            if currentBlock.type != BLT_1WAY:
                                wontFixStatusList.append(status)
                                break

                            #Get block match with status
                            dstBlock = mba.get_mblock(statusList[status])
                            
                            #We delete mov dispatchReg, const instruction
                            currentBlock.make_nop(mins)
                            
                            #Check
                            if dstBlock.serial == currentBlock.serial:
                                break
                            
                            
                            #Insert/replace goto block
                            minsGoto = minsn_t(currentBlock.tail.ea)
                            if currentBlock.tail.opcode == m_goto:
                                currentBlock.tail.l._make_blkref(dstBlock.serial)
                            else:
                                minsGoto.opcode = m_goto
                                minsGoto.l._make_blkref(dstBlock.serial)
                                minsGoto.r = mop_t()
                                minsGoto.d = mop_t()
                                minsGoto.iprops = 0
                                currentBlock.insert_into_block(minsGoto, currentBlock.tail)

                            
                            
                            #Correct predset and succset
                            for oldDstBlockSerial in currentBlock.succset:
                                oldDstBlock = mba.get_mblock(oldDstBlockSerial)
                                oldDstBlock.predset._del(currentBlock.serial)
                            
                            for oldPresetSerial in dstBlock.predset:
                                oldPresetBlock = mba.get_mblock(oldPresetSerial)
                                

                            dstBlock.predset.push_back(currentBlock.serial)
                            
                            currentBlock.succset.clear()
                            currentBlock.succset.push_back(dstBlock.serial)
                            
                            mba.verify(True)
                            
                            if status not in wontFixStatusList:
                                if status not in fixStatusList:
                                    fixStatusList.append(status)
                            
                            #process next block
                            break
                
                #Update to handle control flow status set like this
                #v0 = result != 0 ? 61721445 : 200024545;
                    
                if mins.opcode == m_add and mins.l.t == mop_d and mins.r.t == mop_n and mins.d.t == mop_r and get_mreg_name(mins.d.r, mins.d.size) == dispatchReg:
                    subMins1 = mins.l.d
                    if subMins1.opcode == m_and and subMins1.l.t == mop_d and subMins1.r.t == mop_n:
                        subMins2 = subMins1.l.d
                        if subMins2.opcode == m_neg and subMins2.l.t == mop_d:
                            subMins3 = subMins2.l.d
                            if subMins3.opcode == m_xdu and subMins3.l.t == mop_d:
                                subMins4 = subMins3.l.d
                                if subMins4.opcode == m_setnz and subMins4.l.t == mop_r and subMins4.r.t == mop_n and subMins4.r.value(False) == 0:
                                    status1 = mins.r.value(False)
                                    status2 = (mins.r.value(False) + subMins1.r.value(False)) & 0xFFFFFFFF
                                    if status1 in statusList and status2 in statusList:
                                        #We need to do 2 thing: Fix current block from 1WAY to 2WAY and insert new block. This is an if/else stament
                                        
                                        #TODO: Fix the bug
                                        if currentBlock.tail.opcode == m_goto:
                                            currentBlock.make_nop(currentBlock.tail)
                                        
                                        for oldDstBlockSerial in currentBlock.succset:
                                            oldDstBlock = mba.get_mblock(oldDstBlockSerial)
                                            oldDstBlock.predset._del(currentBlock.serial)
                                        currentBlock.succset.clear()
                                        
                                        newDstBlockSerial = statusList[status2]
                                        if currentBlock.serial not in mba.get_mblock(newDstBlockSerial).predset:
                                            mba.get_mblock(newDstBlockSerial).predset.push_back(currentBlock.serial)
                                        
                                            
                                        insertedJnzIns = minsn_t(currentBlock.tail.ea)
                                        insertedJnzIns.opcode = m_jnz
                                        insertedJnzIns.l = mop_t()
                                        insertedJnzIns.l.make_reg(subMins4.l.r, 4)
                                        insertedJnzIns.r = mop_t()
                                        insertedJnzIns.r.make_number(subMins4.r.value(False), 4)
                                        insertedJnzIns.d = mop_t()
                                        insertedJnzIns.d.make_blkref(newDstBlockSerial)
                                        currentBlock.insert_into_block(insertedJnzIns, currentBlock.tail)
                                        
                                        #I don't know why currentBlock.type = BLT_2WAY throw exception so set it to BLT_NONE (not computed yet)
                                        currentBlock.type = BLT_NONE
                                        currentBlock.succset.push_back(newDstBlockSerial)
                                            
                                        #Because new block will be inserted before current block so we use currentBlock.serial + 1
                                        newInsertedBlock = mba.insert_block(currentBlock.serial + 1)

                                            
                                        #After insert new block, We should update block sertial in statusList
                                        #Old statuses in predset and succset were updated by insert_block even in current block so we don't have to care about it
                                        for s in statusList:
                                            if statusList[s] > currentBlock.serial:
                                                statusList[s] += 1
                                        
                                        currentBlock.succset.push_back(newInsertedBlock.serial)
                                        newInsertedBlock.predset.push_back(currentBlock.serial)
                                        newGotoBlockSerial = statusList[status1]
                                        mba.get_mblock(newGotoBlockSerial).predset.push_back(newInsertedBlock.serial)
                                        insertedGotoIns = minsn_t(currentBlock.tail.ea)
                                        insertedGotoIns.opcode = m_goto
                                        insertedGotoIns.l._make_blkref(newGotoBlockSerial)
                                        insertedGotoIns.r = mop_t()
                                        insertedGotoIns.d = mop_t()
                                        newInsertedBlock.insert_into_block(insertedGotoIns, None)

                                            
                                        newInsertedBlock.succset.push_back(newGotoBlockSerial)
                                        newInsertedBlock.type = BLT_NONE
                                        newInsertedBlock.start = currentBlock.start
                                        newInsertedBlock.end = currentBlock.end
                                        newInsertedBlock.flags = currentBlock.flags
                                        newInsertedBlock.mark_lists_dirty()
                                        
                                        currentBlock.make_nop(mins)
                                        currentBlock.mark_lists_dirty()
                                        
                                        mba.verify(True)
                                        
                                        fixStatusList.append(status1)
                                        fixStatusList.append(status2)
                                            
                                        break
                
                mins = mins.next
        
        #Process final fixStatusList to return
        fixStatusList = list(dict.fromkeys(fixStatusList))
        wontFixStatusList = list(dict.fromkeys(wontFixStatusList))
        finalFixStatusList = [x for x in fixStatusList if x not in wontFixStatusList]
        return finalFixStatusList

    def CleanDispatch(self, mba, dispatchReg, fixStatusList):
        dispatchCleanerObj = CDispatchCleaner()
        dispatchCleanerObj.SetDispatchReg(dispatchReg)
        dispatchCleanerObj.SetFixStatusList(fixStatusList)
        mba.for_all_topinsns(dispatchCleanerObj)
        return dispatchCleanerObj.changed

    def func(self, blk):
        mba = blk.mba

        if mba.maturity != MMAT_GLBOPT1:
            self.isOptimize = False
            return 0

        if self.isOptimize == True:
            return 0

        self.isOptimize = True

        changed = 0
        dispatchDict = self.FindDispatchReg(mba)
        for k, v in dispatchDict.items():
            if v >= 3: #If reg was compare with const at least 3 times, consider it's dispatch reg
                blockStatus = self.FindBlockStatus(mba, k)
                fixStatusList = self.CorrectBlock(mba, k, blockStatus)
                changed += len(fixStatusList)
                changed += self.CleanDispatch(mba, k, fixStatusList)

        #Perform optimize
        mba.mark_chains_dirty()
        mba.optimize_local(0)
        mba.verify(True)

        return changed

class CDispatchFinder(minsn_visitor_t):
    def __init__(self):
        minsn_visitor_t.__init__(self)
        self.dispatchResult = defaultdict()

    def visit_minsn(self):
        #Collect info about mov var, const and jnz/jz var, dst
        ins = self.curins
        if ins.opcode in [m_jz, m_jnz]:
            if ins.r.t == mop_n and ins.l.t == mop_r:
                regName = get_mreg_name(ins.l.r, ins.l.size)
                if regName in self.dispatchResult:
                    self.dispatchResult[regName] += 1
                else:
                    self.dispatchResult[regName] = 1

        return 0

    def GetDispatchList(self):
        return self.dispatchResult

class CDispatchCleaner(minsn_visitor_t):
    def __init__(self):
        minsn_visitor_t.__init__(self)
        self.changed = 0
        self.fixStatusList = None
        self.dispatchReg = ""

    def SetFixStatusList(self, fixStatusList):
        self.fixStatusList = fixStatusList

    def SetDispatchReg(self, dispatchReg):
        self.dispatchReg = dispatchReg

    def visit_minsn(self):
        #Detect jnz/jz var, dst then NOP it
        ins = self.curins
        currentBlock = self.blk
        if ins.opcode in [m_jz, m_jnz]:
            if ins.r.t == mop_n and ins.l.t == mop_r:
                if self.dispatchReg == get_mreg_name(ins.l.r, ins.l.size) and ins.r.value(False) in self.fixStatusList:
                    #Now we clean block
                    jmpDst = ins.d.b
                    if ins.opcode == m_jz:
                        #print("clean " + ins.dstr())
                        self.blk.make_nop(ins)
                        dstBlock = self.mba.get_mblock(jmpDst)
                        dstBlock.predset._del(self.blk.serial)
                        self.blk.succset._del(jmpDst)
                        self.blk.type = BLT_1WAY
                    else: #m_jnz
                        #print("clean " + ins.dstr())
                        self.curins.opcode = m_goto
                        self.curins.l = mop_t()
                        self.curins.l._make_blkref(jmpDst)
                        self.curins.r = mop_t()
                        self.curins.d = mop_t()
                        self.blk.mark_lists_dirty()
                        for serial in self.blk.succset:
                            oldDstBlock = self.mba.get_mblock(serial)
                            if oldDstBlock.serial != jmpDst:
                                oldDstBlock.predset._del(self.blk.serial)

                        self.blk.succset.clear()
                        self.blk.succset.push_back(jmpDst)
                        self.blk.type = BLT_1WAY

                    self.mba.optimize_local(0)
                    self.mba.verify(True)
                    self.changed += 1
        return 0

class CEmotetDeobfuscator(ida_idaapi.plugin_t):
    flags = 0
    comment = "Emotet Deobfuscator (POC)"
    help = ""
    wanted_name = "Emotet Deobfuscator (POC)"
    wanted_hotkey = ""

    def __init__(self):
        self.actived = False
        self.emotetUnCFFObj = CEmotetCFF()
        return

    def init(self):
        print("Emotet Deobfuscator loaded")
        return ida_idaapi.PLUGIN_OK

    def run(self, arg):
        if self.actived == True:
            self.actived = False
            self.emotetUnCFFObj.remove()
            print("Emotet Deobfuscator de-actived")
        else:
            self.actived = True
            self.emotetUnCFFObj.install()
            print("Emotet Deobfuscator actived")

        return True

    def term(self):
        return

def PLUGIN_ENTRY():
    return CEmotetDeobfuscator()


repository: partoftheworlD/IDA7py_FunctionStringAssociate
file: functionstringassociate.py
wanted_name: Strings Function Associate
comment: Associate functions
created_at: 2018-06-12 14:05:39+00:00
pushed_at: 2024-01-15 19:14:46+00:00
description:
It does the same thing as plugin FunctionStringAssociate, but don't need to compile the plugin or download from somewhere for x86 and x64. Now the script just works, in the future it is planned to increase the speed of work and bring it closer to the speed of the plugin.

url:
https://github.com/partoftheworlD/IDA7py_FunctionStringAssociate/blob/0ead9e3cb06dc4a8e40f83ec454e993d7e4d9c39/functionstringassociate.py

content:
import idc
import idautils
import idaapi
import ida_bytes

print("\nStringsFunctionAssociate v0.09 by partoftheworlD! Last Changes <2024-01-15 19:50:02.814613>\n")

class StringException(Exception):
    pass

class Strings_fc:
    def __init__(self):
        self.string_counter = 0

    string_types = ["C",
                    "Pascal",
                    "LEN2",
                    "Unicode",
                    "LEN4",
                    "ULEN2",
                    "ULEN4"]

    def get_string_type(self, addr):
        try:
            type_s = idc.get_str_type(addr)
            return str(ida_bytes.get_strlit_contents(addr, -1, type_s))
        except TypeError:
            raise StringException()        
    
    def clear_comments(self, func_obj):
        idaapi.set_func_cmt(func_obj, '', 1)
        idaapi.set_func_cmt(func_obj, '', 0)

    def save_existing_comments(self, start_func):
        bc = []
        comment = idaapi.get_func_cmt(start_func, 1)
        if comment:
            bc.append(comment)
        comment = idaapi.get_func_cmt(start_func, 0)
        if comment:
            bc.append(comment)
        if len(bc) > 1:
            return bc
        else:
            return None

    def get_strings_per_function(self, start_func):
        strings = []
        fs = ''
        func_obj = idaapi.get_func(start_func)  
        if func_obj:
            bc = self.save_existing_comments(start_func)
            if bc is not None:
                for string in [saved_comment for saved_comment in bc]:
                    strings.append(string)
            self.clear_comments(func_obj)
            for inst_list in idautils.Heads(start_func, idc.find_func_end(start_func)):
                try:
                    for string in [self.get_string_type(xref_addr) for xref_addr in idautils.DataRefsFrom(inst_list)]:
                        if len(string) > 2:
                            strings.append(string[2:-1])
                            self.string_counter += 1
                        else:
                            pass
                except StringException:
                    continue
    
            if strings:
                for c in strings:
                    if '\n' in c:
                        c = c.replace('\n', '')
                    fs += '"' + c + '" '
                idaapi.set_func_cmt(func_obj, 'STR {}# {}'.format(len(strings), fs), 1)
            
        else:
            print("func_obj return 0")
            

    def main(self):
        print("\n[+]Launching...")
        try:            
            for i in idautils.Functions():
                self.get_strings_per_function(i)
            print("\n[+]Well done! Added {} strings in {} functions".format(self.string_counter, idaapi.get_func_qty()))
        except KeyboardInterrupt:
            print("\n[+]Ended by user")


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

    def activate(self, ctx):
        sfc = Strings_fc()
        sfc.main()

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


class Strings_window(idaapi.plugin_t):
    flags = idaapi.PLUGIN_FIX
    comment = 'Associate functions'
    help = 'https://github.com/partoftheworlD/IDA7py_FunctionStringAssociate/'
    wanted_name = 'Strings Function Associate'
    wanted_hotkey = ""

    def init(self):
        try:
            self._install_plugin()
        except Exception:
            form = idaapi.get_current_widget()
        return idaapi.PLUGIN_KEEP

    def _install_plugin(self):
        self.init()

    def term(self):
        pass

    def run(self, arg=0):
        str_Handle = StringsHandler()
        str_Handle.activate(self)


def PLUGIN_ENTRY():
    return Strings_window()


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

content:
from ast import expr_context
from cmath import exp
import collections
from email.policy import default
from uuid import RESERVED_FUTURE
import idaapi
import idc
import ida_ua
import os
import json
import idautils
import ida_kernwin
import ida_name
import ida_hexrays
import ida_funcs
import traceback


icon = b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00*\x00\x00\x00&\x08\x06\x00\x00\x00\xb2\x01\t \x00\x00\x00\x01sRGB\x00\xae\xce\x1c\xe9\x00\x00\x00\x04gAMA\x00\x00\xb1\x8f\x0b\xfca\x05\x00\x00\x00\tpHYs\x00\x00\x0e\xc3\x00\x00\x0e\xc3\x01\xc7o\xa8d\x00\x00\x02TIDATXG\xcd\x98\xcd.\x03Q\x14\xc7\xef\x10\x14-\xd5\xf8Jj!x\x00+,\xc5\x03\x88\x07\xb0\xb0\xb2\xea\x82\'\xc0\x13\x90\xb0\x15\x8f\x80x\x00+k\x89\xd8\x93\xb0\x94\xb0\x13A\xc6\xf9\xb7skz\xe7\xcc\xf4\xdc\xf9\xec/\xf9\xa5\x9d\x99\xce\xbdg\xce\xdc9\xf7N\x15\xb1@\x1e\x93\xf3\xd8\xe81\xaa\xe4\x91\xf7\xd9\xe4\x9etI\x04\xdc\x0b \xb0=\xf2\x89\xbc\xc0\x0e\r\xb2\x89@\xe1;\xb9C\x16\x05\xfaF\x80:\x9ee\xb2\x83KR\x1f\x84\xc8\xf2:\x99\x17\xe8K\xdfY\xed\x01\x19\xc0\x9fU\xbf\xb8\x80\xc0U\xa5\x08\xda\xbe%\xcd~qg\xdbc\xd3\x04\xe3\xc1<A\x9b\xf6\xf8E\x80h\x13\x01q\xfd\xb1\xd9\xd4\xe0\n\xb8\x93\xfcF6 \x00}D\x05\x081FC\xb3\xa9A#\xdc\xc9~1\x96l\x1f8I\x80Zq\xdb\xfe\xa7.J\x04\xbcEF\x81\x00\xf1\x1bI\x80\x10\xe3U\x0c\x1a\xe6\x1a\t\x13\x0f\x1c7apOr7\xad+\x8d4\xab~\xf10"`tf\x96;\x898\xc7\x1at\xc65\x96\x95\x18\x1a\xb1\xa7q\xae\xbeee\xa2\xf2\x97WV\x91\xcd\xae\xe5\xa8\x1b\x81\xb1\xd6glK\x1dp\x1c\xb7\x9fd\x8ea\x01\x92\x98\xc0\xd4\xeaxn\x0e\x97;\xf6G\xb9:Tj\x9e\xc3\x1c\x13\x15w)\x81\xa9\x15\x9d\xdeL\xd5\xdd\xdd\xf2x\xc7~\xceF\xa5\xea\x9e\xd5f\xc2\x02Mu\xa5\x86+\x0etrM\x81\xbe\xcd-\xb9\x1b\xa5\x91\xc01\xed\xf6\xe8\x98\xfbZ_tO\'\xa6\xb9\xe3\xa8\xb1"h\xb8\x89\xf8 OZ_\x83\x9c\xd7f\xd5\xda\xd0\xb0\xb7\xf5\xcf\xca`I\x1d\x8dO\xaa\x92C\xb9\xe4\xc1\xea]\x844P\xb0O>\xb7\xbe\xb6x\xfc\xfeRw_\x9f\xea\x81>\x1b\xe5\xaa\x1aq\xfe\x9bCp\x8d\xcaD\xfb7/\xbf?\xde\x916W\x9e\x99\x80\xf1\xc4\xdd\xc28ZO\x95\xb6\xc4\x99ZMcM\x95\xb6\xd8.XLQ\xdc\xb3|c\xe8 IV\x93.\xbc\xad\x88;\xb5&Zx\xc4%\xce\x82%\xd7lj0\xce\xb8`\xc2Lu\xaa\xb4%\xea\xad\xd5\xb4\x90lj\xd8\xa9\x95\x11Sea\xd9\xd4\x1c\x92\\p~3/\xee\x12\x90\xa9\xa8r\x95Kq\x97\x82\xf1\xc7\x05\ts+\xee\x12\xc2\xb2Z\xe8\x03\x14\x86\xb9`I\xe5=(+\xfcY\xed\xc9ljtV\x0b-\xeeR\xe2\xfc\x81V\x08\x19dR\xa9?"\x80\x16\n\xa6\x0c\x13@\x00\x00\x00\x00IEND\xaeB`\x82'
icon_id = idaapi.load_custom_icon(data=icon, format="png")

class utils:
    def create_dummy_hexrays_const(value):
        num = ida_hexrays.cexpr_t()
        num.op = ida_hexrays.cot_num
        num.n = ida_hexrays.cnumber_t()
        num.n._value = value
        return num

    def get_negative_disass(op):
        # Get negative value based on the size of the operand
        if op.dtype == 0x0: # Byte
            return -(((op.value & 0xff) ^ 0xff) + 1)
        elif op.dtype == 0x1: # Word
            return -(((op.value & 0xffff) ^ 0xffff) + 1)
        elif op.dtype == 0x2: # Dword
            return -(((op.value & 0xffffffff) ^ 0xffffffff) + 1)
        elif op.dtype == 0x7: # Qword
            return -((op.value ^ 0xffffffffffffffff) + 1)

    def get_func_name(ea):
        # Get pretty function name
        func_name = utils.get_pretty_func_name(idc.get_func_name(ea))
        if not func_name:
            func_name = utils.get_pretty_func_name(idc.get_name(ea))
        return func_name

    def get_pretty_func_name(name):
        # Demangle function name
        demangled_name = idc.demangle_name(name, idc.get_inf_attr(idc.INF_SHORT_DN))
        # Return as is if failed to demangle
        if not demangled_name:
            return name
        # Cut arguments
        return demangled_name[:demangled_name.find("(")]

    def prep_func_name(name):
        if name[0] != "." and name[0] != "_":
            # Name does not start with dot or underscore
            return [name,f".{name}",f"_{name}"]
        else:
            return [name[1:],f".{name[1:]}",f"_{name[1:]}"]

class null_after_visitor(ida_hexrays.ctree_visitor_t):
    def __init__(self,decompiled_function,call_xref,func_name,matched):
        self.found_call = False
        ida_hexrays.ctree_visitor_t.__init__(self, ida_hexrays.CV_FAST) # CV_FAST does not keep parents nodes in CTREE
        self.decompiled_function = decompiled_function
        self.call_xref = call_xref
        self.func_name = func_name
        self.insn_counter = 0
        self.matched = matched

    def visit_insn(self, i):
        if self.found_call:
            self.insn_counter += 1
        return 0

    def visit_expr(self, e):
        if e.op == ida_hexrays.cot_call:
            xref_func_name = utils.get_func_name(e.x.obj_ea)
            if not xref_func_name:
                xref_func_name = idc.get_name(e.x.obj_ea)
            if xref_func_name.lower() == self.func_name.lower() and e.ea == self.call_xref:
                self.found_call = True
        if self.found_call and self.insn_counter < 2:
            if e.op == ida_hexrays.cot_asg:
                # The expression is assignment, check the right side of the assign
                if e.y.op == ida_hexrays.cot_num:
                    # The right side is number
                    if e.y.numval() == 0:
                        # Set to null
                        self.matched["set_to_null"] = True
        return 0



class VulFiScanner:
    def __init__(self,custom_rules=None):
        # Init class-wide variables
        self.functions_list = []
        if not custom_rules:
            with open(os.path.join(os.path.dirname(os.path.realpath(__file__)),"vulfi_rules.json"),"r") as rules_file:
                self.rules = json.load(rules_file)
        else:
            self.rules = custom_rules
        with open(os.path.join(os.path.dirname(os.path.realpath(__file__)),"vulfi_prototypes.json"),"r") as proto_file:
            self.prototypes = json.load(proto_file)
        # get pointer size
        info = idaapi.get_inf_structure()
        if info.is_64bit():
            self.ptr_size = 8
        elif info.is_32bit():
            self.ptr_size = 4
        else:
            self.ptr_size = 2
        # Get endianness
        self.endian = "big" if idaapi.cvar.inf.is_be() else "little"
        # Check whether Hexrays can be used
        if not ida_hexrays.init_hexrays_plugin():
            self.hexrays = False
            self.strings_list = idautils.Strings()
        else:
            self.hexrays = True
            #self.strings_list = idautils.Strings()

    def start_scan(self,ignore_addr_list):
        results = []
        ida_kernwin.show_wait_box("VulFi scan running ... ")
        self.prepare_functions_list()
        for rule in self.rules:
            try:
                if rule["function_names"] == ["Array Access"]:
                    xrefs_dict = self.get_array_accesses()
                elif rule["function_names"] == ["Loop Check"]:
                    xrefs_dict = self.get_loops()
                else:
                    xrefs_dict = self.find_xrefs_by_name(rule["function_names"],rule["wrappers"])
            except:
                ida_kernwin.warning("This does not seem like a correct rule file. Aborting scan.")
                return None
            for scanned_function in xrefs_dict:
            # For each function in the rules
                skip_count = 0
                counter = 0 # For UI only
                total_xrefs = len(xrefs_dict[scanned_function]) # For UI only
                for scanned_function_xref_tuple in xrefs_dict[scanned_function]:
                    counter += 1
                    if ida_kernwin.user_cancelled():
                        print("[VulFi] Scan canceled!")
                        ida_kernwin.hide_wait_box()
                        return None
                    if skip_count > 0:
                        skip_count -= 1
                        continue
                    param = []
                    param_count = 0
                    function_call = None
                    scanned_function_xref = scanned_function_xref_tuple[0] # Address
                    scanned_function_name = scanned_function_xref_tuple[1] # Name
                    scanned_function_display_name = scanned_function_xref_tuple[2] # Display Name
                    # Update progress bar
                    ida_kernwin.replace_wait_box(f"Scanning {scanned_function_display_name} ({counter}/{total_xrefs})")
                    # If this rule is already verified for this address, move on
                    if f"{rule['name']}_{hex(scanned_function_xref)}" in ignore_addr_list:
                        continue
                    # For each xref to the function in the rules
                    # If ida_hexrays can be used eval the conditions in the rules file
                    if rule["function_names"] == ["Array Access"]:
                        param = [VulFiScanner.Param(self,scanned_function_xref_tuple[3],scanned_function_xref,scanned_function_name),VulFiScanner.Param(self,scanned_function_xref_tuple[4],scanned_function_xref,scanned_function_name)]
                    elif rule["function_names"] == ["Loop Check"]:
                        param = [VulFiScanner.Param(self,scanned_function_xref_tuple[3],scanned_function_xref,scanned_function_name),VulFiScanner.Param(self,scanned_function_xref_tuple[4],scanned_function_xref,scanned_function_name)]
                    else:
                        params_raw = self.get_xref_parameters(scanned_function_xref,scanned_function_name)
                        if params_raw is None:
                            if self.hexrays:
                                # Likely decompiler failure, we have to skip
                                continue
                            else:
                                # Params were not found, lets mark all XREFS
                                params_raw = []
                        for p in params_raw:
                            param.append(VulFiScanner.Param(self,p,scanned_function_xref,scanned_function_name))
                        param_count = len(param)

                    function_call = VulFiScanner.FunctionCall(self,scanned_function_xref,scanned_function_name)
                    # Name of the function where the xref is located
                    found_in_name = utils.get_func_name(scanned_function_xref)
                    priority = ""
                    # Every xref will be marked as Info in case that params fetch fails
                    try:
                        if not param:
                            priority = "Info"
                        elif eval(rule["mark_if"]["High"],dict(self=self, param=param,param_count=param_count,function_call=function_call)):
                            priority = "High"
                        elif eval(rule["mark_if"]["Medium"],dict(self=self, param=param,param_count=param_count,function_call=function_call)):
                            priority = "Medium"
                        elif eval(rule["mark_if"]["Low"],dict(self=self, param=param,param_count=param_count,function_call=function_call)):
                            priority = "Low"

                    except IndexError:
                        # Decompiler output has fewer parameters than the function prototype
                        # Mark the issue with Info priority
                        priority = "Info"
                        #results.append(list(VulFi.result_window_row(rule["name"],scanned_function_display_name,found_in_name,hex(scanned_function_xref),"Not Checked","Low","")))
                    except Exception:
                        print(traceback.format_exc())
                        ida_kernwin.warning(f"The rule \"{rule}\" is not valid!")
                        continue
                    # If the rule matched and is not wrapped:
                    #if priority and found_in_name and not "wrapped" in scanned_function_display_name:
                    if priority and not "wrapped" in scanned_function_display_name:
                        results.append(list(VulFi.result_window_row(rule["name"],scanned_function_display_name,found_in_name,hex(scanned_function_xref),"Not Checked",priority,"")))
                    elif "wrapped" in scanned_function_display_name and priority:
                        skip_count = 0 # rule for the wrapped function matched so no need to skip calls to the wrapper
                    elif "wrapped" in scanned_function_display_name and priority == "":
                        # Rule for wrapped function did not match, skip the wrapper
                        skip_count = int(scanned_function_display_name[scanned_function_display_name.find("wrapped:") + 8:-1])
        ida_kernwin.hide_wait_box()
        return results
    
    def get_loops(self):
         # Use the trick to extract loop paramters as members of the tuple
         # param[0] will be the counter variable and param[1] will be the check value
         loop_check_list = {"Loop Check":[]}
         if self.hexrays:
            for func in self.functions_list:
                try:
                    decompiled_function = ida_hexrays.decompile(func)
                    code = decompiled_function.pseudocode
                    for tree_item in decompiled_function.treeitems:
                        if tree_item.op == ida_hexrays.cit_while:
                            #print(f"[*] Found 'while' loop at: {hex(tree_item.ea)} ({ida_name.get_ea_name(func)})")
                            op = tree_item.to_specific_type.cwhile.expr.op
                            if (op >= ida_hexrays.cot_land and op <= ida_hexrays.cot_ult) or op in [ida_hexrays.cot_var, ida_hexrays.cot_lnot]:
                                # Not an infinte loop
                                counter, check_val = self.parse_condition(tree_item.to_specific_type.cwhile.expr)
                                loop_check_list["Loop Check"].append((tree_item.ea,"Loop Check","Loop Check",counter,check_val))
                            else:
                                counter, check_val = self.get_loop_check_val_break(tree_item.to_specific_type.cwhile.body)
                                loop_check_list["Loop Check"].append((tree_item.ea,"Loop Check","Loop Check",counter,check_val))
                        elif tree_item.op == ida_hexrays.cit_for:
                            #print(f"[*] Found 'for' loop at: {hex(tree_item.ea)} ({ida_name.get_ea_name(func)})")
                            if tree_item.to_specific_type.cfor.expr.op == ida_hexrays.cot_empty:
                                # Handle endless FOR loop
                                body = tree_item.to_specific_type.cfor.body
                                counter, check_val = self.get_loop_check_val_break(body)
                                loop_check_list["Loop Check"].append((tree_item.ea,"Loop Check","Loop Check",counter,check_val))
                            else:
                                counter = tree_item.to_specific_type.cfor.init.x
                                loop_check_list["Loop Check"].append((tree_item.ea,"Loop Check","Loop Check",counter,self.get_loop_check_val(counter,tree_item.to_specific_type.cfor.expr)))
                        elif tree_item.op == ida_hexrays.cit_do:
                            op = tree_item.to_specific_type.cdo.expr.op
                            if op >= ida_hexrays.cot_land and op <= ida_hexrays.cot_ult or op in [ida_hexrays.cot_var, ida_hexrays.cot_lnot]:
                                # Not an infinte loop
                                counter, check_val = self.parse_condition(tree_item.to_specific_type.cdo.expr)
                                loop_check_list["Loop Check"].append((tree_item.ea,"Loop Check","Loop Check",counter,check_val))
                            else:
                                counter, check_val = self.get_loop_check_val_break(tree_item.to_specific_type.cdo.body)
                                loop_check_list["Loop Check"].append((tree_item.ea,"Loop Check","Loop Check",counter,check_val))
                except Exception as e:
                    print(e)

         return loop_check_list
    
    def parse_condition(self,condition):
        if condition.op == ida_hexrays.cot_lnot:
            #(!x)
            return condition.x, utils.create_dummy_hexrays_const(0)
        if condition.op == ida_hexrays.cot_var or condition.op == ida_hexrays.cot_call or condition.op == ida_hexrays.cot_cast:
            #(x)
            return condition.x, utils.create_dummy_hexrays_const(1)
        if condition.op >= ida_hexrays.cot_eq and condition.op <= ida_hexrays.cot_ult:
            #(x > y)
            return condition.x, condition.y
        if condition.op == ida_hexrays.cot_lor or condition.op == ida_hexrays.cot_land:
            x_cond = self.parse_condition(condition.x)
            y_cond = self.parse_condition(condition.y)
            if y_cond[1].n:
                return y_cond
            return x_cond
    
    def get_loop_check_val(self,counter,comparison):
        x,y = self.parse_condition(comparison)
        if x == counter:
            return y
        # IDA places constant into the y operand for normal comparisons, for everything else we should return X to avoid marking that as a const
        return x

    def get_loop_check_val_break(self,loop_body):
        bd = loop_body.cblock.begin()
        while bd != loop_body.cblock.end():
            if bd.cur.op == ida_hexrays.cit_if:
                break_if = self.find_break(bd.cur.cif)
                if break_if:
                    return self.parse_condition(break_if.expr)
            bd.next()
        
        return None, None

    def find_break(self,if_insn):
        # TODO switch-case
        bt = if_insn.ithen.cblock.begin()
        while bt != if_insn.ithen.cblock.end():
            if bt.cur.op == ida_hexrays.cit_if:
                found_breaks = self.find_break(bt.cur.cif)
                if found_breaks:
                    return found_breaks
            elif bt.cur.op == ida_hexrays.cit_break or bt.cur.op == ida_hexrays.cit_return:
                return if_insn
            bt.next()
            
        if if_insn.ielse:  
            be = if_insn.ielse.cblock.begin()
            while be != if_insn.ielse.cblock.end():
                if be.cur.op == ida_hexrays.cit_if:
                    found_breaks = self.find_break(be.cur.cif)
                    if found_breaks:
                        return found_breaks
                elif be.cur.op == ida_hexrays.cit_break or bt.cur.op == ida_hexrays.cit_return:
                    return if_insn
                be.next()
        return None

    
    def get_array_accesses(self):
        array_access_list = {"Array Access":[]}
        if self.hexrays:
            for func in self.functions_list:
                try:
                    decompiled_function = ida_hexrays.decompile(func)
                    code = decompiled_function.pseudocode
                    for citem in decompiled_function.treeitems:
                        if citem.op == ida_hexrays.cot_idx:
                            # Uses a trick by adding the index expression as a hidden 4th member of tuple
                            array_access_list["Array Access"].append((citem.ea,"Array Access","Array Access",citem.to_specific_type.x,citem.to_specific_type.y))
                except Exception as e:
                    pass
        return array_access_list

    def prepare_functions_list(self):
        self.functions_list = []
        # Gather all functions in all segments
        for segment in idautils.Segments():
            self.functions_list.extend(list(idautils.Functions(idc.get_segm_start(segment),idc.get_segm_end(segment))))

        # Gather imports
        def imports_callback(ea, name, ordinal):
            self.functions_list.append(ea)
            return True

        # For each import
        number_of_imports = idaapi.get_import_module_qty()
        for i in range(0, number_of_imports):
            idaapi.enum_import_names(i, imports_callback)

    # Returns a list of all xrefs to the function with specified name
    def find_xrefs_by_name(self,func_names,wrappers):
        xrefs_list = {}
        insn = ida_ua.insn_t()
        function_names = []
        for name in func_names:
            function_names.extend(utils.prep_func_name(name))

        # Convert function names to expected format (demangle them if possible)
        function_names = list(map(utils.get_pretty_func_name, function_names))

        # For all addresses of functions and imports check if it is function we are looking for
        for function_address in self.functions_list:
            current_function_name = utils.get_func_name(function_address)
            if not current_function_name:
                current_function_name = ida_name.get_ea_name(function_address)
            # If the type is not defined and the name of the function is known set the type
            prototype = self.prototypes.get(current_function_name.lower(), None)
            if prototype is not None:
                idc.SetType(function_address, prototype)
                idaapi.auto_wait()


            # Seach function "as is" and "ingnore case"
            if current_function_name not in function_names:
                current_function_name = current_function_name.lower()
                if current_function_name not in function_names:
                    continue

            # This is the function we are looking for
            if current_function_name not in xrefs_list.keys():
                xrefs_list[current_function_name] = []
            for xref in idautils.XrefsTo(function_address):
                # This rules out any functions within those that we are already looking at (wrapped)
                # It also makes sure that functions that are not xrefed within another function are not displayed
                xref_func_name = utils.get_func_name(xref.frm)
                if not xref_func_name or xref_func_name.lower() in function_names:
                    continue
                # Make sure that the instrution is indeed a call
                if ida_ua.decode_insn(insn,xref.frm) != idc.BADADDR:
                    if (insn.get_canon_feature() & 0x2 == 0x2 or # The instruction is a call
                    (insn.get_canon_feature() & 0xfff == 0x100 and insn.Op1.type in [idc.o_near,idc.o_far])): # The instruction uses first operand which is of type near/far immediate address
                        xref_tuple = (xref.frm,current_function_name,current_function_name)
                        if wrappers:
                            # If we should look for wrappers, do not includ the wrapped xref
                            retrieved_wrappers = self.get_wrapper_xrefs(xref.frm,current_function_name)
                            if not retrieved_wrappers:
                                # No wrappers
                                if not xref_tuple in xrefs_list[current_function_name]:
                                    xrefs_list[current_function_name].append(xref_tuple)
                            else:
                                # Wrappers were found
                                wrapped_tupple = (xref_tuple[0],xref_tuple[1],f"{xref_tuple[2]} (wrapped:{len(retrieved_wrappers)})")
                                xrefs_list[current_function_name].append(wrapped_tupple) # this makes sure that the wrapped function is right before its wrappers
                                xrefs_list[current_function_name].extend(retrieved_wrappers)
                        else:
                            if not xref_tuple in xrefs_list[current_function_name]:
                                xrefs_list[current_function_name].append(xref_tuple)
        return xrefs_list

    # Check if the xref can be a wrapper
    def get_wrapper_xrefs(self,xref,current_function_name):
        if self.hexrays:
            # Hexrays avaialble, we can use decompiler
            return self.get_wrapper_xrefs_hexrays(xref,current_function_name)
        else:
            # Hexrays not available, decompiler cannot be used
            return self.get_wrapper_xrefs_disass(xref)

    def get_wrapper_xrefs_hexrays(self,function_xref,current_function_name):
        wrapper_xrefs = []
        try:
            decompiled_function = ida_hexrays.decompile(function_xref)
        except:
            return wrapper_xrefs
        if decompiled_function:
            # Decompilation is fine
            code = decompiled_function.pseudocode
            for tree_item in decompiled_function.treeitems:
                if tree_item.ea == function_xref and tree_item.op == ida_hexrays.cot_call:
                    # Get called function name
                    xref_func_name = utils.get_func_name(tree_item.to_specific_type.x.obj_ea).lower()
                    if not xref_func_name:
                        xref_func_name = idc.get_name(tree_item.to_specific_type.x.obj_ea).lower()
                    if xref_func_name == current_function_name.lower():
                        # This is the correct call
                        lvars = list(decompiled_function.get_lvars()) # Get lvars
                        arg_objects = list(tree_item.to_specific_type.a)
                        found = True
                        while arg_objects:
                            current_obj = arg_objects.pop(0)
                            if current_obj is None:
                                continue
                            if current_obj.op == ida_hexrays.cot_var:
                                # if variable is not an arg_var we do not have a wrapper
                                if not lvars[current_obj.v.idx].is_arg_var:
                                    found = False
                            else:
                                arg_objects.extend([current_obj.to_specific_type.x,current_obj.to_specific_type.y])
                        # The function is likely a wrapper, get XREFs to it
                        if found:
                            for wrapper_xref in idautils.XrefsTo(decompiled_function.entry_ea):
                                wrapper_xref_func_name = utils.get_func_name(decompiled_function.entry_ea)
                                wrapper_xrefs.append((wrapper_xref.frm,wrapper_xref_func_name,f"{wrapper_xref_func_name} ({current_function_name} wrapper)"))
        return wrapper_xrefs


    # There probably is no architecture-agnostic way on how to do this without decompiler
    def get_wrapper_xrefs_disass(self,xref):
        return []

    # Wrapper for getting function params
    def get_xref_parameters(self,function_xref,scanned_function):
        if self.hexrays:
            # Hexrays avaialble, we can use decompiler
            return self.get_xref_parameters_hexrays(function_xref,scanned_function)
        else:
            # Hexrays not available, decompiler cannot be used
            return self.get_xref_parameters_disass(function_xref)

    # Returns list of ordered paramters from disassembly
    def get_xref_parameters_disass(self,function_xref):
        params_list = []
        # Requires the type to be already assigned to functions
        try:
            for param_ea in idaapi.get_arg_addrs(function_xref):
                if param_ea != idc.BADADDR:
                    param_insn = ida_ua.insn_t()
                    # decode the instruction linked to the parameter
                    if ida_ua.decode_insn(param_insn,param_ea) != idc.BADADDR:
                        if param_insn.get_canon_feature() & 0x00100 == 0x100:
                            # First operand - push
                            params_list.append(param_insn.Op1)
                        elif param_insn.get_canon_feature() & 0x00200 == 0x200:
                            # Second operands - mov/lea
                            params_list.append(param_insn.Op2)
                        elif param_insn.get_canon_feature() & 0x00400 == 0x400:
                            # Third operand
                            params_list.append(param_insn.Op3)
        except:
            return None
        return params_list

    # Returns an ordered list of workable object that represent each parameter of the function from decompiled code
    def get_xref_parameters_hexrays(self,function_xref,scanned_function):
        # Decompile function and find the call
        try:
            decompiled_function = ida_hexrays.decompile(function_xref)
        except:
            return None
        if decompiled_function is None:
            # Decompilation failed
            return None
        index = 0
        code = decompiled_function.pseudocode
        for tree_item in decompiled_function.treeitems:
            if tree_item.ea == function_xref and tree_item.op == ida_hexrays.cot_call:
                xref_func_name = utils.get_func_name(tree_item.to_specific_type.x.obj_ea).lower()
                if not xref_func_name:
                    xref_func_name = idc.get_name(tree_item.to_specific_type.x.obj_ea).lower()
                if xref_func_name == scanned_function.lower():
                    return list(tree_item.to_specific_type.a)
            index += 1
        # Call not found :(
        return None

    class Param:
        def __init__(self,scanner,param,call_xref,scanned_function):
            self.scanner_instance = scanner
            self.param = param
            self.call_xref = call_xref
            self.scanned_function = scanned_function


        def used_as_index(self):
            if self.scanner_instance.hexrays:
                return self.used_as_index_hexrays()
            else:
                return False
            
        def used_as_index_hexrays(self):
            decompiled_function = ida_hexrays.decompile(self.call_xref)
            if decompiled_function:
                code = decompiled_function.pseudocode
                for citem in decompiled_function.treeitems:
                    if citem.op == ida_hexrays.cot_idx:
                        if citem.to_specific_type.y == self.param:
                            return True
            return False

        def is_constant(self):
            if self.string_value() == "" and self.number_value() == None:
                if self.param and self.param.op == ida_hexrays.cot_ref:
                    return False
                asgs = self.__get_var_assignments()
                if asgs: # asgs will be empty with no hexrays
                    for asg in asgs:
                        if self.__is_before_call(asg.ea):
                            if self.string_value(asg.y) == "" and self.number_value(asg.y) == None:
                                # One of the assigns is non-const
                                return False
                    return True
                return False
            else:
                return True

        # Returns True if the param is used in any function call specified in the "function_list" parameter
        def used_in_call_before(self,function_list):
            if self.scanner_instance.hexrays:
                return self.used_in_call_before_hexrays(function_list)
            else:
                return self.used_in_call_before_disass(function_list)

        def used_in_call_before_hexrays(self,function_list):
            # get all calls with the parameter
            calls = self.__get_var_arg_calls()
            # prep function list
            tmp_fun_list = []
            for fun in function_list:
                tmp_fun_list.extend(utils.prep_func_name(fun))
            for call in calls:
                if utils.get_func_name(call.x.obj_ea) in tmp_fun_list and self.__is_before_call(call.ea):
                    return True
            return False

        def used_in_call_before_disass(self,function_list):
            return False

        def used_in_call_after(self,function_list):
            if self.scanner_instance.hexrays:
                return self.used_in_call_after_hexrays(function_list)
            else:
                return self.used_in_call_after_disass(function_list)

        def used_in_call_after_hexrays(self,function_list):
            # get all calls with the parameter
            calls = self.__get_var_arg_calls()
            # prep function list
            tmp_fun_list = []
            for fun in function_list:
                tmp_fun_list.extend(utils.prep_func_name(fun))
            for call in calls:
                if utils.get_func_name(call.x.obj_ea) in tmp_fun_list and not self.__is_before_call(call.ea):
                    return True
            return False
        
        def is_sign_compared(self):
            if self.scanner_instance.hexrays:
                return self.is_sign_compared_hexrays()
            else:
                return False


        def is_sign_compared_hexrays(self):
            decompiled_function = ida_hexrays.decompile(self.call_xref)
            if decompiled_function:
                code = decompiled_function.pseudocode
                for citem in decompiled_function.treeitems:
                    if citem.ea == self.call_xref:
                        parent = decompiled_function.body.find_parent_of(citem)
                        while parent:
                            if parent.op == ida_hexrays.cit_if:
                                comps = self.__get_signed_comparisons(parent)
                                if self.__is_var_used_in_comparison(comps):
                                    return True
                            parent = decompiled_function.body.find_parent_of(parent)
            return False
        
        def __is_var_used_in_comparison(self,comp_list):
            if self.param.op == ida_hexrays.cot_cast:
                param = self.param.x
            else:
                param = self.param
            ops = comp_list.copy()
            while ops:
                op = ops.pop()
                if not op:
                    continue
                if param == op:
                    return True
                ops.extend([op.x,op.y])
            return False
        
        def __get_signed_comparisons(self,citem):
            signed_comparisons = []
            signed_ops = [ida_hexrays.cot_sge, ida_hexrays.cot_sle, ida_hexrays.cot_sgt, ida_hexrays.cot_slt]
            exprs = [citem.to_specific_type.cif.expr]
            while exprs:
                expr = exprs.pop()
                if not expr:
                    continue
                if expr.op in signed_ops:
                    signed_comparisons.append(expr)
                else:
                    exprs.extend([expr.x,expr.y])
            return signed_comparisons


        def used_in_call_after_disass(self,function_list):
            return False

        # Simple check whether the given EA is before the call
        def __is_before_call(self,ea):
            func = idaapi.get_func(ea)
            flow = idaapi.FlowChart(func)
            call_block = None
            asg_block = None
            checked_blocks = []
            # Get block of the assignemnt and block of the call
            for block in flow:
                if ea >= block.start_ea and ea <= block.end_ea:
                    asg_block = block
                if self.call_xref >= block.start_ea and self.call_xref <= block.end_ea:
                    call_block = block
                    checked_blocks.append(block.start_ea)
            # If they are in the same block and asg ea is smaller then call_xref ea return True
            if call_block == asg_block:
                if ea < self.call_xref:
                    return True
            else:
                # Blocks are different
                call_preds = list(call_block.preds())
                while call_preds:
                    current_pred = call_preds.pop(0)
                    # Prevent endless loops
                    if current_pred.start_ea in checked_blocks:
                        continue
                    checked_blocks.append(current_pred.start_ea)
                    # Check if we matched the given assign block
                    if current_pred.start_ea == asg_block.start_ea:
                        return True
                    call_preds.extend(list(current_pred.preds()))
            return False

        def __get_var_arg_calls(self):
            calls = []
            # Parameter is cast, get x
            if self.param.op == ida_hexrays.cot_cast:
                param = self.param.x
            else:
                param = self.param
            decompiled_function = ida_hexrays.decompile(self.call_xref)
            code = decompiled_function.pseudocode
            for citem in decompiled_function.treeitems:
                if citem.op == ida_hexrays.cot_call and citem.ea != self.call_xref: # skip calls we are tracing
                    # Potentially interesting call
                    for a in citem.to_specific_type.a:
                        expressions = [a,a.x,a.y,a.z]
                        while expressions:
                            current_expr = expressions.pop(0)
                            if current_expr:
                                if param == current_expr:
                                    # Call operation, add to array
                                    calls.append(citem.to_specific_type)
                                    break # we can break the loop as the variable was found within the arguments
                                expressions.extend([current_expr.x, current_expr.y, current_expr.z])
            return calls

        # Returns list of assign expressions for better accuracy
        def __get_var_assignments(self):
            asg = []
            # Parameter is cast, get x
            if self.param.op == ida_hexrays.cot_cast:
                param = self.param.x
            else:
                param = self.param
            decompiled_function = ida_hexrays.decompile(self.call_xref)
            code = decompiled_function.pseudocode
            for citem in decompiled_function.treeitems:
                if param == citem.to_specific_type:
                    parent = decompiled_function.body.find_parent_of(citem)
                    if parent.op >= ida_hexrays.cot_asg and parent.op <= ida_hexrays.cot_asgumod:
                        # Assign operation, add to array
                        asg.append(parent.to_specific_type)
            return asg

        def string_value(self,expr=None):
            if not expr:
                expr = self.param
            if not expr:
                return None
            if self.scanner_instance.hexrays: # hexrays
                string_val = idc.get_strlit_contents(expr.obj_ea)
                if string_val:
                    return string_val.decode()
                # If it is a cast (could happen)
                elif expr.op == ida_hexrays.cot_cast:
                    # If casted op is object
                    if expr.x.op == ida_hexrays.cot_obj:
                        # If that object points to a string
                        string_val = idc.get_strlit_contents(expr.x.obj_ea)
                        if string_val:
                            return string_val.decode()
                elif expr.op == ida_hexrays.cot_ref:
                    string_val = idc.get_strlit_contents(expr.x.obj_ea)
                    if string_val:
                        return string_val.decode()
                    if expr.x.op == ida_hexrays.cot_idx:
                        string_val = idc.get_strlit_contents(expr.x.x.obj_ea)
                        if string_val:
                            return string_val.decode()
                    # Check whether we are looking at CFString
                    if expr.x.obj_ea != idc.BADADDR:
                        c_str_pointer_value = idc.get_bytes(expr.x.obj_ea +  2* self.scanner_instance.ptr_size, self.scanner_instance.ptr_size)
                        tmp_ea = int.from_bytes(c_str_pointer_value,byteorder=self.scanner_instance.endian)
                        c_str_len = int.from_bytes(idc.get_bytes(expr.x.obj_ea +  3* self.scanner_instance.ptr_size, self.scanner_instance.ptr_size),byteorder=self.scanner_instance.endian)
                        c_string_value = idc.get_strlit_contents(tmp_ea)
                        # Having a string at this position and its length following the pointer suggests CFString struct
                        if c_string_value and len(c_string_value) == c_str_len:
                            # If this evaluates to string we have a constant
                            return c_string_value.decode()
                else:
                    # not a direct string
                    byte_value = idc.get_bytes(expr.obj_ea,self.scanner_instance.ptr_size)
                    if byte_value:
                        tmp_ea = int.from_bytes(byte_value,byteorder=self.scanner_instance.endian)
                        string_val = idc.get_strlit_contents(tmp_ea)
                        if string_val:
                            # If this evaluates to string we have a constant
                            return string_val.decode()
            else: # No hexrays
                if expr.type == 0x2:
                    # Reference
                    if idc.get_strlit_contents(expr.addr):
                        return idc.get_strlit_contents(expr.addr).decode()
                    else:
                        # Reference to reference
                        addr = int.from_bytes(idc.get_bytes(expr.addr, self.scanner_instance.ptr_size),byteorder=self.scanner_instance.endian)
                        if idc.get_strlit_contents(addr):
                            return idc.get_strlit_contents(addr).decode()

                if expr.type == 0x5:
                    if idc.get_strlit_contents(expr.value):
                        return idc.get_strlit_contents(expr.value).decode()
                # Reverse appoach of going from strings to calls
                for c_string in self.scanner_instance.strings_list:
                    for str_xref in idautils.XrefsTo(c_string.ea):
                        func = idaapi.get_func(str_xref.frm)
                        if func:
                            if func.start_ea == idaapi.get_func(self.call_xref).start_ea:
                                # The xref and the call are in the same function
                                # look if those are used within 5 instructions
                                if (str_xref.frm < self.call_xref):
                                    # XREF to STR is before the call to XREF
                                    instr_list = list(idautils.Heads(str_xref.frm,self.call_xref))
                                    if len(instr_list) <= 10:
                                        # No more then 10 instructions between the str XREF and the call
                                        for head in instr_list:
                                            current_insn = ida_ua.insn_t()
                                            if ida_ua.decode_insn(current_insn,head) != idc.BADADDR:
                                                if head != self.call_xref and current_insn.get_canon_feature() & 0x2 == 0x2:
                                                    # If it is not a target call but it is a call its false
                                                    return ""
                                        # If we survived the loop it is True
                                        return str(c_string)
            return ""

        def number_value(self,expr=None):
            if not expr:
                expr = self.param
            if not expr:
                return None
            if self.scanner_instance.hexrays: # hexrays
                # If it is number directly
                if expr.op == ida_hexrays.cot_num:
                    return expr.n._value
                # if it is a float
                elif expr.op == ida_hexrays.cot_fnum:
                    return expr.fpc.fnum.float
                # If it is a cast
                elif expr.op == ida_hexrays.cot_cast:
                    if expr.x.op == ida_hexrays.cot_num:
                        return expr.x.n._value
                    elif expr.x.op == ida_hexrays.cot_fnum:
                        return expr.x.fpc.fnum.float
            else: # No hexrays
                # self.param is op_t
                # 0x5 is immediate
                if expr.type == 0x5:
                    return expr.value
            return None

        def is_const_number(self,expr=None):
            if not expr:
                expr = self.param
            if self.scanner_instance.hexrays:
                if expr.op == ida_hexrays.cot_num or expr.op == ida_hexrays.cot_fnum:
                    return True
                elif expr.op == ida_hexrays.cot_cast and expr.x.op == ida_hexrays.cot_num:
                    return True
            else:
                if expr.type == 0x5:
                    return True
            return False

        # Wrapper for the basic UAF filter
        def set_to_null_after_call(self):
            if self.scanner_instance.hexrays:
                # Hexrays avaialble, we can use decompiler
                return self.set_to_null_after_call_hexrays()
            else:
                # Hexrays not available, decompiler cannot be used
                return self.set_to_null_after_call_disass()

        # Simple function that just checks whether the call is followed by an isntruction that sets the parameter to null
        # Note that this is very primitive to keep sort of architecture agnostic approach
        def set_to_null_after_call_disass(self):
            call_insn = ida_ua.insn_t()
            following_insn = ida_ua.insn_t()
            # First get the call instruction
            if ida_ua.decode_insn(call_insn,self.call_xref) != idc.BADADDR:
                # Now use the call instruction to get EA of next insn
                if ida_ua.decode_insn(following_insn,call_insn.ea + call_insn.size) != idc.BADADDR:
                    # following_insn can be used to evaluate whether Op2 is a constant or not
                    if following_insn.Op2.type == 0x5 and following_insn.Op2.value == 0x0:
                        # This should cover most of the cases where a mov instuction is used
                        return True
            return False

        # Using hexrays to figure out whether the function was followed by a null-set
        def set_to_null_after_call_hexrays(self):
            # get size of the call instruction, go past it and get expression that is linked to that address?
            matched = {"set_to_null":False}
            try:
                decompiled_function = ida_hexrays.decompile(self.call_xref)
            except:
                return matched["set_to_null"]
            if decompiled_function is None:
                # Decompilation failed
                return None
            custom_visitor = null_after_visitor(decompiled_function,self.call_xref,self.scanned_function,matched)
            custom_visitor.apply_to(decompiled_function.body, None)

            return matched["set_to_null"]



    class FunctionCall:
        def __init__(self,scanner,call_xref,scanned_function):
            self.scanner_instance = scanner
            self.call_xref = call_xref
            self.scanned_function = scanned_function

        
        def reachable_from(self,function_name):
            functions = [idaapi.get_func(self.call_xref)]
            checked_xrefs = []
            while functions:
                current_function = functions.pop(0)
                if current_function:
                    if utils.get_func_name(current_function.start_ea) in utils.prep_func_name(function_name):
                        return True
                    for xref in idautils.XrefsTo(current_function.start_ea):
                        if xref.frm not in checked_xrefs:
                            functions.append(idaapi.get_func(xref.frm))
                            checked_xrefs.append(xref.frm)
            return False

        # Check whether the return value of a function is part of some comparison (verification)
        def return_value_checked(self,check_val = None):
            if self.scanner_instance.hexrays:
                # Hexrays avaialble, we can use decompiler
                return self.return_value_checked_hexrays(check_val)
            else:
                # Hexrays not available, decompiler cannot be used
                return self.return_value_checked_disass(check_val)

        # Check if the return value of a function is verified after the call
        def return_value_checked_disass(self,check_val):
            # To stay architecture agnostic, this simply checks whether there is conditional jump within 5 instructions after the call
            for basic_block in idaapi.FlowChart(idaapi.get_func(self.call_xref)):
                if self.call_xref >= basic_block.start_ea and self.call_xref < basic_block.end_ea and len(list(basic_block.succs())) > 1:
                    # xref call belongs to this block
                    insn_counter = 0
                    current_insn = ida_ua.insn_t()
                    if ida_ua.decode_insn(current_insn,self.call_xref) != idc.BADADDR:
                        while insn_counter < 5:
                            following_insn = ida_ua.insn_t()
                            if ida_ua.decode_insn(following_insn,current_insn.ea + current_insn.size) != idc.BADADDR:
                                current_insn = following_insn
                                # Most of the compare instructions
                                if current_insn.get_canon_feature() & 0xffff == 0x300:
                                    if check_val is not None:
                                        for op in current_insn.ops:
                                            if op.type == 0x5:
                                                negative_value = utils.get_negative_disass(op)
                                                if op.value == check_val or negative_value == check_val:
                                                    return True
                                                else:
                                                    return False
                                    else:
                                        # Likely comparison instruction hit
                                        return True
                                insn_counter += 1
                            else:
                                break
                            # Jumped out of the block within 5 instructions
                            if current_insn.ea >= basic_block.end_ea:
                                return True
            return False

        def return_value_checked_hexrays(self,check_val):
            try:
                decompiled_function = ida_hexrays.decompile(self.call_xref)
            except:
                return None
            if decompiled_function is None:
                # Decompilation failed
                return None
            code = decompiled_function.pseudocode
            index = 0
            for tree_item in decompiled_function.treeitems:
                if tree_item.ea == self.call_xref and tree_item.op == ida_hexrays.cot_call:
                    xref_func_name = utils.get_func_name(tree_item.to_specific_type.x.obj_ea).lower()
                    if not xref_func_name:
                        xref_func_name = idc.get_name(tree_item.to_specific_type.x.obj_ea).lower()
                    if xref_func_name == self.scanned_function.lower():
                        parent = decompiled_function.body.find_parent_of(tree_item)
                        if parent.op == ida_hexrays.cot_cast: # return value is casted
                            parent = decompiled_function.body.find_parent_of(parent)
                        if (parent.op >= ida_hexrays.cot_eq and parent.op <= ida_hexrays.cot_ult):
                            if check_val is not None:
                                parent = parent.to_specific_type
                                if parent.y and (parent.y.n or parent.y.fpc): # Check if there is a Y operand and if it is a number, if it is, get its value
                                    if parent.y.n: # int
                                        value = parent.y.n._value
                                        negative_value = -((value ^ 0xffffffffffffffff) + 1)
                                    else: # float
                                        value = parent.y.fpc.fnum.float
                                        negative_value = value
                                    if negative_value == check_val or value == check_val:
                                        return True
                                    else:
                                        return False
                            else:
                                return True
                        elif parent.op == ida_hexrays.cit_if or parent.op == ida_hexrays.cot_lnot or parent.op == ida_hexrays.cot_lor or parent.op == ida_hexrays.cot_land:
                            if check_val is not None:
                                if not parent.to_specific_type.cif.expr.y:
                                    # There is no Y, likely checked against 0: if(func_call())
                                    if check_val == 0:
                                        return True
                                    else:
                                        return False
                            else:
                                return True
                        elif parent.op == ida_hexrays.cot_asg:
                            # return value is assigned to the variable/global
                            # Look through the rest of the function and find any if comparison with this variable and const number
                            if parent.to_specific_type.x.v or parent.to_specific_type.x.op == ida_hexrays.cot_obj:
                                for sub_tree_item in list(decompiled_function.treeitems)[index:]:
                                    if sub_tree_item.to_specific_type.op >= 22 and sub_tree_item.to_specific_type.op <= 31:
                                        # Comparison operator
                                        if (sub_tree_item.to_specific_type.x.v and sub_tree_item.to_specific_type.x.v.idx == parent.to_specific_type.x.v.idx) or (sub_tree_item.to_specific_type.x.obj_ea and sub_tree_item.to_specific_type.x.obj_ea == parent.to_specific_type.x.obj_ea):
                                            if check_val is not None:
                                                numeric_val = sub_tree_item.to_specific_type.y
                                                if numeric_val and (numeric_val.n or numeric_val.fpc): # Check if there is a Y operand and if it is a number, if it is, get its value
                                                    if numeric_val.n: # int
                                                        value = numeric_val.n._value
                                                        negative_value = -((value ^ 0xffffffffffffffff) + 1)
                                                    else: # float
                                                        value = numeric_val.fpc.fnum.float
                                                        negative_value = value
                                                    if negative_value == check_val or value == check_val:
                                                        return True
                                                    else:
                                                        return False
                                            else:
                                                return True
                                        else:
                                            # Look for embedded assigns in cit_if (look until cit_if or cit_block is found)
                                            embedded_parent = decompiled_function.body.find_parent_of(sub_tree_item)
                                            while True:
                                                if embedded_parent.op == ida_hexrays.cit_if:
                                                    return True
                                                elif embedded_parent.op == ida_hexrays.cit_block:
                                                    return False
                                                else:
                                                    embedded_parent = decompiled_function.body.find_parent_of(embedded_parent)
                                    elif sub_tree_item.op == ida_hexrays.cit_if or sub_tree_item.op == ida_hexrays.cot_lnot or sub_tree_item.op == ida_hexrays.cot_lor or sub_tree_item.op == ida_hexrays.cot_land:
                                        if check_val is not None:
                                            if not parent.is_expr():
                                                if not parent.to_specific_type.cif.expr.y:
                                                    # There is no Y, likely checked against 0: if(func_call())
                                                    if check_val == 0:
                                                        return True
                                                    else:
                                                        return False
                                        else:
                                            return True


                index += 1
            return False

class vulfi_form_t(ida_kernwin.Form):

    def __init__(self,function_name):
        F = ida_kernwin.Form
        F.__init__(
            self,
            r"""STARTITEM {id:rule_name_str}
BUTTON YES* Run
BUTTON CANCEL Cancel
Custom VulFi rule

{FormChangeCb}
Add custom rule to trace function: {function_name}
Custom rule name:
<#Name of the rule#:{rule_name_str}>
Custom Rule:
<#Rule as desribed in README#:{rule_str}>

""", {
            'function_name': F.StringLabel(function_name),
            'rule_name_str': F.StringInput(),
            'rule_str': F.StringInput(),
            'FormChangeCb': F.FormChangeCb(self.OnFormChange)
        })

    def OnFormChange(self, fid):
        return 1

class VulFi_Single_Function(idaapi.action_handler_t):
    result_window_title = "VulFi Results"
    result_window_columns_names = ["IssueName","FunctionName","FoundIn", "Address","Status", "Priority","Comment"]
    result_window_columns_sizes = [15,20,20,8,8,5,30]
    result_window_columns = [ list(column) for column in zip(result_window_columns_names,result_window_columns_sizes)]
    result_window_row = collections.namedtuple("VulFiResultRow",result_window_columns_names)

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

    # Called when the button is clicked
    def activate(self, ctx):
        custom_rule = custom_rule_name = ""
        if not idc.get_type(self.function_ea) and (idaapi.get_func(self.function_ea) and not idc.get_type(idaapi.get_func(self.function_ea).start_ea)):
            # If the type of the function is not set notify the user
            answer = ida_kernwin.ask_buttons("Yes","No","Cancel",1,f"You should first set type for the function. Continue without type anyway?")
            if answer != 1:
                return

        # Show the form
        function_name = idc.get_func_name(self.function_ea)
        if not function_name:
            function_name = idc.get_name(self.function_ea)
        f = vulfi_form_t(function_name)
        # Compile (in order to populate the controls)
        f.Compile()
        # Execute the form
        ok = f.Execute()
        # If the form was confirmed
        if ok == 1:
            custom_rule_name = f.rule_name_str.value
            custom_rule = f.rule_str.value
        else:
            # Cancel
            return
        # Dispose the form
        f.Free()

        if not custom_rule or not custom_rule_name:
            ida_kernwin.warning("Both rule name and the rule have to be filled!")
            return
        # Craft a temporary rule here:
        tmp_rule = [{"name":f"{custom_rule_name}","function_names":[function_name.lower(),f".{function_name.lower()}"],"wrappers":False,"mark_if":{"High":custom_rule,"Medium":"False","Low":"False"}}]
        vulfi_scanner = VulFiScanner(tmp_rule)
        rows = []
        marked_addrs = []
        vulfi_data = {}
        # Load stored data
        node = idaapi.netnode()
        node.create("vulfi_data")
        if node.getblob(1,"S"):
            vulfi_data = json.loads(node.getblob(1,"S"))
        else:
            vulfi_data = {}
        for item in vulfi_data:
            rows.append([vulfi_data[item]["name"],vulfi_data[item]["function"],vulfi_data[item]["in"],vulfi_data[item]["addr"],vulfi_data[item]["status"],vulfi_data[item]["priority"],vulfi_data[item]["comment"]])
            marked_addrs.append(vulfi_data[item]["addr"])

        # Run the scan for selected function
        print("[VulFi] Started the scan ...")
        scan_result = vulfi_scanner.start_scan(marked_addrs)
        if scan_result is None:
            return
        rows.extend(scan_result)
        print("[VulFi] Scan done!")
        # Save the results
        for item in rows:
            vulfi_data[f"{item[3]}_{item[0]}"] = {"name":item[0],"function":item[1],"in":item[2],"addr":item[3],"status":item[4],"priority":item[5],"comment":item[6]}
        node.setblob(json.dumps(vulfi_data).encode("ascii"),1,"S")
        # Construct and show the form
        results_window = VulFiEmbeddedChooser(self.result_window_title,self.result_window_columns,rows,icon_id)
        results_window.AddCommand("Mark as False Positive", flags=4, menu_index=-1, icon=icon_id, emb=None, shortcut=None)
        results_window.AddCommand("Mark as Suspicious", flags=4, menu_index=-1, icon=icon_id, emb=None, shortcut=None)
        results_window.AddCommand("Mark as Vulnerable", flags=4, menu_index=-1, icon=icon_id, emb=None, shortcut=None)
        results_window.AddCommand("Set Vulfi Comment", flags=4, menu_index=-1, icon=icon_id, emb=None, shortcut=None)
        results_window.AddCommand("Remove Item(s)", flags=4, menu_index=-1, icon=icon_id, emb=None, shortcut=None)
        results_window.AddCommand("Export Results", flags=4, menu_index=-1, icon=icon_id, emb=None, shortcut=None)
        results_window.Show()
        hooks.set_chooser(results_window)

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

class vulfi_main_form_t(ida_kernwin.Form):

    def __init__(self):
        F = ida_kernwin.Form
        F.__init__(
            self,
            r"""STARTITEM {id:rDefault}
BUTTON YES* Run
BUTTON CANCEL Cancel
Custom VulFi rule

{FormChangeCb}
<##What rule set to use?##Default rules:{rDefault}>
<Custom rules:{rCustom}>
<Import previous results (JSON):{rImport}>{cType}>
<#Select a file to open#Browse to open:{iFileOpen}>

""", {
            'iFileOpen': F.FileInput(open=True),
            'cType': F.RadGroupControl(("rDefault", "rCustom","rImport")),
            'FormChangeCb': F.FormChangeCb(self.OnFormChange)
        })

    def OnFormChange(self, fid):
        if fid == -1 or fid == self.cType.id:
            if self.GetControlValue(self.cType) == 0:
                self.EnableField(self.iFileOpen, False)
            else:
                self.EnableField(self.iFileOpen, True)
        return 1

class VulFi(idaapi.action_handler_t):
    result_window_title = "VulFi Results"
    result_window_columns_names = ["IssueName","FunctionName","FoundIn", "Address","Status", "Priority","Comment"]
    result_window_columns_sizes = [15,20,20,8,8,5,30]
    result_window_columns = [ list(column) for column in zip(result_window_columns_names,result_window_columns_sizes)]
    result_window_row = collections.namedtuple("VulFiResultRow",result_window_columns_names)
    def __init__(self):
        idaapi.action_handler_t.__init__(self)

    # Called when the button is clicked
    def activate(self, ctx):
        answer = 0
        skip_scan = False
        rows = []
        vulfi_data = {}
        marked_addrs = []
        # Load stored data
        node = idaapi.netnode()
        node.create("vulfi_data")
        if node.getblob(1,"S"):
            vulfi_data = json.loads(node.getblob(1,"S"))
        else:
            vulfi_data = {}
        if vulfi_data:
            answer = ida_kernwin.ask_buttons("Load Existing","Scan Again","Cancel",1,f"Previous scan results found.")
        else:
            answer = 2
        if vulfi_data and answer == 1:
            for item in vulfi_data:
                rows.append([vulfi_data[item]["name"],vulfi_data[item]["function"],vulfi_data[item]["in"],vulfi_data[item]["addr"],vulfi_data[item]["status"],vulfi_data[item]["priority"],vulfi_data[item]["comment"]])
                marked_addrs.append(f'{vulfi_data[item]["name"]}_{vulfi_data[item]["addr"]}')
            print("[VulFi] Loading previous data.")
        elif answer == -1:
            # Cancel
            return
        else:
            # Show the form
            f = vulfi_main_form_t()
            # Compile (in order to populate the controls)
            f.Compile()
            # Execute the form
            ok = f.Execute()
            # If the form was confirmed
            if ok == 1:
                if f.cType.value == 0:
                    # Default scan
                    vulfi_scanner = VulFiScanner()
                elif f.cType.value == 1:
                    try:
                        with open(os.path.join(f.iFileOpen.value),"r") as rules_file:
                            vulfi_scanner = VulFiScanner(json.load(rules_file))
                    except:
                        ida_kernwin.warning("Failed to load custom rules!")
                        return
                else:
                    try:
                        with open(os.path.join(f.iFileOpen.value),"r") as import_file:
                            import_data = json.load(import_file)
                        for item in import_data["issues"]:
                            rows.append([item["IssueName"],item["FunctionName"],item["FoundIn"],item["Address"],item["Status"],item["Priority"],item["Comment"]])
                        skip_scan = True
                    except:
                        ida_kernwin.warning("Failed to load custom rules!")
                        return
            else:
                return

            for item in vulfi_data:
                rows.append([vulfi_data[item]["name"],vulfi_data[item]["function"],vulfi_data[item]["in"],vulfi_data[item]["addr"],vulfi_data[item]["status"],vulfi_data[item]["priority"],vulfi_data[item]["comment"]])
                marked_addrs.append(f'{vulfi_data[item]["name"]}_{vulfi_data[item]["addr"]}')
            # Run the scan
            if not skip_scan:
                print("[VulFi] Started the scan ...")
                scan_result = vulfi_scanner.start_scan(marked_addrs)
                if scan_result is None:
                    return
                rows.extend(scan_result)
                print("[VulFi] Scan done!")
            # Save the results
            for item in rows:
                vulfi_data[f"{item[3]}_{item[0]}"] = {"name":item[0],"function":item[1],"in":item[2],"addr":item[3],"status":item[4],"priority":item[5],"comment":item[6]}
            node.setblob(json.dumps(vulfi_data).encode("ascii"),1,"S")

        # Construct and show the form
        results_window = VulFiEmbeddedChooser(self.result_window_title,self.result_window_columns,rows,icon_id)
        results_window.AddCommand("Mark as False Positive", flags=4, menu_index=-1, icon=icon_id, emb=None, shortcut=None)
        results_window.AddCommand("Mark as Suspicious", flags=4, menu_index=-1, icon=icon_id, emb=None, shortcut=None)
        results_window.AddCommand("Mark as Vulnerable", flags=4, menu_index=-1, icon=icon_id, emb=None, shortcut=None)
        results_window.AddCommand("Set Vulfi Comment", flags=4, menu_index=-1, icon=icon_id, emb=None, shortcut=None)
        results_window.AddCommand("Remove Item(s)", flags=4, menu_index=-1, icon=icon_id, emb=None, shortcut=None)
        results_window.AddCommand("Export Results", flags=4, menu_index=-1, icon=icon_id, emb=None, shortcut=None)
        results_window.Show()
        hooks.set_chooser(results_window)


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

class vulfi_export_form_t(ida_kernwin.Form):

    def __init__(self):
        F = ida_kernwin.Form
        F.__init__(
            self,
            r"""STARTITEM {id:rJSON}
BUTTON YES* Save
BUTTON CANCEL Cancel
VulFi Results Export

{FormChangeCb}
<##Choose format for export##JSON:{rJSON}>
<CSV:{rCSV}>{cType}>
<#Select the output file#Select the output file:{iFileOpen}>

""", {
            'iFileOpen': F.FileInput(save=True),
            'cType': F.RadGroupControl(("rJSON", "rCSV")),
            'FormChangeCb': F.FormChangeCb(self.OnFormChange)
        })

    def OnFormChange(self,fid):
        return 1


class VulFiEmbeddedChooser(ida_kernwin.Choose):
    def __init__(self,title,columns,items,icon,embedded=False):
        ida_kernwin.Choose.__init__(self,title,columns,embedded=embedded,width=100,flags=ida_kernwin.Choose.CH_MULTI + ida_kernwin.Choose.CH_CAN_REFRESH)
        self.items = items
        self.icon = icon
        self.delete = False
        self.comment = False
        self.export = False

    def GetItems(self):
        return self.items

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

    def OnRefresh(self,n):
        for item in self.items:
            item[2] = utils.get_func_name(int(item[3],16))
        if self.delete:
            for i in reversed(n):
                self.items.pop(i)
            self.delete = False
            self.save()
        if self.comment:
            if len(n) == 1:
                comment = ida_kernwin.ask_str(self.items[n[0]][6],1,f"Enter the comment: ")
            else:
                comment = ida_kernwin.ask_str("",1,f"Enter the comment: ")
            for i in n:
                self.items[i][6] = comment
            self.comment = False
            self.save()
        if self.export:
            self.export = False
            self.vulfi_export()
            
        return n

    def save(self):
        # On close dumps the results
        vulfi_dict = {}
        for item in self.items:
            vulfi_dict[f"{item[3]}_{item[0]}"] = {"name":item[0],"function":item[1],"in":item[2],"addr":item[3],"status":item[4],"priority":item[5],"comment":item[6]}
        node = idaapi.netnode()
        node.create("vulfi_data")
        # Set the blob
        node.setblob(json.dumps(vulfi_dict).encode("ascii"),1,"S")

    def OnCommand(self,number,cmd_id):
        # Cmd_ids: 0 - FP, 1 - Susp, 2 - Vuln
        if cmd_id < 3:
            if cmd_id == 0:
                status = "False Positive"
            if cmd_id == 1:
                status = "Suspicious"
            if cmd_id == 2:
                status = "Vulnerable"
            # Item at index #3 is status
            self.items[number][4] = status
        if cmd_id == 3:
            # Comment
            self.comment = True
        if cmd_id == 4:
            # Delete selected items
            self.delete = True
        if cmd_id == 5:
            # Export
            self.export = True
            
        self.Refresh()
        # Save the data after every change
        self.save()

    def vulfi_export(self):
        # Show the form
        f = vulfi_export_form_t()
        # Compile (in order to populate the controls)
        f.Compile()
        # Execute the form
        ok = f.Execute()
        # If the form was confirmed
        if ok == 1:
            # Get file name
            file_name = f.iFileOpen.value
            if file_name:
                if f.cType.value == 0:
                    # JSON
                    # Pretify 
                    tmp_json = {"issues":[]}
                    for item in self.items:
                        tmp_json["issues"].append({
                            "IssueName": item[0],
                            "FunctionName": item[1],
                            "FoundIn": item[2],
                            "Address": item[3],
                            "Status": item[4],
                            "Priority": item[5],
                            "Comment": item[6]
                        })
                    with open(file_name,"w") as out_file:
                        json.dump(tmp_json, out_file)
                    ida_kernwin.info(f"Results exported in JSON format to {file_name}")
                else:
                    #CSV
                    csv_string = "IssueName,FunctionName,FoundIn,Address,Status,Priority,Comment\n"
                    for item in self.items:
                        csv_string += f"{item[0]},{item[1]},{item[2]},{item[3]},{item[4]},{item[5]},{item[6]}\n"
                    with open(file_name,"w") as out_file:
                        out_file.write(csv_string)
                    ida_kernwin.info(f"Results exported in comma-separated CSV file to {file_name}")
        

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

    def OnSelectLine(self,number):
        # By default change to first selected line
        row = VulFi.result_window_row(*self.items[number[0]])
        destination = row.Address
        ida_kernwin.jumpto(int(destination,16))

    def OnGetLineAttr(self, number):
        if self.items[number][4] == "False Positive":
            return (0x9D9D9D,0)
        elif self.items[number][4] == "Suspicious":
            return (0xd0FF,0)
        elif self.items[number][4] == "Vulnerable":
            return (0xFF,0)

    def OnGetLine(self,number):
        try:
            return self.items[number]
        except:
            self.Refresh()
            return None


class vulfi_fetch_t(idaapi.plugin_t):
    comment = "Vulnerability Finder"
    help = "This script helps to reduce the amount of work required when inspecting potentially dangerous calls to functions such as 'memcpy', 'strcpy', etc."
    wanted_name = "VulFi"
    wanted_hotkey = ""
    flags = idaapi.PLUGIN_KEEP


    def init(self):
        vulfi_desc = idaapi.action_desc_t(
            'vulfi:fetch',   # The action name. This acts like an ID and must be unique
            'VulFi',  # The action text.
            VulFi(),   # The action handler.
            '',      # Optional: the action shortcut
            'Make VulFi fetch the potentially interesting places in binary.',  # Optional: the action tooltip (available in menus/toolbar)
            icon_id)           # Optional: the action icon (shows when in menus/toolbars)
        idaapi.register_action(vulfi_desc)
        idaapi.attach_action_to_menu("Search", "vulfi:fetch", idaapi.SETMENU_APP)

    def run(self):
        pass

    def term(self):
        pass



class Hooks(idaapi.UI_Hooks):
    def __init__(self):
        self.chooser = None
        idaapi.UI_Hooks.__init__(self)


    def finish_populating_widget_popup(self, form, popup):
        action_text = f"Add '{utils.get_func_name(idc.here())}' function to VulFi"
        function_ea = idc.here()
        try:
            # Get selected symbol
            selected_symbol, _ = ida_kernwin.get_highlight(ida_kernwin.get_current_viewer())
            # Check if it is a function name
            for function in idautils.Functions():
                if utils.get_func_name(function) in utils.prep_func_name(selected_symbol):
                    action_text = f"Add '{utils.get_func_name(function)}' function to VulFi"
                    function_ea = function
        except:
            pass
        action_desc = idaapi.action_desc_t(
        'vulfi:get_one',   # The action name. This acts like an ID and must be unique
        action_text,  # The action text.
        VulFi_Single_Function(function_ea),   # The action handler.
        '',      # Optional: the action shortcut
        'Make VulFi look for all interesting refences of this function.',  # Optional: the action tooltip (available in menus/toolbar)
        icon_id)           # Optional: the action icon (shows when in menus/toolbars)
        idaapi.unregister_action("vulfi:get_one")
        idaapi.register_action(action_desc)
        if ida_kernwin.get_widget_type(form) == idaapi.BWN_DISASM or ida_kernwin.get_widget_type(form) == idaapi.BWN_PSEUDOCODE:
            idaapi.attach_action_to_popup(form, popup, "vulfi:get_one", "")

    def current_widget_changed(self, widget, prev_widget):
        title = ida_kernwin.get_widget_title(widget)
        if title and "VulFi" in title and self.chooser:
            self.chooser.Refresh()

    def set_chooser(self,chooser):
        self.chooser = chooser

# Run the hooks
hooks = Hooks()
hooks.hook()

def PLUGIN_ENTRY():
    return vulfi_fetch_t()


repository: Still34/IDAscope7
file: IDAscope.py
description: Fork of simpliFiRE.IDAscope; updated to be compatible with IDA 7.x
url: https://github.com/Still34/IDAscope7/blob/65a5a09c207cdb5d1c26a102da9979a9959c2ab9/IDAscope.py
wanted_name: IDAscope
created_at: 2020-10-28 14:57:25+00:00
pushed_at: 2021-01-27 12:52:11+00:00
content:
#!/usr/bin/python
########################################################################
# Copyright (c) 2012
# Daniel Plohmann <daniel.plohmann<at>gmail<dot>com>
# Alexander Hanel <alexander.hanel<at>gmail<dot>com>
# All rights reserved.
########################################################################
#
#  This file is part of IDAscope
#
#  IDAscope is free software: you can redistribute it and/or modify it
#  under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful, but
#  WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
#  General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program.  If not, see
#  <http://www.gnu.org/licenses/>.
#
########################################################################

import idascope.core.helpers.QtShim as QtShim
QtGui = QtShim.get_QtGui()
QtCore = QtShim.get_QtCore()
QtWidgets = QtShim.get_QtWidgets()

import idc
import idaapi
from idaapi import plugin_t
from ida_kernwin import PluginForm

import idascope.config as config
from idascope.core.helpers.ClassCollection import ClassCollection
from idascope.core.structures.IDAscopeConfiguration import IDAscopeConfiguration
from idascope.core.SemanticIdentifier import SemanticIdentifier
from idascope.core.DocumentationHelper import DocumentationHelper
from idascope.core.WinApiProvider import WinApiProvider
from idascope.core.CryptoIdentifier import CryptoIdentifier
from idascope.core.YaraScanner import YaraScanner
from idascope.core.IdaProxy import IdaProxy
from idascope.core.SemanticExplorer import SemanticExplorer
from idascope.widgets.FunctionInspectionWidget import FunctionInspectionWidget
from idascope.widgets.WinApiWidget import WinApiWidget
from idascope.widgets.CryptoIdentificationWidget import CryptoIdentificationWidget
from idascope.widgets.YaraScannerWidget import YaraScannerWidget
from idascope.widgets.SemanticExplorerWidget import SemanticExplorerWidget

################################################################################
# Core of the IDAscope GUI.
################################################################################

HOTKEYS = None
IDASCOPE = None
NAME = "simpliFiRE.IDAscope v1.2.1"


class IDAscopeForm(PluginForm):
    """
    This class contains the main window of IDAscope
    Setup of core modules and widgets is performed in here.
    """

    def __init__(self):
        super(IDAscopeForm, self).__init__()
        global HOTKEYS
        HOTKEYS = []
        self.idascope_widgets = []
        self.cc = ClassCollection()
        self.ensureRootPathSanity(config.configuration)
        self.config = IDAscopeConfiguration(config.configuration, self.cc)
        self.icon = self.cc.QIcon(self.config.icon_file_path + "idascope.png")

    def ensureRootPathSanity(self, configuration):
        try:
            root_dir = configuration["paths"]["idascope_root_dir"]
            if not self.cc.os.path.exists(root_dir) or not "IDAscope.py" in self.cc.os.listdir(root_dir):
                print "[!] IDAscope.py is not present in root directory specified in \"config.py\", " \
                     + "trying to resolve path..."
                resolved_pathname = self.cc.os.path.dirname(self.cc.sys.argv[0])
                if "IDAscope.py" in self.cc.os.listdir(resolved_pathname):
                    print "[+] IDAscope root directory successfully resolved."
                    configuration["paths"]["idascope_root_dir"] = resolved_pathname
                else:
                    print "[-] IDAscope.py is not resolvable!"
                    raise Exception()
        except:
            print "[!] IDAscope config is broken. Could not locate root directory. " \
                 + "Try setting the field \"idascope_root_dir\" to the path where \"IDAscope.py\" is located."
            self.cc.sys.exit(-1)

    def setupSharedModules(self):
        """
        Setup shared IDAscope modules.
        """
        time_before = self.cc.time.time()
        print ("[/] setting up shared modules...")
        self.semantic_explorer = SemanticExplorer(self)
        self.documentation_helper = DocumentationHelper(self.config)
        self.semantic_identifier = SemanticIdentifier(self.config)
        self.winapi_provider = WinApiProvider(self.config)
        self.crypto_identifier = CryptoIdentifier(self.config)
        self.yara_scanner = YaraScanner(self.config)
        print ("[\\] this took %3.2f seconds.\n" % (self.cc.time.time() - time_before))

    def setupWidgets(self):
        """
        Setup IDAscope widgets.
        """
        time_before = self.cc.time.time()
        print ("[/] setting up widgets...")
        self.idascope_widgets.append(SemanticExplorerWidget(self))
        self.idascope_widgets.append(FunctionInspectionWidget(self))
        self.idascope_widgets.append(WinApiWidget(self))
        self.idascope_widgets.append(CryptoIdentificationWidget(self))
        self.idascope_widgets.append(YaraScannerWidget(self))
        self.setupIDAscopeForm()
        print ("[\\] this took %3.2f seconds.\n" % (self.cc.time.time() - time_before))

    def setupIDAscopeForm(self):
        """
        Orchestrate the already initialized widgets in tabs on the main window.
        """
        self.tabs = self.cc.QTabWidget()
        self.tabs.setTabsClosable(False)
        for widget in self.idascope_widgets:
            self.tabs.addTab(widget, widget.icon, widget.name)
        layout = self.cc.QVBoxLayout()
        layout.addWidget(self.tabs)
        self.parent.setLayout(layout)

    def OnCreate(self, form):
        """
        When creating the form, setup the shared modules and widgets
        """
        self.printBanner()
        # compatibility with IDA < 6.9
        try:
            self.parent = self.FormToPySideWidget(form)
        except Exception as exc:
            self.parent = self.FormToPyQtWidget(form)
        self.parent.setWindowIcon(self.icon)
        self.setupSharedModules()
        self.setupWidgets()

    def printBanner(self):
        banner = "#############################################\n" \
               + "  ___ ____    _                             \n" \
               + " |_ _|  _ \  / \   ___  ___ ___  _ __   ___ \n" \
               + "  | || | | |/ _ \ / __|/ __/ _ \\| '_ \\ / _ \\\n" \
               + "  | || |_| / ___ \\\\__ \\ (_| (_) | |_) |  __/\n" \
               + " |___|____/_/   \\_\\___/\\___\\___/| .__/ \\___|\n" \
               + "                                |_|         \n" \
               + "#############################################\n" \
               + " by Daniel Plohmann and Alexander Hanel      \n" \
               + "#############################################\n"
        print banner
        print ("[+] Loading simpliFiRE.IDAscope")

    def OnClose(self, form):
        """
        Perform cleanup.
        """
        global IDASCOPE
        del IDASCOPE

    def Show(self):
        if idc.retrieve_input_file_md5() is None:
            return
        else:
            return PluginForm.Show(self,
                NAME,
                options=(PluginForm.WCLS_CLOSE_LATER | PluginForm.WOPN_RESTORE | PluginForm.WCLS_SAVE))

################################################################################
# functionality offered to IDAscope's widgets
################################################################################

    def setTabFocus(self, widget_name):
        """
        Can be used by IDAscope widgets to set focus to a widget, identified by name.
        @param widget_name: A widget name
        @type widget_name: str
        """
        for widget in self.idascope_widgets:
            if widget.name == widget_name:
                tab_index = self.tabs.indexOf(widget)
                self.tabs.setCurrentIndex(tab_index)
        return

    def registerHotkey(self, shortcut, py_function_pointer):
        """
        Can be used by IDAscope widgets to register hotkeys.
        Uses a global list HOTKEYS of function pointers that link to the desired functionality.
        Right now, linked functions cannot take parameters and should scrape all information they need by themselves.
        @param shortcut: A string describing a shortcut, e.g. "ctrl+F3"
        @type shortcut: str
        @param py_function_pointer: a python function that shall be called when the shortcut is triggered.
        @type py_function_pointer: a pointer to a python function
        """
        global HOTKEYS
        hotkey_index = len(HOTKEYS)
        hotkey_name = "idascope_HOTKEY_%d" % hotkey_index
        HOTKEYS.append(py_function_pointer)
        self.cc.ida_proxy.compile_idc_text('static %s() { RunPythonStatement("HOTKEYS[%d]()"); }' % (hotkey_name, hotkey_index))
        self.cc.ida_proxy.AddHotkey(shortcut, hotkey_name)

################################################################################
# Usage as plugin
################################################################################


def PLUGIN_ENTRY():
    return IDAscopePlugin()


class IDAscopePlugin(plugin_t):
    """
    Plugin version of IDAscope. Use this to deploy IDAscope via IDA plugins folder.
    """
    flags = idaapi.PLUGIN_UNL
    comment = NAME
    help = "IDAscope - Different tools to ease reverse engineering."
    wanted_name = "IDAscope"
    wanted_hotkey = "Ctrl-F4"

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

    def run(self, arg=0):
        # Create form
        f = IDAscopeForm()
        # Show the form
        f.Show()
        return

    def term(self):
        pass

################################################################################
# Usage as script
################################################################################


def main():
    global IDASCOPE
    try:
        IDASCOPE
        IDASCOPE.OnClose(IDASCOPE)
        print ("reloading IDAscope")
        IDASCOPE = IDAscopeForm()
        return
    except Exception:
        IDASCOPE = IDAscopeForm()

    if IDASCOPE.config.idascope_plugin_only:
        print "IDAscope: configured as plugin-only mode, ignoring main function of script. " \
             + "This can be changed in \"idascope/config.py\"."
    else:
        IDASCOPE.Show()


if __name__ == "__main__":
    main()


repository: TheBinKing/VulnLink
file: VulnLink/VulnLink.py
wanted_name: VulnLink
comment: Vulnerability Finder
created_at: 2023-11-19 05:03:21+00:00
pushed_at: 2024-08-24 07:01:31+00:00
url:
https://github.com/TheBinKing/VulnLink/blob/1dcc32e835ad36b775a8c1300a48ec8891728aab/VulnLink/VulnLink.py

content:
from cmath import exp
import collections
from email.policy import default
from uuid import RESERVED_FUTURE
import idaapi
import idc
import os
import json
import idautils
import ida_kernwin

from TargetFuncFinder import analyze_dangerous_functions_in_binary
from TargetFuncChain import get_function_call_chains
from util import *



icon = b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00*\x00\x00\x00&\x08\x06\x00\x00\x00\xb2\x01\t \x00\x00\x00\x01sRGB\x00\xae\xce\x1c\xe9\x00\x00\x00\x04gAMA\x00\x00\xb1\x8f\x0b\xfca\x05\x00\x00\x00\tpHYs\x00\x00\x0e\xc3\x00\x00\x0e\xc3\x01\xc7o\xa8d\x00\x00\x02TIDATXG\xcd\x98\xcd.\x03Q\x14\xc7\xef\x10\x14-\xd5\xf8Jj!x\x00+,\xc5\x03\x88\x07\xb0\xb0\xb2\xea\x82\'\xc0\x13\x90\xb0\x15\x8f\x80x\x00+k\x89\xd8\x93\xb0\x94\xb0\x13A\xc6\xf9\xb7skz\xe7\xcc\xf4\xdc\xf9\xec/\xf9\xa5\x9d\x99\xce\xbdg\xce\xdc9\xf7N\x15\xb1@\x1e\x93\xf3\xd8\xe81\xaa\xe4\x91\xf7\xd9\xe4\x9etI\x04\xdc\x0b \xb0=\xf2\x89\xbc\xc0\x0e\r\xb2\x89@\xe1;\xb9C\x16\x05\xfaF\x80:\x9ee\xb2\x83KR\x1f\x84\xc8\xf2:\x99\x17\xe8K\xdfY\xed\x01\x19\xc0\x9fU\xbf\xb8\x80\xc0U\xa5\x08\xda\xbe%\xcd~qg\xdbc\xd3\x04\xe3\xc1<A\x9b\xf6\xf8E\x80h\x13\x01q\xfd\xb1\xd9\xd4\xe0\n\xb8\x93\xfcF6 \x00}D\x05\x081FC\xb3\xa9A#\xdc\xc9~1\x96l\x1f8I\x80Zq\xdb\xfe\xa7.J\x04\xbcEF\x81\x00\xf1\x1bI\x80\x10\xe3U\x0c\x1a\xe6\x1a\t\x13\x0f\x1c7apOr7\xad+\x8d4\xab~\xf10"`tf\x96;\x898\xc7\x1at\xc65\x96\x95\x18\x1a\xb1\xa7q\xae\xbeee\xa2\xf2\x97WV\x91\xcd\xae\xe5\xa8\x1b\x81\xb1\xd6glK\x1dp\x1c\xb7\x9fd\x8ea\x01\x92\x98\xc0\xd4\xeaxn\x0e\x97;\xf6G\xb9:Tj\x9e\xc3\x1c\x13\x15w)\x81\xa9\x15\x9d\xdeL\xd5\xdd\xdd\xf2x\xc7~\xceF\xa5\xea\x9e\xd5f\xc2\x02Mu\xa5\x86+\x0etrM\x81\xbe\xcd-\xb9\x1b\xa5\x91\xc01\xed\xf6\xe8\x98\xfbZ_tO\'\xa6\xb9\xe3\xa8\xb1"h\xb8\x89\xf8 OZ_\x83\x9c\xd7f\xd5\xda\xd0\xb0\xb7\xf5\xcf\xca`I\x1d\x8dO\xaa\x92C\xb9\xe4\xc1\xea]\x844P\xb0O>\xb7\xbe\xb6x\xfc\xfeRw_\x9f\xea\x81>\x1b\xe5\xaa\x1aq\xfe\x9bCp\x8d\xcaD\xfb7/\xbf?\xde\x916W\x9e\x99\x80\xf1\xc4\xdd\xc28ZO\x95\xb6\xc4\x99ZMcM\x95\xb6\xd8.XLQ\xdc\xb3|c\xe8 IV\x93.\xbc\xad\x88;\xb5&Zx\xc4%\xce\x82%\xd7lj0\xce\xb8`\xc2Lu\xaa\xb4%\xea\xad\xd5\xb4\x90lj\xd8\xa9\x95\x11Sea\xd9\xd4\x1c\x92\\p~3/\xee\x12\x90\xa9\xa8r\x95Kq\x97\x82\xf1\xc7\x05\ts+\xee\x12\xc2\xb2Z\xe8\x03\x14\x86\xb9`I\xe5=(+\xfcY\xed\xc9ljtV\x0b-\xeeR\xe2\xfc\x81V\x08\x19dR\xa9?"\x80\x16\n\xa6\x0c\x13@\x00\x00\x00\x00IEND\xaeB`\x82'
icon_id = idaapi.load_custom_icon(data=icon, format="png")

class utils:
    def get_negative_disass(op):
        # Get negative value based on the size of the operand
        if op.dtype == 0x0: # Byte
            return -(((op.value & 0xff) ^ 0xff) + 1)
        elif op.dtype == 0x1: # Word
            return -(((op.value & 0xffff) ^ 0xffff) + 1)
        elif op.dtype == 0x2: # Dword
            return -(((op.value & 0xffffffff) ^ 0xffffffff) + 1)
        elif op.dtype == 0x7: # Qword
            return -((op.value ^ 0xffffffffffffffff) + 1)

    def get_func_name(ea):
        # Get pretty function name
        func_name = utils.get_pretty_func_name(idc.get_func_name(ea))
        if not func_name:
            func_name = utils.get_pretty_func_name(idc.get_name(ea))
        return func_name

    def get_pretty_func_name(name):
        # Demangle function name
        demangled_name = idc.demangle_name(name, idc.get_inf_attr(idc.INF_SHORT_DN))
        # Return as is if failed to demangle
        if not demangled_name:
            return name
        # Cut arguments
        return demangled_name[:demangled_name.find("(")]

    def prep_func_name(name):
        if name[0] != "." and name[0] != "_":
            # Name does not start with dot or underscore
            return [name,f".{name}",f"_{name}"]
        else:
            return [name[1:],f".{name[1:]}",f"_{name[1:]}"]
        
# implement the ui of VulnLikn Call chains Results       
class VulnLink_Single_Function(idaapi.action_handler_t):
    result_window_title = "VulnLink Call Chain Results"
    result_window_columns_names = ["FuncName","FuncAddr", "ChainNo"]
    result_window_columns_sizes = [15,20,8]
    result_window_columns = [ list(column) for column in zip(result_window_columns_names,result_window_columns_sizes)]
    result_window_row = collections.namedtuple("VulnLinkCallChainResultRow",result_window_columns_names)

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

    # Called when the button is clicked
    def activate(self, ctx):
        # Show the form
        function_name = idc.get_func_name(self.function_ea)
        if not function_name:
            function_name = idc.get_name(self.function_ea)
        
        explorer = HexRaysCodeXplorer()
        chains = get_function_call_chains(function_name)
        rows = []
        f = open("output.txt", "w")
        chain_score = {}
        all_rows = []
        # 两步交互获取所有调用链的危险系数和危险描述
        for i, chain in enumerate(chains):
            f.write("chain {} start\n".format(i))
            call_chains = []
            func_infos = []
            rows = []
            for func in chain:
                rows.append([func[0], str(func[1]), i])
                f.write("{}, {}\n".format(func[0], func[1]))
                assembly_code = explorer.get_function_asm_code(func[1])
                f.write("assemply_code: {}\n".format(assembly_code))
                pseudo_code = explorer.get_function_pseudocode(func[1])
                f.write("pseudo_code: {}\n".format(pseudo_code))
                # 根据汇编码和伪代码获取函数基本信息
                func_info = get_function_signature(func[0], assembly_code, pseudo_code)
                f.write("func_info: {}\n".format(func_info))
                call_chains.append(func[1])
                func_infos.append(func_info)
            
            f.write("call_chains: {}\n".format(call_chains))
            f.write("func_infos: {}\n".format(func_infos))
            # 根据调用链上各个函数的基本信息获取调用链的危险系数和危险描述
            score, desc = get_call_chain_danger_score(call_chains, func_infos)
            f.write("score={}, desc={}\n".format(score, desc))
            rows.append(["score={}".format(score), '', ''])
            rows.append(["desc={}".format(desc), '', ''])   
            all_rows.append(rows)
            chain_score[i] = score
            
        # 按危险系数从高到低排序
        all_rows = sorted(all_rows, key=lambda rows: -chain_score[rows[0][2]])
        # Construct and show the form
        items = []
        for i, rows in enumerate(all_rows):
            for row in rows:
                new_row = []
                for each in row:
                    new_row.append(str(each))
                items.append(new_row)
            items.append(['', '', ''])
        results_window = VulnLinkCallChainsEmbeddedChooser(self.result_window_title,self.result_window_columns,items,icon_id, all_rows)
        results_window.Show()
        hooks.set_chains_chooser(results_window)

    # This action is always available.
    def update(self, ctx):
        return idaapi.AST_ENABLE_ALWAYS
    
    
# implement form of single call chain results
class VulnLinkCallChainsEmbeddedChooser(ida_kernwin.Choose):
    titles = ["FuncName","FuncAddr", "Status", "Priority","Comment"]
    result_window_row = collections.namedtuple("VulnLinkCallChainResultRow",titles)
    def __init__(self,title,columns,items,icon, chains, embedded=False):
        ida_kernwin.Choose.__init__(self,title,columns,embedded=embedded,width=100,flags=ida_kernwin.Choose.CH_MULTI + ida_kernwin.Choose.CH_CAN_REFRESH)
        self.items = items
        self.chains = chains
        self.icon = icon

    def GetItems(self):
        return self.items

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

    def OnRefresh(self, n):
        items = []
        for i, rows in enumerate(self.chains):
            for row in rows:
                new_row = []
                for each in row:
                    new_row.append(str(each))
                items.append(new_row)
            items.append(['', '', ''])
        self.items = items

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

    def OnSelectLine(self,number):
        # By default change to first selected line
        row = VulnLink_Single_Function.result_window_row(*self.items[number[0]])
        chainNo = int(row.ChainNo)
        columns_sizes = [15,20,8, 12, 30]
        columns = [ list(column) for column in zip(self.titles, columns_sizes)]
        rows = []
        for func in self.chains[chainNo]:
            rows.append([func[0], str(func[1]), '', '', ''])
        results_window = VulnLinkCallChainEmbeddedChooser('single call chain', columns, rows,icon_id, self.chains, chainNo)
        results_window.AddCommand("Mark as False Positive", flags=4, menu_index=-1, icon=icon_id, emb=None, shortcut=None)
        results_window.AddCommand("Mark as Suspicious", flags=4, menu_index=-1, icon=icon_id, emb=None, shortcut=None)
        results_window.AddCommand("Mark as Vulnerable", flags=4, menu_index=-1, icon=icon_id, emb=None, shortcut=None)
        results_window.AddCommand("Set Vulnlink Comment", flags=4, menu_index=-1, icon=icon_id, emb=None, shortcut=None)
        results_window.AddCommand("clean this chain", flags=4, menu_index=-1, icon=icon_id, emb=None, shortcut=None)
        results_window.AddCommand("clean all chains contains this node", flags=4, menu_index=-1, icon=icon_id, emb=None, shortcut=None)
        results_window.Show()
        hooks.set_chain_chooser(results_window)
      
        
    def OnGetLine(self,number):
        try:
            return self.items[number]
        except:
            self.Refresh()
            return None
         
# # implement ui of single call chain results 
class VulnLinkCallChainEmbeddedChooser(ida_kernwin.Choose):
    def __init__(self,title,columns,items,icon,chains, chainNo, embedded=False):
        ida_kernwin.Choose.__init__(self,title,columns,embedded=embedded,width=100,flags=ida_kernwin.Choose.CH_MULTI + ida_kernwin.Choose.CH_CAN_REFRESH)
        self.items = items
        self.icon = icon
        self.comment = False
        self.chains = chains
        self.chainNo = chainNo

    def GetItems(self):
        return self.items

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

    def OnRefresh(self,n):
        for item in self.items:
            item[0] = utils.get_func_name(int(item[1],16))
        if self.comment:
            if len(n) == 1:
                comment = ida_kernwin.ask_str(self.items[n[0]][4],1,f"Enter the comment: ")
            else:
                comment = ida_kernwin.ask_str("",1,f"Enter the comment: ")
            for i in n:
                self.items[i][4] = comment
            self.comment = False
            #self.save()
            
        return n

    def OnCommand(self,number,cmd_id):
        # Cmd_ids: 0 - FP, 1 - Susp, 2 - Vuln
        if cmd_id < 3:
            if cmd_id == 0:
                status = "False Positive"
            if cmd_id == 1:
                status = "Suspicious"
            if cmd_id == 2:
                status = "Vulnerable"
            # Item at index #3 is status
            self.items[number][4] = status
        if cmd_id == 3:
            # Comment
            self.comment = True
        if cmd_id == 4:
            # Delete selected items
            del self.chains[self.chainNo]
            self.Close()
        if cmd_id == 5:
            self.del_relate_chains(number)
            self.Close()
        self.Refresh()



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

    def OnSelectLine(self,number):
        # By default change to first selected line
        row = VulnLinkCallChainsEmbeddedChooser.result_window_row(*self.items[number[0]])
        destination = row.FuncAddr
        ida_kernwin.jumpto(int(destination,10))

    def OnGetLineAttr(self, number):
        if self.items[number][2] == "False Positive":
            return (0x9D9D9D,0)
        elif self.items[number][2] == "Suspicious":
            return (0xd0FF,0)
        elif self.items[number][2] == "Vulnerable":
            return (0xFF,0)

    def OnGetLine(self,number):
        try:
            return self.items[number]
        except:
            self.Refresh()
            return None
        
    def del_relate_chains(self, number):
        func_name = self.items[number][0]
        func_addr = int(self.items[number][1])
        deleted = []
        for i, chain in enumerate(self.chains):
            for node in chain:
                if node[0] == func_name and node[1] == func_addr:
                    deleted.append(i)
                    break
        while deleted:
            self.chains.pop(deleted.pop())
               

class VulnLinkScanner:
    def __init__(self, custom_rules_file=None):
        if not custom_rules_file:
            self.rules_file = 'function_content.json'
        else:
            self.rules_file = custom_rules_file

    def start_scan(self):
        ida_kernwin.show_wait_box("VulnLink scan running ... ")
        rules_file_path = os.path.join(os.path.abspath(__file__), "..", self.rules_file)
        target_result, target_result_record = analyze_dangerous_functions_in_binary(rules_file_path)
        ida_kernwin.hide_wait_box()
        res = []
        for each_result in target_result_record:
            print(each_result)
            result = []
            result.append(each_result['Binary file name'])
            result.append(each_result['FuncName'])
            result.append(each_result['Calladd'])
            result.append(each_result['FoundIn'])
            result.extend(['', '', ''])
            res.append(result)
        return target_result, res

  
class vulnlink_form_t(ida_kernwin.Form):

    def __init__(self,function_name):
        F = ida_kernwin.Form
        F.__init__(
            self,
            r"""STARTITEM {id:rule_name_str}
BUTTON YES* Run
BUTTON CANCEL Cancel
Custom VulnLink rule

{FormChangeCb}
Add custom rule to trace function: {function_name}
Custom rule name:
<#Name of the rule#:{rule_name_str}>
Custom Rule:
<#Rule as desribed in README#:{rule_str}>


""", {
            'function_name': F.StringLabel(function_name),
            'rule_name_str': F.StringInput(),
            'rule_str': F.StringInput(),
            'FormChangeCb': F.FormChangeCb(self.OnFormChange)
        })

    def OnFormChange(self, fid):
        return 1


class vulnlink_main_form_t(ida_kernwin.Form):

    def __init__(self):
        F = ida_kernwin.Form
        F.__init__(
            self,
            r"""STARTITEM {id:rDefault}
BUTTON YES* Run
BUTTON CANCEL Cancel
Custom VulnLink rule

{FormChangeCb}
<##What rule set to use?##Default rules:{rDefault}>
<Custom rules:{rCustom}>
<Import previous results (JSON):{rImport}>{cType}>
<#Select a file to open#Browse to open:{iFileOpen}>

""", {
            'iFileOpen': F.FileInput(open=True),
            'cType': F.RadGroupControl(("rDefault", "rCustom","rImport")),
            'FormChangeCb': F.FormChangeCb(self.OnFormChange)
        })

    def OnFormChange(self, fid):
        if fid == -1 or fid == self.cType.id:
            if self.GetControlValue(self.cType) == 0:
                self.EnableField(self.iFileOpen, False)
            else:
                self.EnableField(self.iFileOpen, True)
        return 1

class VulnLink(idaapi.action_handler_t):
    result_window_title = "VulnLink Results"
    result_window_columns_names = ["BinaryFileName","FuncName","Calladd", "FoundIn","Status", "Priority","Comment"]
    result_window_columns_sizes = [15,20,20,8,8,5,30]
    result_window_columns = [ list(column) for column in zip(result_window_columns_names,result_window_columns_sizes)]
    result_window_row = collections.namedtuple("VulnLinkResultRow",result_window_columns_names)
    def __init__(self):
        idaapi.action_handler_t.__init__(self)

    # Called when the button is clicked
    def activate(self, ctx):
        answer = 0
        skip_scan = False
        rows = []
        vulnlink_data = {}
        # Load stored data
        node = idaapi.netnode()
        node.create("vulnlink_data")
        if node.getblob(1,"S"):
            vulnlink_data = json.loads(node.getblob(1,"S"))
        else:
            vulnlink_data = {}
        if vulnlink_data:
            answer = ida_kernwin.ask_buttons("Load Existing","Scan Again","Cancel",1,f"Previous scan results found.")
        else:
            answer = 2
        if vulnlink_data and answer == 1:
            for item in vulnlink_data:
                # mod load previous data
                rows.append([vulnlink_data[item]["BinaryFileName"],vulnlink_data[item]["FuncName"],vulnlink_data[item]["Calladd"],vulnlink_data[item]["FoundIn"],vulnlink_data[item]["status"],vulnlink_data[item]["priority"],vulnlink_data[item]["comment"]])
                
            print("[VulnLink] Loading previous data.")
        elif answer == -1:
            # Cancel
            return
        else:
            # Show the form
            f = vulnlink_main_form_t()
            # Compile (in order to populate the controls)
            f.Compile()
            # Execute the form
            ok = f.Execute()
            # If the form was confirmed
            if ok == 1:
                # choose rules
                if f.cType.value == 0:
                    # Default scan
                    vulnlink_scanner = VulnLinkScanner()
                elif f.cType.value == 1:
                    try:
                        vulnlink_scanner = VulnLinkScanner(f.iFileOpen.value)
                    except:
                        ida_kernwin.warning("Failed to load custom rules!")
                        return
                else:
                    try:
                        with open(os.path.join(f.iFileOpen.value),"r") as import_file:
                            import_data = json.load(import_file)
                        for item in import_data["issues"]:
                            rows.append([item["BinaryFileName"],item["FuncName"],item["Calladd"],item["FoundIn"],item["Status"],item["Priority"],item["Comment"]])
                        skip_scan = True
                    except:
                        ida_kernwin.warning("Failed to load custom data files")
                        return
            else:
                return

            for item in vulnlink_data:
                rows.append([vulnlink_data[item]["BinaryFileName"],vulnlink_data[item]["FuncName"],vulnlink_data[item]["Calladd"],vulnlink_data[item]["FoundIn"],vulnlink_data[item]["status"],vulnlink_data[item]["priority"],vulnlink_data[item]["comment"]])
            # Run the scan
            if not skip_scan:
                print("[VulnLink] Started the scan ...")
                dangerous_funcs, scan_result = vulnlink_scanner.start_scan()
                if scan_result is None:
                    return
                rows.extend(scan_result)
                print("[VulnLink] Scan done!")
            # Save the results
            for item in rows:
                vulnlink_data[f"{item[3]}_{item[0]}"] = {"BinaryFileName":item[0],"FuncName":item[1],"Calladd":item[2],"FoundIn":item[3],"status":item[4],"priority":item[5],"comment":item[6]}
            node.setblob(json.dumps(vulnlink_data).encode("ascii"),1,"S")

        # Construct and show the form
        results_window = VulnLinkEmbeddedChooser(self.result_window_title,self.result_window_columns,rows,icon_id)
        results_window.AddCommand("Mark as False Positive", flags=4, menu_index=-1, icon=icon_id, emb=None, shortcut=None)
        results_window.AddCommand("Mark as Suspicious", flags=4, menu_index=-1, icon=icon_id, emb=None, shortcut=None)
        results_window.AddCommand("Mark as Vulnerable", flags=4, menu_index=-1, icon=icon_id, emb=None, shortcut=None)
        results_window.AddCommand("Set VulnLink Comment", flags=4, menu_index=-1, icon=icon_id, emb=None, shortcut=None)
        results_window.AddCommand("Remove Item(s)", flags=4, menu_index=-1, icon=icon_id, emb=None, shortcut=None)
        results_window.AddCommand("Export Results", flags=4, menu_index=-1, icon=icon_id, emb=None, shortcut=None)
        results_window.Show()
        hooks.set_func_chooser(results_window)

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

class vulnlink_export_form_t(ida_kernwin.Form):

    def __init__(self):
        F = ida_kernwin.Form
        F.__init__(
            self,
            r"""STARTITEM {id:rJSON}
BUTTON YES* Save
BUTTON CANCEL Cancel
VulnLink Results Export

{FormChangeCb}
<##Choose format for export##JSON:{rJSON}>
<CSV:{rCSV}>{cType}>
<#Select the output file#Select the output file:{iFileOpen}>

""", {
            'iFileOpen': F.FileInput(save=True),
            'cType': F.RadGroupControl(("rJSON", "rCSV")),
            'FormChangeCb': F.FormChangeCb(self.OnFormChange)
        })

    def OnFormChange(self,fid):
        return 1


class VulnLinkEmbeddedChooser(ida_kernwin.Choose):
    def __init__(self,title,columns,items,icon,embedded=False):
        ida_kernwin.Choose.__init__(self,title,columns,embedded=embedded,width=100,flags=ida_kernwin.Choose.CH_MULTI + ida_kernwin.Choose.CH_CAN_REFRESH)
        self.items = items
        self.icon = icon
        self.delete = False
        self.comment = False
        self.export = False

    def GetItems(self):
        return self.items

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

    def OnRefresh(self,n):
        for item in self.items:
            item[2] = utils.get_func_name(int(item[2], 16))
        if self.comment:
            if len(n) == 1:
                comment = ida_kernwin.ask_str(self.items[n[0]][6],1,f"Enter the comment: ")
            else:
                comment = ida_kernwin.ask_str("",1,f"Enter the comment: ")
            for i in n:
                self.items[i][6] = comment
            self.comment = False
            self.save()
        if self.export:
            self.export = False
            self.vulnlink_export()
            
        return n

    def save(self):
        # On close dumps the results
        vulnlink_dict = {}
        for item in self.items:
            vulnlink_dict[f"{item[3]}_{item[0]}"] = {"name":item[0],"function":item[1],"in":item[2],"addr":item[3],"status":item[4],"priority":item[5],"comment":item[6]}
        node = idaapi.netnode()
        node.create("vulnlink_data")
        # Set the blob
        node.setblob(json.dumps(vulnlink_dict).encode("ascii"),1,"S")

    def OnCommand(self,number,cmd_id):
        # Cmd_ids: 0 - FP, 1 - Susp, 2 - Vuln
        if cmd_id < 3:
            if cmd_id == 0:
                status = "False Positive"
            if cmd_id == 1:
                status = "Suspicious"
            if cmd_id == 2:
                status = "Vulnerable"
            # Item at index #3 is status
            self.items[number][4] = status
        if cmd_id == 3:
            # Comment
            self.comment = True
        if cmd_id == 4:
            # Delete selected items
            self.delete = True
        if cmd_id == 5:
            # Export
            self.export = True
            
        self.Refresh()
        # Save the data after every change
        self.save()

    def vulnlink_export(self):
        # Show the form
        f = vulnlink_export_form_t()
        # Compile (in order to populate the controls)
        f.Compile()
        # Execute the form
        ok = f.Execute()
        # If the form was confirmed
        if ok == 1:
            # Get file name
            file_name = f.iFileOpen.value
            if file_name:
                if f.cType.value == 0:
                    # JSON
                    # Pretify 
                    tmp_json = {"issues":[]}
                    for item in self.items:
                        tmp_json["issues"].append({
                            "BinaryFileName": item[0],
                            "FuncName": item[1],
                            "Calladd": item[2],
                            "FoundIn": item[3],
                            "Status": item[4],
                            "Priority": item[5],
                            "Comment": item[6]
                        })
                    with open(file_name,"w") as out_file:
                        json.dump(tmp_json, out_file)
                    ida_kernwin.info(f"Results exported in JSON format to {file_name}")
                else:
                    #CSV
                    csv_string = "IssueName,FunctionName,FoundIn,Address,Status,Priority,Comment\n"
                    for item in self.items:
                        csv_string += f"{item[0]},{item[1]},{item[2]},{item[3]},{item[4]},{item[5]},{item[6]}\n"
                    with open(file_name,"w") as out_file:
                        out_file.write(csv_string)
                    ida_kernwin.info(f"Results exported in comma-separated CSV file to {file_name}")
        

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

    def OnSelectLine(self,number):
        # By default change to first selected line
        row = VulnLink.result_window_row(*self.items[number[0]])
        destination = row.Calladd
        ida_kernwin.jumpto(int(destination,16))

    def OnGetLineAttr(self, number):
        if self.items[number][4] == "False Positive":
            return (0x9D9D9D,0)
        elif self.items[number][4] == "Suspicious":
            return (0xd0FF,0)
        elif self.items[number][4] == "Vulnerable":
            return (0xFF,0)

    def OnGetLine(self,number):
        try:
            return self.items[number]
        except:
            self.Refresh()
            return None


class vulnlink_fetch_t(idaapi.plugin_t):
    comment = "Vulnerability Finder"
    help = "This script helps to reduce the amount of work required when inspecting potentially dangerous calls to functions such as 'memcpy', 'strcpy', etc."
    wanted_name = "VulnLink"
    wanted_hotkey = ""
    flags = idaapi.PLUGIN_KEEP

    def init(self):
        vulnlink_desc = idaapi.action_desc_t(
            'vulnlink:fetch',   # The action name. This acts like an ID and must be unique
            'VulnLink',  # The action text.
            VulnLink(),   # The action handler.
            '',      # Optional: the action shortcut
            'Make VulnLink fetch the potentially interesting places in binary.',  # Optional: the action tooltip (available in menus/toolbar)
            icon_id)           # Optional: the action icon (shows when in menus/toolbars)
        idaapi.register_action(vulnlink_desc)
        idaapi.attach_action_to_menu("Search", "vulnlink:fetch", idaapi.SETMENU_APP)

    def run(self):
        pass

    def term(self):
        pass



class Hooks(idaapi.UI_Hooks):
    def __init__(self):
        self.chains_chooser = None
        self.chain_chooser = None
        self.func_chooser = None
        idaapi.UI_Hooks.__init__(self)

    def finish_populating_widget_popup(self, form, popup):
        action_text = f"get all call chains for function '{utils.get_func_name(idc.here())}'"
        function_ea = idc.here()
        try:
            # Get selected symbol
            selected_symbol, _ = ida_kernwin.get_highlight(ida_kernwin.get_current_viewer())
            # Check if it is a function name
            for function in idautils.Functions():
                if utils.get_func_name(function) in utils.prep_func_name(selected_symbol):
                    action_text = f"Add '{utils.get_func_name(function)}' function to VulnLink"
                    function_ea = function
        except:
            pass
        action_desc = idaapi.action_desc_t(
        'vulnlink:get_one',   # The action name. This acts like an ID and must be unique
        action_text,  # The action text.
        VulnLink_Single_Function(function_ea),   # The action handler.
        '',      # Optional: the action shortcut
        'Make VulnLink look for all interesting refences of this function.',  # Optional: the action tooltip (available in menus/toolbar)
        icon_id)           # Optional: the action icon (shows when in menus/toolbars)
        idaapi.unregister_action("vulnlink:get_one")
        idaapi.register_action(action_desc)
        if ida_kernwin.get_widget_title(form) == 'VulnLink Results':
            idaapi.attach_action_to_popup(form, popup, "vulnlink:get_one", "")
        

    def current_widget_changed(self, widget, prev_widget):
        title = ida_kernwin.get_widget_title(widget)
        if title and title == 'VulnLink Call Chain Results' and self.chains_chooser:
            self.chains_chooser.Refresh()
        elif title and title == 'single call chain' and self.chain_chooser:
            self.chain_chooser.Refresh()
        elif self.func_chooser:
            self.func_chooser.Refresh()

    def set_func_chooser(self, func_chooser):
        self.func_chooser = func_chooser
        
    def set_chains_chooser(self, chains_chooser):
        self.chains_chooser = chains_chooser
        
    def set_chain_chooser(self, chain_chooser):
        self.chain_chooser = chain_chooser


# Run the hooks
hooks = Hooks()
hooks.hook()


def PLUGIN_ENTRY():
    return vulnlink_fetch_t()


