# plugins

repository: zmrbak/IDAPython
file: 配套代码/ida10/python/examples/core/auto_instantiate_widget_plugin.py
description: 在线课程《IDAPython脚本化软件逆向分析入门》课程源码
wanted_name: autoinst
created_at: 2022-05-02 02:13:34+00:00
pushed_at: 2022-05-14 10:41:43+00:00
url:
https://github.com/zmrbak/IDAPython/blob/478278339ba0c150b3ace657296233d8a89b26ea/%E9%85%8D%E5%A5%97%E4%BB%A3%E7%A0%81/ida10/python/examples/core/auto_instantiate_widget_plugin.py

content:
"""
summary: better integrating custom widgets in the desktop layout

description:
  This is an example demonstrating how one can create widgets from a plugin,
  and have them re-created automatically at IDA startup-time or at desktop load-time.

  This example should be placed in the 'plugins' directory of the
  IDA installation, for it to work.

  There are 2 ways to use this example:
  1) reloading an IDB, where the widget was opened
     - open the widget ('View > Open subview > ...')
     - save this IDB, and close IDA
     - restart IDA with this IDB
       => the widget will be visible

  2) reloading a desktop, where the widget was opened
     - open the widget ('View > Open subview > ...')
     - save the desktop ('Windows > Save desktop...') under, say, the name 'with_auto'
     - start another IDA instance with some IDB, and load that desktop
       => the widget will be visible

keywords: desktop
"""

import ida_idaapi
import ida_kernwin

title = "Auto-instantiable at IDA startup"

# -----------------------------------------------------------------------
class auto_inst_t(ida_kernwin.simplecustviewer_t):
    def __init__(self):
        ida_kernwin.simplecustviewer_t.__init__(self)

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

        text = __doc__
        for l in text.split("\n"):
            self.AddLine(l)
        return True

# -----------------------------------------------------------------------
auto_inst = None

# -----------------------------------------------------------------------
def register_open_action():
    """
    Provide the action that will create the widget
    when the user asks for it.
    """
    class create_widget_t(ida_kernwin.action_handler_t):
        def activate(self, ctx):
            if ida_kernwin.find_widget(title) is None:
                global auto_inst
                auto_inst = auto_inst_t()
                assert(auto_inst.Create())
                assert(auto_inst.Show())

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

    action_name = "autoinst:create"
    ida_kernwin.register_action(
        ida_kernwin.action_desc_t(
            action_name,
            title,
            create_widget_t()))
    ida_kernwin.attach_action_to_menu(
        "View/Open subviews/Strings",
        action_name,
        ida_kernwin.SETMENU_APP)

# -----------------------------------------------------------------------
auto_inst_hooks = None
def register_autoinst_hooks():
    """
    Register hooks that will create the widget when IDA
    requires it because of the IDB/desktop
    """
    class auto_inst_hooks_t(ida_kernwin.UI_Hooks):
        def create_desktop_widget(self, ttl, cfg):
            if ttl == title:
                global auto_inst
                auto_inst = auto_inst_t()
                assert(auto_inst.Create())
                return auto_inst.GetWidget()

    global auto_inst_hooks
    auto_inst_hooks = auto_inst_hooks_t()
    auto_inst_hooks.hook()

# -----------------------------------------------------------------------
class auto_inst_plugin_t(ida_idaapi.plugin_t):
    flags = 0
    comment = "This plugin creates a widget that will be recreated automatically if needed, either at startup or when loading a desktop that requires it"
    help = "No help, really"
    wanted_name = "autoinst"
    wanted_hotkey = ""

    def init(self):
        register_open_action()
        register_autoinst_hooks()

    def run(self, arg):
        pass

    def term(self):
        pass


def PLUGIN_ENTRY():
    return auto_inst_plugin_t()

comment:
This plugin creates a widget that will be recreated automatically if needed, either at startup or when loading a desktop that requires it


repository: aliyunav/Finger
file: finger_plugin.py
description: A tool for recognizing function symbol
url: https://github.com/aliyunav/Finger/blob/f9c493baad8c540b66f135f8aea64145b2293f95/finger_plugin.py
wanted_name: Finger
created_at: 2021-10-14 07:52:09+00:00
pushed_at: 2024-04-08 07:40:01+00:00
content:
import idc
import idaapi
import idautils
import traceback
from finger_sdk import client, ida_func


class FingerManager:
    def __init__(self):
        self.url = "https://sec-lab.aliyun.com/finger/recognize/"
        self.headers = {'content-type': 'application/json'}
        self.timeout = 5
        self.client = None


    def recognize_function(self, start_ea):
        func_symbol = None
        try:
            self.client = client.Client(self.url, self.headers, self.timeout)
            func_feat = ida_func.get_func_feature(start_ea)
            if func_feat:
                func_id, res = self.client.recognize_function(func_feat)
                if res and res[func_id]:
                    func_symbol = res[func_id]
        except Exception as e:
            print(traceback.format_exc())
        if func_symbol:
            func_symbol = str(func_symbol)  # python2 unicode to str
        return func_symbol


    def recognize_selected_function(self, funcs):
        for pfn in funcs:
            func_name = idc.get_func_name(pfn.start_ea)
            func_symbol = self.recognize_function(pfn.start_ea)
            if func_symbol:
                idc.set_color(pfn.start_ea, idc.CIC_FUNC, 0x98FF98)
                idaapi.set_name(pfn.start_ea, func_symbol, idaapi.SN_FORCE)
                idaapi.update_func(pfn)
                print("[+]Recognize %s: %s" %(func_name, func_symbol))
            else:
                print("[-]%s recognize failed" %(func_name))


    def recognize_function_callback(self, menupath):
        ea = idaapi.get_screen_ea()
        pfn = idaapi.get_func(ea)
        if pfn:
            func_name = idc.get_func_name(pfn.start_ea)
            func_symbol = self.recognize_function(pfn.start_ea)
            if func_symbol:
                idc.set_color(pfn.start_ea, idc.CIC_FUNC, 0x98FF98)
                idaapi.set_name(pfn.start_ea, func_symbol, idaapi.SN_FORCE)
                idaapi.update_func(pfn)
                print("[+]Recognize %s: %s" %(func_name, func_symbol))
            else:
                print("[-]%s recognize failed" %(func_name))
        else:
            print("[-]0x%x is not a function" %ea)


    def recognize_functions_callback(self, menupath):
        funcs = []
        for ea in idautils.Functions():
            funcs.append(idaapi.get_func(ea))
        self.recognize_selected_function(funcs)


class FingerUIManager:
    class UIHooks(idaapi.UI_Hooks):
        def finish_populating_widget_popup(self, widget, popup):
            if idaapi.get_widget_type(widget) == idaapi.BWN_FUNCS:
                idaapi.attach_action_to_popup(widget, popup, "Finger:RecognizeSelected", "Finger/")
            if idaapi.get_widget_type(widget) == idaapi.BWN_DISASM:
                idaapi.attach_action_to_popup(widget, popup, "Finger:RecognizeFunction", "Finger/")


    class ActionHandler(idaapi.action_handler_t):
        def __init__(self, name, label, shortcut=None, tooltip=None, icon=-1, flags=0):
            idaapi.action_handler_t.__init__(self)
            self.name = name
            self.action_desc = idaapi.action_desc_t(name, label, self, shortcut, tooltip, icon, flags)

        def register_action(self, callback,  menupath=None):
            self.callback = callback
            if not idaapi.register_action(self.action_desc):
                return False
            if menupath and not idaapi.attach_action_to_menu(menupath, self.name, idaapi.SETMENU_APP):
                return False
            return True

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

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


    def __init__(self, name):
        self.name = name
        self.mgr = FingerManager()
        self.hooks = FingerUIManager.UIHooks()

    def register_actions(self):
        menupath = self.name
        idaapi.create_menu(menupath, self.name, "Help")

        action = FingerUIManager.ActionHandler("Finger:RecognizeFunctions", "Recognize all functions", "")
        action.register_action(self.mgr.recognize_functions_callback, menupath)
        action = FingerUIManager.ActionHandler("Finger:RecognizeFunction", "Recognize function", "")
        action.register_action(self.mgr.recognize_function_callback, menupath)
        recognize_action = FingerUIManager.ActionHandler("Finger:RecognizeSelected", "Recognize function")
        if recognize_action.register_action(self.selected_function_callback):
            self.hooks.hook()
            return True
        return False


    def selected_function_callback(self, ctx):
        funcs = map(idaapi.getn_func, ctx.chooser_selection)
        if ctx.action == "Finger:RecognizeSelected":
            self.mgr.recognize_selected_function(funcs)


def check_ida_version():
    if idaapi.IDA_SDK_VERSION < 700:
        print("[-]Finger support 7.x IDA, please update your IDA version.")
        return False
    return True


class FingerPlugin(idaapi.plugin_t):
    wanted_name = "Finger"
    comment, help, wanted_hotkey = "", "", ""
    flags = idaapi.PLUGIN_FIX | idaapi.PLUGIN_HIDE | idaapi.PLUGIN_MOD

    def init(self):
        if check_ida_version():
            idaapi.msg("[+]Finger plugin starts\n")
            manager = FingerUIManager(FingerPlugin.wanted_name)
            if manager.register_actions():
                return idaapi.PLUGIN_OK
        return idaapi.PLUGIN_SKIP

    def run(self, ctx):
        return

    def term(self):
        return


def PLUGIN_ENTRY():
    return FingerPlugin()


repository: sean2077/jumptohex-ida
file: jumptohex.py
comment: Jump to Hex View
created_at: 2024-06-16 00:38:23+00:00
pushed_at: 2024-06-16 06:07:34+00:00
description:
IDA plugin for jumping from the designated address of IDA View to the corresponding location of Hex View.

url:
https://github.com/sean2077/jumptohex-ida/blob/d8a5088f8cd4e4505291d9309b039c9a5df1c9e4/jumptohex.py

content:
############################################################################################
##
## IDA Jump to Hex View!
##
## Available for IDA 7+ and Python 3.8+
##
## To install:
##      Copy script into plugins directory, i.e: C:\Program Files\<ida version>\plugins
##
## To run:
##      Right-click on an address in the disassembly and select "Jump to Hex View"
##
############################################################################################

__AUTHOR__ = "@sean2077"

PLUGIN_NAME = "Jump to Hex View"
PLUGIN_HOTKEY = "Ctrl+Shift+J"
VERSION = "1.0.0"

ACTION_PREFIX = "sean2077"

import ida_kernwin
import idaapi
import idc


def jump_to_hex_action():
    ea = idc.get_screen_ea()
    if ea == idaapi.BADADDR:
        idaapi.warning("Invalid address selected")
        return

    idaapi.msg(f"Jumping to Hex View at address: {ea:#x}")

    # Check if Hex View is open
    hex_view = idaapi.find_widget("Hex View-1")
    if hex_view is None:
        idaapi.refresh_idaview_anyway()
        ida_kernwin.process_ui_action("ToggleDump")
        hex_view = idaapi.find_widget("Hex View-1")

    if hex_view:
        ida_kernwin.activate_widget(hex_view, True)
        ida_kernwin.jumpto(ea)
    else:
        idaapi.warning("Failed to open Hex view, please open it manually from the View menu.")


class JumpToHex(idaapi.plugin_t):
    flags = idaapi.PLUGIN_PROC | idaapi.PLUGIN_HIDE
    comment = "Jump to Hex View"
    help = "Right-click an address and select 'Jump to Hex'"
    wanted_name = PLUGIN_NAME
    wanted_hotkey = PLUGIN_HOTKEY

    def init(self):
        self._init_action_jump_to_hex()
        self._init_hooks()
        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):
        self._hooks.unhook()
        self._del_action_jump_to_hex()
        idaapi.msg("%s terminated...\n" % self.wanted_name)

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

    ACTION_JUMP_TO_HEX = f"{ACTION_PREFIX}:jump_to_hex"

    def _init_action_jump_to_hex(self):
        action_desc = idaapi.action_desc_t(
            self.ACTION_JUMP_TO_HEX,
            "Jump to Hex",
            IDACtxEntry(jump_to_hex_action),
            PLUGIN_HOTKEY,
            "Jump to hex view at the current address",
            0,
        )
        assert idaapi.register_action(action_desc), "Action registration failed"

    def _del_action_jump_to_hex(self):
        idaapi.unregister_action(self.ACTION_JUMP_TO_HEX)


class Hooks(idaapi.UI_Hooks):
    def finish_populating_widget_popup(self, widget, popup):
        inject_jump_to_hex_actions(widget, popup, idaapi.get_widget_type(widget))
        return 0


def inject_jump_to_hex_actions(form, popup, form_type):
    if form_type == idaapi.BWN_DISASM:
        idaapi.attach_action_to_popup(form, popup, JumpToHex.ACTION_JUMP_TO_HEX, "Jump to Hex", idaapi.SETMENU_APP)
    return 0


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


def PLUGIN_ENTRY():
    return JumpToHex()


repository: josephH00/ida-InTooDeep
file: InTooDeep.py
description: IDA Pro Folder Management Made Easy
comment: Make it easier to manage folders in IDA
created_at: 2024-08-11 11:05:43+00:00
pushed_at: 2024-12-24 09:58:18+00:00
url:
https://github.com/josephH00/ida-InTooDeep/blob/639258049e8033a939320d9371f961e83db6dfad/InTooDeep.py

content:
import idaapi
import ida_kernwin
import ida_dirtree
import ida_funcs
import idc
from PyQt5 import QtWidgets, sip, QtCore

PLUGIN_NAME = "In Too Deep"
moveHereCustomIconId = -1
newFolderCustomIconId = -1


def getTreeForWidget(widget):
    mapping = {
        idaapi.BWN_FUNCS: ida_dirtree.DIRTREE_FUNCS,
        
        idaapi.BWN_PSEUDOCODE: ida_dirtree.DIRTREE_FUNCS,
        
        idaapi.BWN_DISASM: ida_dirtree.DIRTREE_FUNCS,
        
        idaapi.BWN_IMPORTS: ida_dirtree.DIRTREE_IMPORTS,
        
        idaapi.BWN_LOCTYPS: ida_dirtree.DIRTREE_LOCAL_TYPES,
            
        idaapi.BWN_NAMES: ida_dirtree.DIRTREE_NAMES,
    }
    
    # Retain compatibility for removed Enums and Structs views in IDA versions 8.x and lower
    if idaapi.IDA_SDK_VERSION < 900:
        mapping[idaapi.BWN_ENUMS] = ida_dirtree.DIRTREE_ENUMS
        mapping[idaapi.BWN_STRUCTS] = ida_dirtree.DIRTREE_STRUCTS
    
    return mapping.get(widget, None)


def getBytesFromQStyleIcon(icon):
    pixmap = QtWidgets.QWidget().style().standardIcon(icon).pixmap(100)

    byteArray = QtCore.QByteArray()
    buff = QtCore.QBuffer(byteArray)
    buff.open(QtCore.QIODevice.WriteOnly)
    pixmap.save(buff, "PNG")

    return byteArray.data()


def getFuncUnderCursor(widget, cursorEA):
    vu = idaapi.get_widget_vdui(widget)

    if vu:
        cursorEA = vu.item.get_ea()
    else:
        opNum = idaapi.get_opnum()
        if opNum != -1:
            # Cursor under operand value with function
            opAddr = idc.get_operand_value(cursorEA, opNum)
            f = idaapi.get_func(opAddr)
            if f and f.start_ea == opAddr:
                return opAddr

    # Cursor directly over a function
    f = idaapi.get_func(cursorEA)
    if f and f.start_ea == cursorEA:
        return cursorEA


def getFuncAbsPath(ea):
    tree = ida_dirtree.get_std_dirtree(ida_dirtree.DIRTREE_FUNCS)
    funcName = ida_funcs.get_func_name(ea)

    absPath = ""

    class treeTraversal(ida_dirtree.dirtree_visitor_t):
        def visit(self, cursor, direntry):
            if direntry.isdir or funcName != tree.get_entry_name(direntry):
                return 0

            nonlocal absPath
            absPath = tree.get_abspath(cursor)

            return -1

    tree.traverse(treeTraversal())
    return absPath


class MoveHereActionHandler(idaapi.action_handler_t):
    def __init__(self, destPath, srcPaths=None):
        ida_kernwin.action_handler_t.__init__(self)
        self.destPath = destPath
        self.srcPaths = srcPaths

    def activate(self, ctx): 
        widgetType = idaapi.get_widget_type(ctx.widget)
        treeType = getTreeForWidget(widgetType)
        tree = ida_dirtree.get_std_dirtree(treeType)
        
        if not ctx.dirtree_selection and not (widgetType == idaapi.BWN_PSEUDOCODE or widgetType == idaapi.BWN_DISASM):
            print("Please enable 'Show folders' for this view before continuing")
            return
       
        selectedItems = []
        if widgetType == idaapi.BWN_DISASM or widgetType == idaapi.BWN_PSEUDOCODE:
            ea = getFuncUnderCursor(ctx.widget, ctx.cur_ea)
            selectedItems = [getFuncAbsPath(ea)]
        elif self.srcPaths:
             selectedItems = self.srcPaths
        else:
            selectedItems = [tree.get_abspath(i) for i in ctx.dirtree_selection]

        for i in selectedItems:
            itemName = i.split("/")[-1]
            tree.rename(i, self.destPath + "/" + itemName) 

    def update(self, ctx):
        pass


class CreateNewFolderActionHandler(idaapi.action_handler_t):
    def __init__(self, destPath):
        ida_kernwin.action_handler_t.__init__(self)
        self.destPath = destPath

    def activate(self, ctx):        
        folderName = ida_kernwin.ask_str("", ida_kernwin.HIST_IDENT, "Please enter folder name")
        newFolderPath = self.destPath + "/" + folderName

        widgetType = idaapi.get_widget_type(ctx.widget)
        treeType = getTreeForWidget(widgetType)
        tree = ida_dirtree.get_std_dirtree(treeType)
        
        absPathOfSelectedItems = None
        if widgetType == idaapi.BWN_DISASM or widgetType == idaapi.BWN_PSEUDOCODE:
            # The disassembly and pseudocode views don't populate ctx.dirtree_selection, but ida_dirtree.DIRTREE_FUNCS can still be modified without the show folders view being active
            ea = getFuncUnderCursor(ctx.widget, ctx.cur_ea)
            absPathOfSelectedItems = [getFuncAbsPath(ea)]
        else:        
            if not ctx.dirtree_selection:
                print("Please enable 'Show folders' for this view before continuing")
                return
            
            absPathOfSelectedItems = [tree.get_abspath(i) for i in ctx.dirtree_selection] # Creation of a new dir will invalidate previous ctx

        tree.mkdir(newFolderPath)

        mh = MoveHereActionHandler(newFolderPath, srcPaths=absPathOfSelectedItems)
        mh.activate(ctx)

    def update(self, ctx):
        pass


class ContextMenuHooks(idaapi.UI_Hooks):
    def finish_populating_widget_popup(self, widget, popup):
        widgetType = idaapi.get_widget_type(widget)
        treeType = getTreeForWidget(widgetType)
        if not treeType:
            return  # Don't show menu for views that don't support folders
        
        tree = ida_dirtree.get_std_dirtree(treeType)

        if widgetType == idaapi.BWN_DISASM or widgetType == idaapi.BWN_PSEUDOCODE:
            if getFuncUnderCursor(widget, idaapi.get_screen_ea()) == None:
                return  # Don't show menu in the disassembly or pseudocode view if cursor is not over a function

        class treeTraversal(ida_dirtree.dirtree_visitor_t):
            def visit(self, cursor, direntry):
                global moveHereCustomIconId, newFolderCustomIconId

                if not direntry.isdir:
                    return 0

                absPath = tree.get_abspath(cursor)

                mhHandler = MoveHereActionHandler(absPath)
                cnfHandler = CreateNewFolderActionHandler(absPath)

                if absPath == "/":
                    absPath = "/[root]"

                path = "Move To" + absPath + "/"

                idaapi.attach_dynamic_action_to_popup(
                    widget,
                    popup,
                    idaapi.action_desc_t(
                        None, 
                        "Here", 
                        mhHandler, 
                        None,
                        None, 
                        moveHereCustomIconId
                    ),
                    path,
                    idaapi.SETMENU_APP,
                )

                idaapi.attach_dynamic_action_to_popup(
                    widget,
                    popup,
                    idaapi.action_desc_t(
                        None,
                        "Create New Folder",
                        cnfHandler,
                        None,
                        None,
                        newFolderCustomIconId,
                    ),
                    path,
                    idaapi.SETMENU_APP,
                )

                return 0

        tree.traverse(treeTraversal())


class PluginIconHook(idaapi.UI_Hooks):
    def updated_actions(self):
        global PLUGIN_NAME, moveHereCustomIconId
        ida_kernwin.update_action_icon("Edit/Plugins/" + PLUGIN_NAME, moveHereCustomIconId)


class InTooDeepPlugin(idaapi.plugin_t):
    comment = "Make it easier to manage folders in IDA"
    help = ""
    wanted_name = PLUGIN_NAME
    wanted_hotkey = ""
    flags = 0

    def init(self):
        global PLUGIN_NAME, moveHereCustomIconId, newFolderCustomIconId

        moveHereCustomIconId = idaapi.load_custom_icon(
            data=getBytesFromQStyleIcon(QtWidgets.QStyle.SP_DirClosedIcon), format="png"
        )
        newFolderCustomIconId = idaapi.load_custom_icon(
            data=getBytesFromQStyleIcon(QtWidgets.QStyle.SP_DirOpenIcon), format="png"
        )

        self.contextMenuHook = ContextMenuHooks()
        self.contextMenuHook.hook()
        
        self.pluginIconHook = PluginIconHook()
        self.pluginIconHook.hook()

        print("{} - Hooks registered".format(PLUGIN_NAME))
        return idaapi.PLUGIN_KEEP

    def run(*args):
        pass

    def term(self):
        global moveHereCustomIconId, newFolderCustomIconId

        idaapi.free_custom_icon(moveHereCustomIconId)
        idaapi.free_custom_icon(newFolderCustomIconId)

        self.contextMenuHook.unhook()
        self.pluginIconHook.unhook()


def PLUGIN_ENTRY():
    return InTooDeepPlugin()


repository: cmu-sei/pharos
file: tools/ooanalyzer/ida/OOAnalyzer.py
description: Automated static analysis tools for binary programs
wanted_name: OOAnalyzer
comment: IDA OOAnalyzer plugin
created_at: 2015-06-12 17:51:42+00:00
pushed_at: 2025-03-27 20:13:57+00:00
url:
https://github.com/cmu-sei/pharos/blob/1a4d4b9b315bf6356d7ffc42e32a786711fd06a8/tools/ooanalyzer/ida/OOAnalyzer.py

content:
'''Created on Jan 8, 2014

@version: 2.0
@organization: CERT Malicious Code Team
@author: jsg

This is the IDA plugin to apply OOAnalyzer output into IDA
Pro. Version 2.0 is the first version that works with OOAnalyzer.

'''
try:
    from PyQt5.QtWidgets import *
    from PyQt5.QtGui import *
    from PyQt5.QtCore import *
    import six
    import logging
    import zlib
    import traceback
    import re
    import copy
    import json
    from abc import ABCMeta, abstractmethod
    import sys
    import codecs
    import ida_bytes
    import ida_idp
    import ida_kernwin
    import ida_name
    import ida_struct
    import ida_ua
    import ida_xref
    import idaapi
    import idaapi
    import idautils
    import idc
except ImportError:
    print("Could not import IDA Python modules")
    sys.exit(-1)

# In lieu of using "six", let's make out own Python2/3 changes ourselves.
python_version_2 = (sys.version_info[0] == 2)

PLUGIN_VERSION = 2.0

EXPECTED_JSON_VERSION = "2.2.0"

BLOB_SIZE = 1024
OUR_NETNODE = "$ com.williballenthin"
INT_KEYS_TAG = 'M'
STR_KEYS_TAG = 'N'
STR_TO_INT_MAP_TAG = 'O'
INT_TO_INT_MAP_TAG = 'P'
OOA_JSON = "OOA_JSON"
APPLIED_CLASS_NAMES = "OOA_APPLIED"

logger = logging.getLogger(__name__)

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

def ida_hexify(value):
    '''
    convert hex string (0xNNNN...) to IDA-like hex value (NNNh)
    '''
    val = value

    # convert integers to strings and make them "IDA-like"
    if not isinstance(value, str):
        val = hex(value)

    if val.find("0x") != -1:
        ret_val = val[val.find("0x") + 2:]
        if ret_val.find('L') != -1:  # in newer versions of IDA an 'L' is appended to immediates
            ret_val = ret_val[:-1]
        return ret_val
    else:
        return val


def sanitize_name(name):
    if name != None:
        # strip off non-alphanumeric characters because IDA kind of sucks
        return re.sub('[^0-9a-zA-Z]+', '_', name)
    return None


def generate_vftable_name(classes, vft, cls, off):
    '''
    Generate vftable names of the form X_Y_... where Y is the parent of X. This captures the
    lineage of overwritten vftptrs
    '''

    lineage = []
    c = cls

    while c is not None:

        if c in lineage:
            print("WARNING: Class %s is in it's own lineage. This shouldn't happen" % c)
            break

        lineage.append(c)

        # add the current parent
        if off in c.parents:

            # go to the next parent at this offset. This will capture over-written vftptrs
            c = next(
                (pc for pc in classes if c.parents[off].name == pc.name), None)

            if c is None:
                print("WARNING: Unable to find correct parent")
                assert False

        else:
            c = None

    vft_parts = []
    for x in lineage:
        if x.ida_name != "":
            vft_parts.append(x.ida_name)
        else:
            vft_parts.append(x.name)

    vft_name = "_".join(str(y) for y in vft_parts)

    print("Vritual function table name %s_vftable" % vft_name)
    return "%s_vftable" % vft_name

# The code below was taken from:
# 
# https://github.com/williballenthin/ida-netnode/tree/master
#  
# To manage storage in an IDA IDB.

class NetnodeCorruptError(RuntimeError):
    pass

class Netnode(object):
    """
    A netnode is a way to persistently store data in an IDB database.
    The underlying interface is a bit weird, so you should read the IDA
      documentation on the subject. Some places to start:

      - https://www.hex-rays.com/products/ida/support/sdkdoc/netnode_8hpp.html
      - The IDA Pro Book, version 2

    Conceptually, this netnode class represents is a key-value store
      uniquely identified by a namespace.

    This class abstracts over some of the peculiarities of the low-level
      netnode API. Notably, it supports indexing data by strings or
      numbers, and allows values to be larger than 1024 bytes in length.

    This class supports keys that are numbers or strings.
    Values must be JSON-encodable. They can not be None.

    Implementation:
     (You don't have to worry about this section if you just want to
        use the library. Its here for potential contributors.)

      The major limitation of the underlying netnode API is the fixed
        maximum length of a value. Values must not be larger than 1024
        bytes. Otherwise, you must use the `blob` API. We do that for you.

      The first enhancement is transparently zlib-encoding all values.

      To support arbitrarily sized values with keys of either int or str types,
        we store the values in different places:

        - integer keys with small values: stored in default supval table
        - integer keys with large values: the data is stored in the blob
           table named 'M' using an internal key. The link from the given key
           to the internal key is stored in the supval table named 'P'.
        - string keys with small values: stored in default hashval table
        - string keys with large values: the data is stored in the blob
           table named 'N' using an integer key. The link from string key
           to int key is stored in the supval table named 'O'.
    """

    def __init__(self, netnode_name):
        self._netnode_name = netnode_name
        self._n = idaapi.netnode(netnode_name, 0, True)

    @staticmethod
    def _decompress(data):
        '''
        args:
          data (bytes): the data to decompress

        returns:
          bytes: the decompressed data.
        '''
        return zlib.decompress(data)

    @staticmethod
    def _compress(data):
        '''
        args:
          data (bytes): the data to compress

        returns:
          bytes: the compressed data.
        '''
        return zlib.compress(data)

    @staticmethod
    def _encode(data):
        '''
        args:
          data (object): the data to serialize to json.

        returns:
          bytes: the ascii-encoded serialized data buffer.
        '''
        return json.dumps(data).encode("ascii")

    @staticmethod
    def _decode(data):
        '''
        args:
          data (bytes): the ascii-encoded json serialized data buffer.

        returns:
          object: the deserialized object.
        '''
        return json.loads(data.decode("ascii"))

    def _intdel(self, key):
        assert isinstance(key, six.integer_types)

        did_del = False
        storekey = self._n.supval(key, INT_TO_INT_MAP_TAG)
        if storekey is not None:
            storekey = int(storekey)
            self._n.delblob(storekey, INT_KEYS_TAG)
            self._n.supdel(key, INT_TO_INT_MAP_TAG)
            did_del = True
        if self._n.supval(key) is not None:
            self._n.supdel(key)
            did_del = True

        if not did_del:
            raise KeyError("'{}' not found".format(key))

    def _get_next_slot(self, tag):
        '''
        get the first unused supval table key, or 0 if the
         table is empty.
        useful for filling the supval table sequentially.
        '''
        slot = self._n.suplast(tag)
        if slot is None or slot == idaapi.BADNODE:
            return 0
        else:
            return slot + 1

    def _intset(self, key, value):
        assert isinstance(key, six.integer_types)
        assert value is not None

        try:
            self._intdel(key)
        except KeyError:
            pass

        if len(value) > BLOB_SIZE:
            storekey = self._get_next_slot(INT_KEYS_TAG)
            self._n.setblob(value, storekey, INT_KEYS_TAG)
            self._n.supset(key, str(storekey).encode(
                'utf-8'), INT_TO_INT_MAP_TAG)
        else:
            self._n.supset(key, value)

    def _intget(self, key):
        assert isinstance(key, six.integer_types)

        storekey = self._n.supval(key, INT_TO_INT_MAP_TAG)
        if storekey is not None:
            storekey = int(storekey.decode('utf-8'))
            v = self._n.getblob(storekey, INT_KEYS_TAG)
            if v is None:
                raise NetnodeCorruptError()
            return v

        v = self._n.supval(key)
        if v is not None:
            return v

        raise KeyError("'{}' not found".format(key))

    def _strdel(self, key):
        assert isinstance(key, (str))

        did_del = False
        storekey = self._n.hashval(key, STR_TO_INT_MAP_TAG)
        if storekey is not None:
            storekey = int(storekey.decode('utf-8'))
            self._n.delblob(storekey, STR_KEYS_TAG)
            self._n.hashdel(key, STR_TO_INT_MAP_TAG)
            did_del = True
        if self._n.hashval(key):
            self._n.hashdel(key)
            did_del = True

        if not did_del:
            raise KeyError("'{}' not found".format(key))

    def _strset(self, key, value):
        assert isinstance(key, (str))
        assert value is not None

        try:
            self._strdel(key)
        except KeyError:
            pass

        if len(value) > BLOB_SIZE:
            storekey = self._get_next_slot(STR_KEYS_TAG)
            self._n.setblob(value, storekey, STR_KEYS_TAG)
            self._n.hashset(key, str(storekey).encode(
                'utf-8'), STR_TO_INT_MAP_TAG)
        else:
            self._n.hashset(key, bytes(value))

    def _strget(self, key):
        assert isinstance(key, (str))

        storekey = self._n.hashval(key, STR_TO_INT_MAP_TAG)
        if storekey is not None:
            storekey = int(storekey.decode('utf-8'))
            v = self._n.getblob(storekey, STR_KEYS_TAG)
            if v is None:
                raise NetnodeCorruptError()
            return v

        v = self._n.hashval(key)
        if v is not None:
            return v

        raise KeyError("'{}' not found".format(key))

    def __getitem__(self, key):
        if isinstance(key, str):
            v = self._strget(key)
        elif isinstance(key, six.integer_types):
            v = self._intget(key)
        else:
            raise TypeError("cannot use {} as key".format(type(key)))

        data = self._decompress(v)
        return self._decode(data)

    def __setitem__(self, key, value):
        '''
        does not support setting a value to None.
        value must be json-serializable.
        key must be a string or integer.
        '''
        assert value is not None

        v = self._compress(self._encode(value))
        if isinstance(key, str):
            self._strset(key, v)
        elif isinstance(key, six.integer_types):
            self._intset(key, v)
        else:
            raise TypeError("cannot use {} as key".format(type(key)))

    def __delitem__(self, key):
        if isinstance(key, str):
            self._strdel(key)
        elif isinstance(key, six.integer_types):
            self._intdel(key)
        else:
            raise TypeError("cannot use {} as key".format(type(key)))

    def get(self, key, default=None):
        try:
            return self[key]
        except (KeyError, zlib.error):
            return default

    def __contains__(self, key):
        try:
            if self[key] is not None:
                return True
            return False
        except (KeyError, zlib.error):
            return False

    def _iter_int_keys_small(self):
        # integer keys for all small values
        i = None
        if using_ida7api:
            i = self._n.supfirst()
        else:
            i = self._n.sup1st()
        while i != idaapi.BADNODE:
            yield i
            if using_ida7api:
                i = self._n.supnext(i)
            else:
                i = self._n.supnxt(i)

    def _iter_int_keys_large(self):
        # integer keys for all big values
        if using_ida7api:
            i = self._n.supfirst(INT_TO_INT_MAP_TAG)
        else:
            i = self._n.sup1st(INT_TO_INT_MAP_TAG)
        while i != idaapi.BADNODE:
            yield i
            if using_ida7api:
                i = self._n.supnext(i, INT_TO_INT_MAP_TAG)
            else:
                i = self._n.supnxt(i, INT_TO_INT_MAP_TAG)

    def _iter_str_keys_small(self):
        # string keys for all small values
        if using_ida7api:
            i = self._n.hashfirst()
        else:
            i = self._n.hash1st()
        while i != idaapi.BADNODE and i is not None:
            yield i
            if using_ida7api:
                i = self._n.hashnext(i)
            else:
                i = self._n.hashnxt(i)

    def _iter_str_keys_large(self):
        # string keys for all big values
        if using_ida7api:
            i = self._n.hashfirst(STR_TO_INT_MAP_TAG)
        else:
            i = self._n.hash1st(STR_TO_INT_MAP_TAG)
        while i != idaapi.BADNODE and i is not None:
            yield i
            if using_ida7api:
                i = self._n.hashnext(i, STR_TO_INT_MAP_TAG)
            else:
                i = self._n.hashnxt(i, STR_TO_INT_MAP_TAG)

    def iterkeys(self):
        for key in self._iter_int_keys_small():
            yield key

        for key in self._iter_int_keys_large():
            yield key

        for key in self._iter_str_keys_small():
            yield key

        for key in self._iter_str_keys_large():
            yield key

    def keys(self):
        return [k for k in list(self.iterkeys())]

    def itervalues(self):
        for k in list(self.keys()):
            yield self[k]

    def values(self):
        return [v for v in list(self.itervalues())]

    def iteritems(self):
        for k in list(self.keys()):
            yield k, self[k]

    def items(self):
        return [(k, v) for k, v in list(self.iteritems())]

    def kill(self):
        self._n.kill()
        self._n = idaapi.netnode(self._netnode_name, 0, True)


class PyOOAnalyzer(object):
    '''
    Responsible for parsing and applying JSON class specification to an IDB.
    '''

    def __init__(self):
        '''
        Constructor for the PyOOAnalyzer Plugin
        '''
        self.__json_filename = None
        self.__json_data = None
        self.__classes = []
        self.__vfcalls = []

        self.__member_usages_found = False
        self.__vcall_usages_found = False
        self.__classes_found = False

        self.__is_parsed = False

        # this is the results of parsing
        self.__parse_results = {"NumUsages": 0}

        self.__applied_classes = {}

        return

    def set_classes(self, other_classes):
        self.__classes = copy.deepcopy(other_classes)

    def __str__(self):
        '''
        Convert a class to a string
        '''
        s = ""
        for c in self.__classes:
            s += "%s\n" % c
        return s

    def set_json_file(self, jsn_file):
        '''
        Set the JSON file
        '''

        self.__json_filename = jsn_file
        result = True
        if jsn_file != None:

            try:
                self.__json_data = open(jsn_file).read()
            except:
                result = False

        # the file was opened, it is eligible for parsing
        if result == True:
            self.__is_parsed = False

        return result

    def is_parsed(self):
        '''
        Return flag indicate if the JSON file is parsed
        '''
        return self.__is_parsed

    def get_classes(self):
        return self.__classes

    def find_name(self, ea):
        for cls in self.__classes:
            for method in cls.methods:
                if method.start_ea == ea:
                    return method.method_name
        return None

    def get_json_filename(self):
        return self.__json_filename

    def get_json_data(self):
        return self.__json_data

    def set_json_data(self, data):
        self.__json_data = data

    def get_parse_results(self):
        '''
        Return results of parsing
        '''
        if len(self.__parse_results) == 0:
            return None

        return self.__parse_results

    def parse(self):
        """
        Parse the JSON object into a Python dictionary. Parse methods are responsible for properly
        setting the type of data. addresses are always base 16 and offsets are always in base 10
        """

        def _byteify(data):
            # if this is a unicode string, return its string representation
            if python_version_2:
                if isinstance(data, basestring):
                    return str(data.encode('utf-8').decode())
            else:
                if isinstance(data, str):
                    return str(data.encode('utf-8').decode())

            # if this is a list of values, return list of byteified values
            if isinstance(data, list):
                return [_byteify(item) for item in data]

            # if this is a dictionary, return dictionary of byteified keys and values
            # but only if we haven't already byteified it
            if isinstance(data, dict):
                if python_version_2:
                    return {_byteify(key): _byteify(value) for key, value in data.iteritems()}
                else:
                    return {_byteify(key): _byteify(value) for key, value in data.items()}

            # if it's anything else, return it in its original form
            return data

        if self.__json_data == None:
            return [False, "File not specified"]

        data = None
        try:
            data = json.loads(self.__json_data, object_hook=_byteify)
        except Exception as e:
            print("JSON parsing error: %s" % e)
            return [False, "JSON parsing error: %s" % e]

        file_version = data["version"] if "version" in data else "unknown"
        if file_version != EXPECTED_JSON_VERSION:
            return [False, "Expected JSON version '%s' but found '%s' instead. Aborting." % (EXPECTED_JSON_VERSION, file_version)]

        if "filemd5" in data:
            jsonMd5 = data["filemd5"]
            # Some IDA give a binary encoded version, and some give it in ascii already!
            idaMd5 = idautils.GetInputFileMD5()
            idaMd52 = codecs.encode(idaMd5, 'hex').decode("ascii")
            if jsonMd5.lower() != idaMd5.lower() and jsonMd5.lower() != idaMd52.lower():
                print(jsonMd5, idaMd5)
                if ida_kernwin.ask_yn(ida_kernwin.ASKBTN_YES, "There was a hash mismatch.  This JSON may be for a different file '%s'.  Do you want to import anyway?" % data["filename"]) == ida_kernwin.ASKBTN_NO:
                    return [False, "User aborted after hash mismatch"]

        print("Parsing JSON structures ...")
        self.__parse_structs(data["structures"])
        print("Completed Parsing JSON structures file: %s classes found" %
              len(self.__classes))

        self.__parse_vf_calls(data['vcalls'])
        print("Completed Parsing JSON class usages")

        self.__is_parsed = True
        return [True, "OK"]

    def apply(self):
        '''
        Apply classes, then usages
        '''

        print("Applying class structures")

        self.__apply_structs()

        print("Applying class virtual function calls")

        self.__apply_vfcalls()

        return True

    def apply_struct(self, cls_to_apply):
        '''
        Apply classes, then usages
        '''
        result = False
        for c in self.__classes:
            if c.ida_name == cls_to_apply.ida_name:
                if not c.applied:

                    print("Applying class %s" % c.ida_name)

                    self.apply_class(c)

                    c.applied = True
                    result = True
                    break
                else:
                    print("Class %s already applied" % c.name)
                    result = False
                    break

        return result

    def __parse_structs(self, structs):
        '''
        Parse the structures defined in the JSON file. The order here matters:

           * vftables must be parsed before methods because vftables contain additional methods
        '''
        self.__parse_results["NumClasses"] = len(structs)

        for s in structs.values():
            # use the ID of the current struct, if it exists
            c = PyClassStructure(s)

            #  The class must be created for members to work
            self.__classes.append(c)

            print("Found class %s" % c.ida_name)

        for s in structs.values():
            # Now fill out the structures
            for c in self.__classes:
                if c.name == s['name']:

                    cid = ida_struct.get_struc_id(c.ida_name)
                    if cid != idaapi.BADADDR:
                        c.id = cid

                    self.__parse_methods(c, s['methods'])
                    self.__parse_vftables(c, s['vftables'])
                    self.__parse_members(c, s['members'])

        # this point each class structure is loaded
        self.__classes_found = True

        return

    def __parse_members(self, cls, members):
        '''
        Add each member specified
        '''

        print("Parsing %d members for class %s ..." %
              (len(members), cls.ida_name))

        for m in members.values():
            offset = int(m['offset'], 16)
            if offset > 0x7FFFFFFF:
                offset -= 0x100000000

            if offset < 0:
                print("WARNING: Negative member offset (%d) found - discarding" % offset)
                continue

            # valid member - create it
            mem = PyClassMember(cls, m['type'], m['size'])
            mem.cls = cls
            mem.base = m['base']
            mem.member_name = m['name']

            self.__parse_member_usages(cls, m['usages'])

            if mem.member_type == ida_bytes.FF_STRUCT:

                mem.class_member_name = m['struc']  # This could be mangled

                # Search for the parsed class information
                cls_mem = None
                for c in self.__classes:
                    if mem.class_member_name == c.name:
                        cls_mem = c
                        break

                assert cls_mem is not None, "Unable to find class member struct %s" % mem.class_member_name

                mem.class_member = cls_mem

                if m['parent']:
                    # this is a parent
                    mem.is_parent = True
                    cls.add_parent(cls_mem, offset)
                    print("   - Found class parent, name: '%s', type: %s" %
                          (mem.member_name, mem.class_member.ida_name))
                else:
                    print("   - Found class member, name: '%s', type: %s" %
                          (mem.member_name, mem.class_member.ida_name))

            else:
                mem.class_member = None
                print("   - Found member '%s' @ offset 0x%x" %
                      (mem.member_name, offset))

            # offset of this member
            mem.offset = offset

            # add the new member if it is not from the base class
            if not mem.base:
                cls.add_member(mem, offset)

            if cls.id != idaapi.BADADDR:
                mid = idc.get_member_id(cls.id, offset)
                if mid != -1:
                    mem.id = mid

        print("Members parsed")

        return

    def __parse_methods(self, cls, methods):
        '''
        Parse the methods defined in the JSON file
        '''

        print("Parsing %d methods for class %s ..." %
              (len(methods), cls.ida_name))

        for m in methods.values():

            meth = PyClassMethod()

            meth.cls = cls
            meth.is_virtual = False  # no traditional methods are virtual

            meth.is_ctor = False
            meth.is_dtor = False
            if m['type'] == "ctor":
                meth.is_ctor = True
            elif m['type'] == "dtor" or m['type'] == "deldtor":
                meth.is_dtor = True

            #  this method is imported
            meth.is_import = False
            if m['import']:
                meth.is_import = True

            meth.start_ea = int(m['ea'], 16)

            meth.method_name = m['name']
            meth.demangled_name = m['demangled_name'] if m['demangled_name'] != "" else None
            meth.userdef_name = False

            flags = ida_bytes.get_flags(meth.start_ea)

            # In vs2010/Lite/oo, there is a function tail at 0x403ed0.  Unfortunately another
            # function tail calls that chunk, and get_func(0x403ed0) refers to the larger one.
            # This results in IDA helpfully saying that the function at 0x403ed0 is named
            # sub_402509.  Thanks IDA!
            idafunc = idaapi.get_func(meth.start_ea)
            is_func = idafunc is not None and idafunc.start_ea == meth.start_ea

            has_name = is_func \
                and ida_bytes.has_name(flags) \
                and ida_bytes.has_user_name(flags) \
                and idc.get_func_name(meth.start_ea) != ""
            # Does IDA have a name for this function?
            if has_name:
                meth.method_name = idc.get_func_name(meth.start_ea)
                meth.demangled_name = idc.demangle_name(meth.method_name,
                                                        idc.get_inf_attr(idc.INF_SHORT_DN))
                meth.demangled_name = meth.demangled_name if meth.demangled_name != "" else None
                meth.userdef_name = True

            print("   - Method %s parsed" % meth.method_name)

            cls.add_method(meth)

        print("All methods parsed")

        return

    def __parse_vftables(self, cls, vftables):
        '''
        Parse the vftables defined in the JSON file
        '''

        print("Parsing %s vftables for class %s ..." %
              (len(vftables), cls.ida_name))

        for v in vftables.values():

            vftptr_off = int(v['vftptr'], 16)
            if vftptr_off < 0:
                print(
                    "   WARNING: Negative virtual function pointer offset found (%d) - discarding" % vftptr_off)
                continue

            vft = PyClassVftable()
            vft.cls = cls
            vft.length = int(v['length'])
            vft.start_ea = int(v['ea'], 16)

            for entry in v['entries'].values():
                start_ea = int(entry['ea'], 16)

                # Look up method instead of duplicating code
                meth = next(
                    method for method in cls.methods if method.start_ea == start_ea)
                assert meth is not None

                vft.add_virtual_function(meth, int(entry['offset']))

            print("   - Adding vtable %x at offset %d" %
                  (vft.start_ea, vftptr_off))
            print("")

            vft_name = generate_vftable_name(
                self.__classes, vft, cls, vftptr_off)
            vftid = ida_struct.get_struc_id(vft_name)
            if vftid != idaapi.BADADDR:
                vft.id = vftid
            cls.add_vftable(vft, vftptr_off)

        print("Vftables parsed")

        return

    def __parse_member_usages(self, cls, mem_usages):
        '''
        Parse member usages and enqueue them for processing.
        '''
        self.__parse_results["NumUsages"] += len(mem_usages)
        for ea in mem_usages:
            mu = PyMemberUsage()
            mu.ida_name = cls.ida_name
            mu.ea = int(ea, 16)
            mu.cid = idaapi.BADADDR
            cls.add_usage(mu)
        return

    def __parse_vf_calls(self, vcalls):
        '''
        Parse virtual function calls
        '''
        print("Parsing %s virtual function calls ..." % (len(vcalls)))
        self.__parse_results["NumVcalls"] = len(vcalls)

        for call_ea, targets in vcalls.items():
            vc = PyVirtualFunctionCallUsage()
            vc.call_ea = int(call_ea, 16)

            for t in targets['targets']:
                vc.add_target_ea(int(t, 16))

            print("Parsed virtual function call: %s" % vc)

            self.__vfcalls.append(vc)

        return

    # The following methods concern applying the parsed JSON file to an IDB

    def apply_class(self, cls):
        '''
        Apply the class
        '''

        print("================================================")
        print("Applying class '%s'" % cls.ida_name)

        if cls.ida_name in self.__applied_classes:
            print("class '%s' is already applied" % cls.ida_name)
            return cls.id

        cls.id = ida_struct.add_struc(0, cls.ida_name, 0)
        cls.ptr = ida_struct.get_struc(cls.id)
        # Save the ID of this struct
        self.__applied_classes[cls.ida_name] = cls

        if cls.demangled_name != "":
            ida_struct.set_struc_cmt(
                cls.id, "Original mangled name: %s" % cls.name, 1)

        if len(cls.members) > 0:
            print("Applying %d members ..." % len(cls.members))
            for moff, mem in cls.members.items():
                self.__apply_member(cls, mem, moff)
            print("Members applied")
        else:
            print("Class with no members not allowed by IDA Pro, creating dummy member")

            idc.add_struc_member(
                cls.id, "EMPTY", 0, ida_bytes.FF_BYTE | ida_bytes.FF_DATA, 0xffffffff, 1)
        if len(cls.parents) > 0:
            print("Applying %d parents ..." % len(cls.parents))
            for poff, par in cls.parents.items():
                self.__apply_parent(cls, par, poff)
            print("Parents applied")

        if len(cls.vftables) > 0:
            print("Applying %d vftables ..." % len(cls.vftables))
            for voff, vft in cls.vftables.items():
                self.__apply_vftable(cls, vft, voff)
            print("Vftables applied")

        if len(cls.methods) > 0:
            print("Applying %d methods ..." % len(cls.methods))
            for n in cls.methods:
                self.__apply_method(n)
            print("Methods applied")

        if len(cls.usages) > 0:
            print("Applying %d usages ..." % len(cls.usages))
            for u in cls.usages:
                u.cid = cls.id
                u.apply()
            print("Methods applied")

        print("Applied class '%s'" % cls.ida_name)
        print("================================================\n")

        cls.applied = True

        return cls.id

    def __apply_member(self, cls, mem, off):
        '''
        Apply members of this class to the IDB
        '''

        if mem.is_parent:
            return

        nbytes = idaapi.BADADDR
        if mem.member_type == ida_bytes.FF_STRUCT:
            if mem.class_member.ida_name not in self.__applied_classes:
                self.apply_class(mem.class_member)
                print("Applying class member: %s" %
                      str(mem.class_member.ida_name))

            nbytes = ida_struct.get_struc_size(
                ida_struct.get_struc(mem.class_member.id))
            mem_type = ida_bytes.FF_STRUCT | ida_bytes.FF_DATA

            idc.add_struc_member(cls.id, mem.member_name,
                                 off, mem_type, mem.class_member.id, nbytes)

        else:
            # non-struct member
            nbytes = mem.size
            idc.add_struc_member(cls.id, mem.member_name,
                                 off, mem.member_type, 0xffffffff, nbytes)

        # save the member ID
        mem.id = idc.get_member_id(cls.id, off)
        if mem.id == -1:
            print("WARNING: Adding member at offset %d to %s failed" %
                  (off, cls.ida_name))

        return

    def apply_method(self, method):
        self.__apply_method(method)

    def __apply_method(self, method):
        '''
        Apply methods of this class to the IDB
        '''

        method_name = method.demangled_name if method.demangled_name is not None else method.method_name

        cmt = cmt = "%s::%s" % (method.cls.ida_name, method_name)

        if method.is_virtual == True:
            cmt = "virtual %s" % cmt

        if method.is_ctor == True:
            cmt += " (constructor)"

        if method.is_dtor == True:
            cmt += " (destructor)"

        idafunc = idaapi.get_func(method.start_ea)
        is_func = idafunc is not None and idafunc.start_ea == method.start_ea
        if is_func:
            # Check for existing name
            if idc.get_name_ea_simple(method_name) != idaapi.BADADDR:
                method_name += "_%x" % method.start_ea

            print("Renaming %s to %s" %
                  (idc.get_func_name(method.start_ea), method_name))

            idc.set_name(method.start_ea, method_name,
                         ida_name.SN_NOCHECK | ida_name.SN_FORCE)
        else:
            print("Not renaming the function at %#x %s because it is not a function in IDA!" % (
                method.start_ea, method_name))

        idc.set_func_cmt(method.start_ea, cmt, 1)

        return

    def __apply_vftable(self, cls, vft, off):
        '''
        Apply the parsed vftable(s) for this class
        '''

        vft_name = generate_vftable_name(self.__classes, vft, cls, off)

        print("Applying vftable: %s" % (vft_name))

        idc.set_name(vft.start_ea, "%s_%x" %
                     (vft_name, vft.start_ea), ida_name.SN_CHECK)

        vftid = ida_struct.add_struc(0, vft_name, 0)

        # this will likely occur if the structure name is already present in the IDB.
        if vftid == idaapi.BADADDR:
            print("Vtable already exists!")
            n = 1  # simply create an indexed struct
            while vftid == idaapi.BADADDR:
                vft_name = "%s_%d" % (vft_name, n)

                vftid = ida_struct.add_struc(0, vft_name, 0)
                n += 1

        vft.id = vftid
        vft.ptr = ida_struct.get_struc(vft.id)

        ida_struct.set_struc_cmt(
            vftid, vft.cls.ida_name + " virtual function table", 1)

        print("VFtable length = %d" % vft.length)

        for vt_index in range(0, vft.length):
            vt_off = vt_index * 4

            # Where the entry points to
            ea = idc.get_wide_dword(vft.start_ea + vt_off)

            # Look up the name in the virtual function table entries.  This only contains
            # functions believed to be defined on the current class, and will probably be
            # removed.
            if vt_index in vft.virtual_functions:
                name = vft.virtual_functions[vt_index].method_name
            else:
                name = self.find_name(ea)
                if name is None:
                    # Last resort: Use whatever name IDA has
                    name = idc.get_func_name(ea)

            # print("Virtual function entry index %d using name %s" % (vt_index, name))

            # Mark the entry as a dword
            ida_bytes.create_data(vft.start_ea + vt_off,
                                  ida_bytes.FF_DWORD, 4, idaapi.BADADDR)

            res = idc.add_struc_member(vft.id,
                                       str(name),
                                       vt_off,
                                       ida_bytes.FF_DWORD,
                                       0xffffffff,
                                       4)
            if res == ida_struct.STRUC_ERROR_MEMBER_NAME:
                name = "%s_%d" % (name, vt_index)
                res = idc.add_struc_member(vft.id,
                                           str(name),
                                           vt_off,
                                           ida_bytes.FF_DWORD,
                                           0xffffffff,
                                           4)

            # Add comment with the original value
            idc.set_member_cmt(vftid, vt_off, ida_hexify(ea), 1)

        # add the vftptr :or rename existing member to vftptr
        vptr_name = "vftptr_0x%s" % ida_hexify(off)

        if idc.get_member_name(cls.id, off) == None:
            # there is no member at this offset
            idc.add_struc_member(cls.id, vptr_name, off,
                                 ida_bytes.FF_DWORD, 0xffffffff, 4)
            idc.set_member_cmt(cls.id, off, vft_name, 1)
        else:
            # There is a member at this offset

            if (idc.get_member_flag(cls.id, off) & ida_bytes.DT_TYPE) != ida_bytes.FF_STRUCT:
                # there is already a member defined at this offset, change it to a vftptr
                # ida_struct.set_member_name(cls.ptr, off, vptr_name)
                idc.set_member_type(
                    cls.id, off, ida_bytes.FF_DWORD, 0xFFFFFFFF, 4)

        # set the vtable name on the global class list
        for c in self.__classes:
            if c.ida_name == cls.ida_name:
                c.vftables[off].vtable_name = vft_name
                break

        return

    def __apply_parent(self, cls, parent, off):
        '''
        Apply the parent classes to the IDB. This may (should) overwrite a virtual function pointer
        '''
        print("Applying parent: %s" % parent.ida_name)

        if parent.id == idaapi.BADADDR:
            self.apply_class(parent)

            # add the variable for the parent
        nbytes = ida_struct.get_struc_size(parent.id)
        name = "parent_%s" % ida_hexify(off)

        print("Parent ID: %s size: %d" % (parent.id, nbytes))

        # apply the parent, deleting any empty members :or vftptrs
        result = idc.add_struc_member(
            cls.id, name, off, ida_bytes.FF_STRUCT | ida_bytes.FF_DATA, parent.id, nbytes)
        if result == ida_struct.STRUC_ERROR_MEMBER_OFFSET:
            idc.del_struc_member(cls.id, off)
            idc.add_struc_member(
                cls.id, name, off, ida_bytes.FF_STRUCT | ida_bytes.FF_DATA, parent.id, nbytes)

        return

    def get_applied_classes(self):
        '''
        Return a list of classes that have been applied
        '''
        return [cls for name, cls in self.__applied_classes.items()]

    def apply_all_structs(self):
        '''
        Apply the parsed classes. The basic algorithm is to iteratively apply
        structs for classes without parents. Eventually all structs will have
        their parents applied
        '''

        print("OOAnalyzer has %d classes to apply:" % len(self.__classes))

        for c in self.__classes:
            if c.demangled_name != "":
                print("   * %s (Demangled name: %s)" %
                      (c.name, c.demangled_name))
            else:
                print("   * %s" % (c.name))

        worklist_count = len(self.__classes)
        # while all classes are not complete

        while worklist_count > 0:

            # idaapi.show_wait_box("Applying class %s ..." % c.demangled_name)

            for c in self.__classes:

                # Has this class been applied
                if c.applied == True:
                    continue

                do_apply = True
                missing_member = ""

                # if this class has members, check for sub classes
                if c.members:
                    for o, m in c.members.items():

                        if m == None:
                            continue
                        if m.member_type == ida_bytes.FF_STRUCT:

                            if m.class_member == None:
                                continue

                            if m.class_member.ida_name not in self.__applied_classes:
                                # don't apply members with structs that are not yet applied
                                do_apply = False
                                missing_member = m.class_member.ida_name
                                break

                if do_apply:
                    print("Applying class '%s'" % c.ida_name)
                    self.apply_class(c)
                    c.applied = True
                    worklist_count -= 1

                else:
                    print("Delaying application of class '%s' because cannot find member class '%s'" % (
                        c.ida_name, missing_member))

            # idaapi.hide_wait_box()

        return

    def apply_vfcalls(self):
        '''
        Apply the different usages for each class
        '''

        for u in self.__vfcalls:
            # apply the virtual function call
            u.apply()

        return

# =============================================================================
# End class


class PyClassUsage(object):

    OP_REG = ida_ua.o_reg     # 1
    OP_MEM_REF = ida_ua.o_mem     # 2
    OP_REG_INDEX = ida_ua.o_phrase  # 3
    OP_REG_INDEX_DISP = ida_ua.o_displ   # 4
    OP_IMMEDIATE = ida_ua.o_imm     # 5

    '''
   This interface represents the usages for classes.
   '''

    def apply(self, usage):
        raise NotImplementedError("Must extend")

# =============================================================================
# End class


class PyVirtualFunctionCallUsage(PyClassUsage):
    '''
    This class represents a virtual function call
    '''

    def __init__(self):
        '''
        Constructor
        '''
        self.__call_ea = None
        self.__targets = []

        return

    def __str__(self):

        tstr = ""
        for t in self.targets:
            tstr += ("%x " % t)

        return "Call %x -> Target(s): %s" % (self.__call_ea, tstr)

    def get_call_ea(self):
        return self.__call_ea

    def get_target_ea(self):
        return self.__targets

    def set_call_ea(self, value):
        self.__call_ea = value

    def add_target_ea(self, t):
        self.__targets.append(t)

    def apply(self):
        '''
        Apply virtual function call by making an xref to the right place in the IDB
        '''
        print("================================================")
        for t in self.__targets:
            print("Applying virtual function call %x -> %x" %
                  (self.__call_ea, t))

            ida_xref.add_cref(self.__call_ea, t,
                              ida_xref.XREF_USER | ida_xref.fl_F)

            cmt = idc.GetCommentEx(self.__call_ea, 0)
            if cmt != None:
                cmt = cmt + ", %s" % ida_hexify(t)
                print("Multiple virtual function calls detected at %x (%s)" %
                      (self.__call_ea, cmt))

            else:
                cmt = "%s" % ida_hexify(t)

            idc.set_cmt(self.__call_ea, cmt, 0)

        print("================================================")

        return True

    # properties of virtual function calls

    call_ea = property(get_call_ea, set_call_ea, None, None)
    targets = property(get_target_ea, None, None, None)


# =============================================================================
# End class

class PyMemberUsage(PyClassUsage):

    def __init__(self, name=None):
        self.__ea = None
        self.__ida_name = name
        self.__cid = idaapi.BADADDR

        return

    def apply(self):
        '''
        Apply a usage
        '''

        if self.__ida_name is None:
            return False

        if self.__cid == idaapi.BADADDR:
            print("*** Cannot apply because bad ID")
            return False

        print("================================================")
        print("Applying class %s to %x" % (self.__ida_name, self.__ea))
        print("================================================")

        # there must be two and only two operands, assume 1 and check for 2
        n = 0
        if idc.get_operand_type(self.__ea, 1) in [self.OP_REG_INDEX, self.OP_REG_INDEX_DISP]:
            n = 1

        idc.op_stroff(self.__ea, n, self.__cid, 0)

        return True

    def get_cid(self):
        return self.__cid

    def get_ida_name(self):
        return self.__ida_name

    def get_ea(self):
        return self.__ea

    def set_cid(self, value):
        self.__cid = value

    def set_ea(self, value):
        self.__ea = value

    def set_ida_name(self, value):
        self.__ida_name = value

    ida_name = property(get_ida_name, set_ida_name, None, None)
    cid = property(get_cid, set_cid, None, None)
    ea = property(get_ea, set_ea, None, None)

# =============================================================================
# End class


class PyVtableUsage(PyClassUsage):

    def __init__(self, cls):
        self.__call_ea = None
        self.__cls = cls
        self.__vf_offset = None

        return

    def get_call_ea(self):
        return self.__call_ea

    def get_cls(self):
        return self.__cls

    def get_vf_offset(self):
        return self.__vf_offset

    def set_call_ea(self, value):
        self.__call_ea = value

    def set_cls(self, value):
        self.__cls = value

    def set_vf_offset(self, value):
        self.__vf_offset = value

    def del_call_ea(self):
        del self.__call_ea

    def del_cls(self):
        del self.__cls

    def del_vf_offset(self):
        del self.__vf_offset

    def apply(self):
        pass

    call_ea = property(get_call_ea, set_call_ea,
                       del_call_ea, "call_ea's docstring")
    cls = property(get_cls, set_cls, del_cls, "class_name's docstring")
    vf_offset = property(get_vf_offset, set_vf_offset,
                         del_vf_offset, "vf_offset's docstring")

# =============================================================================
# End class


# =============================================================================
# End class

class PyClassMethod(object):
    '''
    This class represents a class method
    '''

    def __init__(self, meth_name=None, cls=None, ea=None, ctor=False, dtor=False, virt=False, imp=False, udn=False):

        # Set the formatter for this object
        self.__start_ea = None
        if ea != None:
            self.__start_ea = int(ea)

        self.__method_name = None
        if meth_name != None:
            self.__method_name = meth_name

        self.__cls = None
        if cls != None:
            self.__cls = cls

        self.__userdef_name = udn
        self.__is_ctor = ctor
        self.__is_dtor = dtor
        self.__is_virtual = virt
        self.__is_import = imp

        return

    def get_best_method_name(self):
        if self.demangled_name is not None:
            s = self.demangled_name
            return re.sub(r'^.*?::', '', s)
        else:
            return self.method_name

    def __str__(self):
        return "%s @ %x" % (self.__method_name, self.__start_ea)

    def get_cls(self):
        return self.__cls

    def get_is_virtual(self):
        return self.__is_virtual

    def get_is_import(self):
        return self.__is_import

    def get_is_ctor(self):
        return self.__is_ctor

    def get_is_dtor(self):
        return self.__is_dtor

    def get_method_name(self):
        return self.__method_name

    def get_start_ea(self):
        return self.__start_ea

    def get_userdef_name(self):
        return self.__userdef_name

    def set_is_virtual(self, value):
        self.__is_virtual = value

    def set_is_import(self, value):
        self.__is_import = value

    def set_cls(self, value):
        self.__cls = value

    def set_is_ctor(self, value):
        self.__is_ctor = value

    def set_is_dtor(self, value):
        self.__is_dtor = value

    def set_method_name(self, value):
        self.__method_name = value

    def set_start_ea(self, value):
        self.__start_ea = value

    def set_userdef_name(self, value):
        self.__userdef_name = value

    # properties for class methods
    userdef_name = property(get_userdef_name, set_userdef_name, None, None)
    method_name = property(get_method_name, set_method_name, None, None)
    is_virtual = property(get_is_virtual, set_is_virtual, None, None)
    is_import = property(get_is_import, set_is_import, None, None)
    is_ctor = property(get_is_ctor, set_is_ctor, None, None)
    is_dtor = property(get_is_dtor, set_is_dtor, None, None)
    cls = property(get_cls, set_cls, None, None)
    start_ea = property(get_start_ea, set_start_ea, None, None)

# =============================================================================
# End class


class PyClassVftable(object):

    def __init__(self, vft_ea=None, cls=None):
        '''
        Represents a VFtable
        '''

        self.__vftable_ea = None
        if vft_ea != None:
            self.__vftable_ea = vft_ea

        self.__cls = None
        if cls != None:
            self.__cls = cls

        self.__vfuncs = {}
        self.__vtable_name = None

        self.__id = idaapi.BADADDR

        return

    def __str__(self):
        t = "Vftable  %s @ %x\n" % (self.__vtable_name, self.__vftable_ea)
        for k, v in self.__vfuncs.items():
            t += "\tvirtual %x @ %x\n" % (v.start_ea, k)
        return t

    def has_vfuncs(self):
        return not self.__vfuncs

    def get_vftable_ea(self):
        return self.__vftable_ea

    def get_vftable_name(self):
        return self.__vtable_name

    def get_cls(self):
        return self.__cls

    def set_vftable_ea(self, value):
        self.__vftable_ea = value

    def set_vftable_name(self, value):
        self.__vtable_name = value

    def set_cls(self, value):
        self.__cls = value

    def add_virtual_function(self, value, offset):
        self.__vfuncs[offset] = value

    def get_virtual_functions(self):
        return self.__vfuncs

    def set_id(self, value):
        self.__id = value
        return

    def get_id(self):
        return self.__id

    # Properties
    start_ea = property(get_vftable_ea, set_vftable_ea, None, None)
    vtable_name = property(get_vftable_name, set_vftable_name, None, None)
    cls = property(get_cls, set_cls, None, None)
    virtual_functions = property(get_virtual_functions, None, None, None)
    id = property(get_id, set_id, None, None)

# =============================================================================
# End class


class PyClassMember(object):
    '''
    Represents a class member
    '''

    def __init__(self, cls, mem_type, size):
        '''
        initialize class member
        '''

        self.__cls = cls
        self.__type_size = size
        self.set_type(mem_type, size)

        self.__member_name = None
        self.__class_mem = None

        self.__id = idaapi.BADADDR

        self.__is_parent = False

        self.__is_vftptr = False

        self.__offset = -1

        return

    def __str__(self):
        '''
        Return member as string
        '''
        return "%s (%x)" % (self.__member_name, self.__id)

    # property getters/setters

    def get_member_name(self):
        return self.__member_name

    def get_type(self):
        return self.__type

    def get_cls(self):
        return self.__cls

    def get_class_member(self):
        return self.__class_mem

    def set_class_member(self, value):
        self.__class_mem = value
        if self.__type_size == None and self.__class_mem != None:
            # ida_struct.get_struc_size(idc.GetStrucIdByName(self.value))
            self.__type_size = self.__class_mem.size()

        return

    def set_member_name(self, value):
        self.__member_name = value
        return

    def get_type_size(self):
        return self.__type_size

    def set_type(self, value, size):
        '''
        the ID for the struct member_type must be supplied
        the member_type can either by primitive :or a structure. If it is a struct :or
        string, then size is required
        '''

        if value == 'vftptr':
            self.__is_vftptr = True
            self.__type = ida_bytes.FF_DWORD

        elif value == 'struc':
            self.__type = ida_bytes.FF_STRUCT

        elif value == '':
            d = {1: ida_bytes.FF_BYTE,
                 2: ida_bytes.FF_WORD,
                 4: ida_bytes.FF_DWORD}
            self.__type = d.get(size, None)
            if self.__type is None:
                # XXX: We should really have a marker here so that in apply_member we create a
                # dummy type of the appropriate size.
                print("WARNING: Unknown type for size %d" % size)
                self.__type = ida_bytes.FF_BYTE
        else:
            print("WARNING: Unknown type %s" % value)
            assert (False)

        return

    def set_cls(self, value):
        self.__cls = value

    def set_id(self, value):
        self.__id = value
        return

    def set_isvftptr(self, value):
        self.__is_vftptr = value
        return

    def get_isvftptr(self):
        return self.__is_vftptr

    def get_id(self):
        return self.__id

    def get_size(self):
        return self.__size

    def set_size(self, value):
        self.__size = value
        return

    def get_isparent(self):
        return self.__is_parent

    def set_isparent(self, value):
        self.__is_parent = value

    def get_member_offset(self):
        return self.__offset

    def set_member_offset(self, value):
        self.__offset = value

    # Properties for class members

    member_name = property(get_member_name, set_member_name, None, None)
    is_parent = property(get_isparent, set_isparent, None, None)
    is_vftptr = property(get_isvftptr, set_isvftptr, None, None)
    member_type = property(get_type, set_type, None, None)
    class_member = property(get_class_member, set_class_member, None, None)
    cls = property(get_cls, set_cls, None, None)
    offset = property(get_member_offset, set_member_offset, None, None)
    size = property(get_type_size, None, None, None)
    id = property(get_id, set_id, None, None)

# =============================================================================
# End class


class PyClassStructure(object):
    '''
    Represents C++ class/structure
    '''

    def __init__(self, struc_def):
        '''
        Builds PyClassStructure
        '''
        self.__methods = []
        self.__vftables = {}

        # members and parents take up space
        self.__members = {}
        self.__parents = {}

        self.__vfcalls = []
        self.__vfcalls = []

        self.__id = idaapi.BADADDR   # This is set when applied

        self.__applied = False

        self.__size = struc_def['size']

        self.__demangled_name = struc_def['demangled_name']

        self.__name = struc_def['name']

        # set the IDA name
        self.__ida_name = sanitize_name(self.__name)

        if self.__demangled_name != "":
            self.__ida_name = sanitize_name(self.__demangled_name)

        return

    def __eq__(self, other):
        if other == None:
            return False
        return self.__id == other.__id and self.__name == other.__name

    def __str__(self):
        '''
        String representation of C++ class
        '''
        folks = ""
        for o, p in self.__parents.items():
            folks += "%s" % (p.ida_name)

        s = "class %s" % self.__ida_name
        if len(folks) > 0:
            s += " : %s" % folks

        s += " { \n"
        for m in self.__members:
            s += "\t%s @ %s\n" % (self.__members[m], m)

        for t in self.__methods:
            s += "\t%s\n" % t

        for v in self.__vftables:
            s += "\t%s\n" % str(v)

        return "%s}" % s

    # property getters/setters

    def get_best_name(self):
        dn = self.get_demangled_name()
        if dn is not None and dn != "":
            return dn
        else:
            return self.get_name()

    def get_ida_name(self):
        return self.__ida_name

    def set_ida_name(self, value):
        self.__ida_name = value
        return

    def get_demangled_name(self):
        return self.__demangled_name

    def set_demangled_name(self, value):
        self.__demangled_name = value
        return

    def get_name(self):
        return self.__name

    def set_name(self, value):
        self.__name = value
        return

    def add_member(self, member, off):
        '''
        Add the member, ensuring that the required parameters are specified
        '''
        member.cls = self
        self.__members[off] = member

        return

    def add_method(self, method):

        method.cls = self
        self.__methods.append(method)

        return

    def add_vftable(self, vft, vftptr):
        self.__vftables[vftptr] = vft
        return

    def add_parent(self, parent, off):
        self.__parents[off] = parent
        return

    def add_usage(self, usage):
        self.__vfcalls.append(usage)
        return

    def has_members(self):
        return not self.__members

    def has_methods(self):
        return not self.__methods

    def has_vftable(self):
        return not self.__vftables

    def get_members(self):
        return self.__members

    def get_methods(self):
        return self.__methods

    def get_vftables(self):
        return self.__vftables

    def set_vftables(self, value):
        self.__vftables = value

    def get_parents(self):
        return self.__parents

    def set_id(self, value):
        self.__id = value
        return

    def get_id(self):
        return self.__id

    def set_applied(self, value):
        self.__applied = value
        return

    def get_applied(self):
        return self.__applied

    def get_size(self):
        return self.__size

    def set_size(self, value):
        self.__size = value
        return

    def get_usages(self):
        return self.__vfcalls

    def set_usages(self, value):
        self.__vfcalls = value

    def get_vfcalls(self):
        return self.__vfcalls

    def set_vfcalls(self, value):
        self.__vfcalls = value

    # Properties of class structures

    demangled_name = property(
        get_demangled_name, set_demangled_name, None, None)
    ida_name = property(get_ida_name, set_ida_name, None, None)
    name = property(get_name, set_name, None, None)
    members = property(get_members, None, None, None)
    methods = property(get_methods, None, None, None)
    vftables = property(get_vftables, set_vftables, None, None)
    parents = property(get_parents, None, None, None)
    id = property(get_id, set_id, None, None)
    size = property(get_size, set_size, None, None)
    applied = property(get_applied, set_applied, None, None)
    usages = property(get_usages, set_usages, None, None)
    vfcalls = property(get_vfcalls, set_vfcalls, None, None)

# =============================================================================
# End class


class PyOOAnalyzerMethodTreeItem(QTreeWidgetItem):
    '''
    This is a method element in the class viewer tree
    '''

    def __init__(self, parent, meth):

        super(PyOOAnalyzerMethodTreeItem, self).__init__(parent)
        self.method = meth
        self.setText(0, self.method.get_best_method_name())


class PyOOAnalyzerMemberTreeItem(QTreeWidgetItem):
    '''
    This is a member element in the class viewer tree
    '''

    def __init__(self, parent, mem=None):

        super(PyOOAnalyzerMemberTreeItem, self).__init__(parent)
        self.member = mem
        self.setText(0, self.member.member_name)


class PyOOAnalyzerStructTreeItem(QTreeWidgetItem):
    '''
    This is a class element in the class viewer tree
    '''

    def __init__(self, parent, cls=None):

        super(PyOOAnalyzerStructTreeItem, self).__init__(parent)
        self.class_struct = cls
        self.setText(0, self.class_struct.get_best_name())


class PyOOAnalyzerExpForm(idaapi.PluginForm):
    '''
    This is the main class viewer form.
    '''

    def __init__(self):

        super(PyOOAnalyzerExpForm, self).__init__()

        self.__ooanalyzer = None

        self.cls_tree = None
        self.idbhook = None
        self.idphook = None
        self.__node = Netnode("$ cert.ooanalyzer")

        # the name of the IDA netnode that contains class information
        self.__NODE_NAME = "$OBJD"

        self.__visible = False

        return

    def start_ida_hooks(self):

        if self.idbhook == None:
            self.idbhook = PyOOAnalyzerIDBHooks(self.__ooanalyzer, self)
            self.idbhook.hook()

        if self.idphook == None:
            self.idphook = PyOOAnalyzerIDPHooks(self.__ooanalyzer, self)
            self.idphook.hook()

    def remove_ida_hooks(self):

        if self.idbhook != None:
            self.idbhook.unhook()
            self.idbhook = None

        if self.idphook != None:
            self.idphook.unhook()
            self.idphook = None

    def populate_class_list(self):
        '''
        Initialize the class viewer
        '''

        cls_list = self.__ooanalyzer.get_classes()
        for cls in cls_list:
            cls_entry = PyOOAnalyzerStructTreeItem(self.cls_tree, cls)

            cls_methods = QTreeWidgetItem(cls_entry)
            cls_methods.setText(0, "Methods")

            for m in cls.methods:
                PyOOAnalyzerMethodTreeItem(cls_methods, m)

            cls_members = QTreeWidgetItem(cls_entry)
            cls_members.setText(0, "Members")

            for off, mbr in cls.members.items():
                if mbr.is_parent == False:
                    PyOOAnalyzerMemberTreeItem(cls_members, mbr)

            if len(cls.parents) > 0:
                cls_parents = QTreeWidgetItem(cls_entry)
                cls_parents.setText(0, "Parents")
                for o, p in cls.parents.items():
                    # Parents are really special members that are class types
                    PyOOAnalyzerStructTreeItem(cls_parents, p)

        return

    def navigate(self, event):
        '''
        Enable 1-click navigation to class methods
        '''

        self.cls_tree.blockSignals(True)

        item = self.cls_tree.currentItem()
        if type(item) is PyOOAnalyzerMethodTreeItem:
            idaapi.jumpto(item.method.start_ea)

        self.cls_tree.blockSignals(False)

        return

    def update_class_method(self, cid, old_name, new_name):
        '''
        Update the class viewer in response to a method name change elsewhere in
        the IDB
        '''

        self.cls_tree.blockSignals(True)

        iterator = QTreeWidgetItemIterator(self.cls_tree,
                                           QTreeWidgetItemIterator.NotHidden)

        item = iterator.value()

        terminate = False
        while item and not terminate:
            if type(item) is PyOOAnalyzerStructTreeItem:
                if item.class_struct.id == cid:

                    # Have to check for because parents are stuct items,
                    # thus we may be in a parent
                    if item.childCount() > 0:
                        meth_group = item.child(0)  # should be methods
                        if meth_group != None:
                            if meth_group.text(0) == "Methods":

                                i = 0
                                while i < meth_group.childCount():
                                    meth_item = meth_group.child(i)
                                    if meth_item:

                                        if type(meth_item) is PyOOAnalyzerMethodTreeItem:
                                            if meth_item.text(0) == old_name:
                                                meth_item.setText(0, new_name)
                                                terminate = True
                                                break
                                    i += 1

            iterator += 1
            item = iterator.value()

        self.cls_tree.blockSignals(False)

    def update_class_member(self, cid, old_name, new_name):
        '''
        Update the class viewer in response to a member name change elsewhere in
        the IDB
        '''

        self.cls_tree.blockSignals(True)

        iterator = QTreeWidgetItemIterator(self.cls_tree,
                                           QTreeWidgetItemIterator.NotHidden)

        item = iterator.value()
        terminate = False
        while item and not terminate:
            if type(item) is PyOOAnalyzerStructTreeItem:
                if item.class_struct.id == cid:

                    # Have to check for because parents are stuct items,
                    # thus we may be in a parent
                    if item.childCount() > 0:
                        mbr_group = item.child(1)  # should be members
                        if mbr_group != None:

                            if mbr_group.text(0) == "Members":
                                i = 0
                                while i < mbr_group.childCount():
                                    mem_item = mbr_group.child(i)
                                    if mem_item:
                                        if type(mem_item) is PyOOAnalyzerMemberTreeItem:
                                            if mem_item.text(0) == old_name:
                                                mem_item.setText(0, new_name)
                                                terminate = True
                                                break
                                    i += 1

            iterator += 1
            item = iterator.value()

        self.cls_tree.blockSignals(False)

    def update_class(self, old_name, new_name):
        '''
        Update the class viewer in response to a class name change elsewhere in
        the IDB
        '''

        self.cls_tree.blockSignals(True)

        iterator = QTreeWidgetItemIterator(self.cls_tree,
                                           QTreeWidgetItemIterator.NotHidden)

        item = iterator.value()
        while item:
            if item.text(0) == old_name:
                item.setText(0, new_name)
            else:
                # rename parents of this item
                if item.childCount() == 3:
                    par_group = item.child(3)  # should be members
                    if par_group != None:
                        if par_group.text(0) == "Parents":
                            i = 0
                            while i < par_group.childCount():
                                par_item = mbr_group.child(i)
                                if par_item:
                                    if par_item.text(0) == old_name:
                                        par_item.setText(0, new_name)
            iterator += 1
            item = iterator.value()

        self.cls_tree.blockSignals(False)

        return

    def __edit_member_from_class_viewer(self, item):
        '''
        Edit the class member from the Class Viewer
        '''
        member = item.member
        old_member_name = member.member_name
        new_member_name = ida_kernwin.ask_str(
            old_member_name, 0, "Enter new member name")

        if new_member_name == None:
            self.cls_tree.blockSignals(False)
            return

        # cid = ida_struct.get_struc_id(member.member_name
        cls_members = idautils.StructMembers(member.cls.id)

        for [off, n, s] in cls_members:
            if n == member.member_name:

                if idc.set_member_name(member.cls.id, off, str(new_member_name)) != 0:
                    item.setText(0, str(new_member_name))

                    cls_list = self.__ooanalyzer.get_classes()
                    for c in cls_list:
                        for moff, mem in c.members.items():
                            if mem.member_name == old_member_name:
                                mem.member_name = new_member_name
                else:
                    ida_kernwin.warning(
                        "Cannot rename member to %s, the name already exists or is malformed!" % new_member_name)

        return

    def __edit_method_from_class_viewer(self, item):
        '''
        Handle method change request initiated from Class Viewer
        '''
        method = item.method
        old_method_name = method.method_name
        new_method_name = ida_kernwin.ask_str(
            old_method_name, 0, "Enter new method name")

        if new_method_name == None:
            self.cls_tree.blockSignals(False)
            return

        cls = item.parent().parent().class_struct

        m = next(m for m in cls.methods if m.start_ea == method.start_ea)
        m.method_name = new_method_name
        m.userdef_name = True
        # Rename the function
        self.__ooanalyzer.apply_method(m)
        item.setText(0, m.get_best_method_name())

        # XXX: Rename all vftable structures.
        # # the method is a virtual function
        # if m.is_virtual == True:
        #     cmt = "virtual %s" % cmt

        #     # rename virtual function in vftable structure
        #     for off, vt in c.vftables.items():
        #         for vfoff, vf in vt.virtual_functions.items():
        #             if vf.start_ea == m.start_ea:

        #                 global ignore_renamed
        #                 ignore_renamed = True
        #                 idc.set_member_name(
        #                     vt.id, 4*vfoff, new_method_name)
        #                 ignore_renamed = False

        #                 vf.method_name = new_method_name
        #                 break

        return

    def __edit_class_from_class_viewer(self, item):
        '''
        Handle class name change request initiated from Class Viewer
        '''

        old_name = item.class_struct.ida_name
        cid = idaapi.get_struc_id(str(old_name))

        if (cid is not None) and (cid != 0) and (cid != idaapi.BADADDR):

            new_name = ida_kernwin.ask_str(
                old_name, 0, "Enter new class name:")

            if new_name == None:
                self.cls_tree.blockSignals(False)
                return

            if ida_struct.set_struc_name(cid, new_name) != 0:
                item.setText(0, str(new_name))

                cls_list = self.__ooanalyzer.get_classes()
                for c in cls_list:
                    if c.ida_name == old_name:
                        c.ida_name = new_name

            else:
                ida_kernwin.warning(
                    "Cannot rename class to %s, the name already exists!" % new_name)
        else:
            ida_kernwin.warning("Cannot rename class before it is applied")

        return

    def edit_class_item(self, event):

        self.cls_tree.blockSignals(True)

        item = self.cls_tree.currentItem()

        if type(item) is PyOOAnalyzerMemberTreeItem:
            self.__edit_member_from_class_viewer(item)

        elif type(item) is PyOOAnalyzerMethodTreeItem:
            self.__edit_method_from_class_viewer(item)

        elif type(item) is PyOOAnalyzerStructTreeItem:
            self.__edit_class_from_class_viewer(item)

        self.cls_tree.blockSignals(False)

    def OnCreate(self, form):

        # Get parent widget

        self.parent = self.FormToPyQtWidget(form)

        # Create cls_tree control
        self.cls_tree = QTreeWidget()
        headerItem = QTreeWidgetItem()
        headerItem.setText(0, "Class")
        headerItem.setText(1, "State")

        self.cls_tree.setHeaderItem(headerItem)
        self.cls_tree.setWindowTitle("OOAnalzyer Results")
        self.cls_tree.setColumnWidth(0, 200)

        self.cls_tree.itemSelectionChanged.connect(
            lambda: self.navigate(self.cls_tree.currentItem()))

        # install the context menu
        self.cls_tree.setContextMenuPolicy(Qt.CustomContextMenu)
        self.cls_tree.customContextMenuRequested.connect(self.open_menu)

        # Create layout
        layout = QVBoxLayout()
        layout.addWidget(self.cls_tree)

        self.populate_class_list()
        # Populate PluginForm
        self.parent.setLayout(layout)

        applied_cls = self.__ooanalyzer.get_applied_classes()
        for c in applied_cls:
            self.__mark_applied(c)

        return

    def open_menu(self, position):

        menu = QMenu()

        # can only apply classes
        applyAction = None
        gotoDefAction = None
        renameAction = None

        item = self.cls_tree.currentItem()
        if type(item) is PyOOAnalyzerStructTreeItem:
            cls = item.class_struct
            if cls.applied:
                gotoDefAction = menu.addAction("Open Definition")
                renameAction = menu.addAction("Rename")
            else:
                applyAction = menu.addAction("Apply")
        elif type(item) is PyOOAnalyzerMethodTreeItem:
            if item.method.cls.applied:
                renameAction = menu.addAction("Rename")
        elif type(item) is PyOOAnalyzerMemberTreeItem:
            if item.member.cls.applied:
                renameAction = menu.addAction("Rename")

        action = menu.exec_(QCursor.pos())

        if applyAction is not None and action == applyAction:
            self.apply_action()
        elif gotoDefAction is not None and action == gotoDefAction:
            self.goto_def_action()
        elif renameAction is not None and action == renameAction:
            self.rename_action()

        return

    def apply_class_item(self, item):
        '''
        Apply a class
        '''
        if type(item) is PyOOAnalyzerStructTreeItem:

            applied_classes = self.__ooanalyzer.get_applied_classes()

            if self.__ooanalyzer.apply_struct(item.class_struct):

                new_applied_classes = self.__ooanalyzer.get_applied_classes()

                # mark all the newly applied classes
                delta = [
                    x for x in new_applied_classes if x not in applied_classes]

                for c in delta:
                    self.__mark_applied(c)
        else:
            print("Not applying because not struct")

        return

    def __mark_applied(self, applied_class):
        '''
        Mark the classes that have been applied
        '''
        root = self.cls_tree.invisibleRootItem()
        child_count = root.childCount()

        for i in range(child_count):
            cur_item = root.child(i)
            if type(cur_item) is PyOOAnalyzerStructTreeItem:
                if applied_class.ida_name == cur_item.class_struct.ida_name:
                    cur_item.setBackground(
                        0, QBrush(Qt.green, Qt.Dense7Pattern))
                    font = cur_item.font(0)
                    font.setBold(True)
                    cur_item.setFont(0, font)
                    cur_item.setText(1, "Applied")

    def apply_action(self):
        '''
        Execute the apply action
        '''
        item = self.cls_tree.currentItem()
        self.apply_class_item(item)

    def goto_def_action(self):
        '''
        Execute the goto definition action
        '''
        item = self.cls_tree.currentItem()
        cls = item.class_struct
        if cls.applied == True:
            ida_kernwin.open_structs_window(item.class_struct.id)
        else:
            ida_kernwin.warning(
                "You must apply class '%s' before you can open its definition" % item.class_struct.ida_name)

    def rename_action(self):
        '''
        Execute the rename action
        '''
        item = self.cls_tree.currentItem()
        self.edit_class_item(item)

    def __save_to_idb(self):
        '''
        Save the JSON to a netnode
        '''
               
        self.__node[OOA_JSON] = self.__ooanalyzer.get_json_data()
        self.__node[APPLIED_CLASS_NAMES] = [
            cls.ida_name for cls in self.__ooanalyzer.get_applied_classes()]

        return

    def __load_from_idb(self):
        '''
        Load previously applied JSON
        '''

        if OOA_JSON in self.__node and APPLIED_CLASS_NAMES in self.__node:
            return [self.__node[OOA_JSON], self.__node[APPLIED_CLASS_NAMES]]
        elif OOA_JSON in self.__node:
            return [self.__node[OOA_JSON], None]
        elif APPLIED_CLASS_NAMES in self.__node:
            return [None, self.__node[APPLIED_CLASS_NAMES]]

        return [None, None]

    def __load_from_json_file(self, json_file):
        '''
        Parse the JSON file
        '''
        if json_file != None:

            if self.__ooanalyzer.set_json_file(json_file):
                print("Opened JSON file %s" % json_file)
            else:
                ida_kernwin.warning("Could not open %s" % json_file)
                return None, 0, 0, 0

        if self.__ooanalyzer.is_parsed() == False:  # not parsed yet
            result, msg = self.__ooanalyzer.parse()
            print("Parsed %s %s" % (result, msg))
            if result == False:
                ida_kernwin.warning("Could not parse JSON: %s" % msg)
                return None, 0, 0, 0

        parse_results = self.__ooanalyzer.get_parse_results()
        if parse_results == None:
            return None, 0, 0, 0

        num_classes = 0
        num_vcalls = 0
        num_usages = 0

        if "NumClasses" in parse_results:
            num_classes = parse_results["NumClasses"]
        if "NumVcalls" in parse_results:
            num_vcalls = parse_results["NumVcalls"]
        if "NumUsages" in parse_results:
            num_usages = parse_results["NumUsages"]

        res_str = """
Successfully parsed JSON file: "%s".

The following C++ constructs are ready to apply:

  * %s class structures
  * %s object usages
  * %s virtual function calls

Press \"Yes\" to apply these items to the IDB. Press no to apply them manually (Note that virtual function calls will be resolved automatically)
""" % (self.__ooanalyzer.get_json_filename(), num_classes, num_usages, num_vcalls)

        return res_str, num_classes, num_vcalls, num_usages

    def OnClose(self, form):
        """
        Called when the plugin form is closed
        """

        global clsExpForm

        self.__save_to_idb()

        clsExpForm.remove_ida_hooks()

        self.__visible = False

        clsExpForm = None

        return

    def __setup_form(self):
        '''
        set up form data
        '''

        [json, applied_classes] = self.__load_from_idb()
        if json:
            print("OOAnalyzer class information found in IDB, opening Class Viewer")
          
            # mark previously applied classes
            self.__ooanalyzer = PyOOAnalyzer()
            self.__ooanalyzer.set_json_data(json)
            self.__ooanalyzer.parse()
            if applied_classes:

                # Find the intersection
                intersection = [
                    c for c in self.__ooanalyzer.get_classes() if c.ida_name in applied_classes]

                for c in intersection:
                    self.__ooanalyzer.apply_class(c)

        else:

            # load the class spec from a JSON file

            json_file = ida_kernwin.ask_file(
                0, "*.json", "Open JSON C++ specification file")
            if not json_file:
                return False

            self.__ooanalyzer = PyOOAnalyzer()

            # if everything applies OK, then show the form
            parse_results, nc, nvc, nu = self.__load_from_json_file(json_file)

            if parse_results != None:
                form_msg = """CERT OOAnalyzer v%.1f
%s
""" % (PLUGIN_VERSION, parse_results)

                do_apply = ida_kernwin.ask_yn(ida_kernwin.ASKBTN_YES, form_msg)

                if do_apply == ida_kernwin.ASKBTN_YES:  # yes
                    print("Applying all results")

                    self.__ooanalyzer.apply_all_structs()

                # apply virtual function calls for yes or no
                if do_apply == ida_kernwin.ASKBTN_YES or do_apply == ida_kernwin.ASKBTN_NO:
                    print("Applying virtual function calls")
                    self.__ooanalyzer.apply_vfcalls()
                    print("Applied virtual function calls")

                else:  # cancel
                    print("User cancelled.  Not applying anything")
                    return False

            else:
                # ida_kernwin.warning(
                #     "No C++ data structures found in '%s'" % json_file)
                return False

        return True

    def Show(self):

        if not self.__visible:
            setup_form = False
            try:
                setup_form = self.__setup_form
            except Exception as e:
                ida_kernwin.warning("An exception occurred: %s" % e)
            if self.__setup_form():

                if self.idbhook == None and self.idphook == None:
                    self.start_ida_hooks()

                self.__visible = True
                return idaapi.PluginForm.Show(self, "OOAnalyzer Class Viewer", options=idaapi.PluginForm.WOPN_PERSIST)
            else:
                return -1

    def is_visible(self):
        return self.__visible


class PyOOAnalyzerIDPHooks(ida_idp.IDP_Hooks):

    def __init__(self, objd, frm):

        ida_idp.IDP_Hooks.__init__(self)

        self.__ooanalyzer = objd
        self.__form = frm

        return

    def renamed(self, *args):
        '''
        Something in the database has been renamed
        '''
        ea = args[0]
        new_name = args[1]
        local_name = args[2]
        global ignore_renamed
        if ignore_renamed:
            return ida_idp.IDP_Hooks.renamed(self, ea, new_name, local_name)

        # a function has been renamed
        if idc.get_func_name(ea) != "":
            cls_model = self.__ooanalyzer.get_classes()
            for c in cls_model:
                for m in c.methods:
                    if m.start_ea == ea:

                        self.__form.update_class_method(
                            c.id, m.method_name, new_name)
                        m.method_name = new_name

                        if m.is_virtual:

                            # rename function name in vftable structure
                            for off, vt in c.vftables.items():
                                for vfoff, vf in vt.virtual_functions.items():

                                    if vf.start_ea == m.start_ea:

                                        ignore_renamed = True
                                        idc.set_member_name(
                                            vt.id, 4*vfoff, new_name)

                                        ignore_renamed = False

                                        vf.method_name = new_name
                                        break

                        return ida_idp.IDP_Hooks.renamed(self, ea, new_name, local_name)

        return ida_idp.IDP_Hooks.renamed(self, ea, new_name, local_name)


class PyOOAnalyzerIDBHooks(idaapi.IDB_Hooks):

    def __init__(self, objd, frm):

        self.form = frm
        self.__ooanalyzer = objd

        idaapi.IDB_Hooks.__init__(self)

    def struc_member_renamed(self, *args):
        '''
        A structure member has been renamed
        '''
        global ignore_renamed
        if ignore_renamed:
            return 0

        struc = args[0]
        mbr = args[1]

        class_model = self.__ooanalyzer.get_classes()
        for c in class_model:
            if c.id == struc.id:

                changed_name = idc.get_member_name(c.id, int(mbr.soff))

                for off, m in c.members.items():
                    if off == mbr.soff:
                        if m:
                            self.form.update_class_member(
                                c.id, m.member_name, changed_name)
                            m.member_name = changed_name
                            return 0

        # Update virtual functions if they are changed
        for c in class_model:
            for off, vt in c.vftables.items():
                if struc.id == vt.id:

                    changed_name = idc.get_member_name(vt.id, int(mbr.soff))

                    for vfoff, vf in vt.virtual_functions.items():

                        # this has to be scaled
                        if (vfoff*4) == mbr.soff:
                            if vt:
                                self.form.update_class_method(
                                    c.id, vf.method_name, changed_name)
                                vf.method_name = changed_name

                                ignore_renamed = True
                                idc.set_name(
                                    vf.start_ea, changed_name, ida_name.SN_CHECK)
                                ignore_renamed = False

                                return 0
        return 0

    def struc_renamed(self, *args):
        '''
        A structure was renamed in the IDB
        '''
        renamed_struc = args[0]

        class_model = self.__ooanalyzer.get_classes()
        for c in class_model:
            if c.id == renamed_struc.id:

                new_name = ida_struct.get_struc_name(renamed_struc.id)
                if new_name != None:

                    self.form.update_class(c.ida_name, new_name)

                    c.ida_name = new_name
                    for off, m in c.members.items():
                        m.ida_name = new_name

                    for mtd in c.methods:
                        cmt = "%s::%s" % (c.ida_name, mtd.method_name)
                        if mtd.is_virtual == True:
                            cmt = "virtual %s" % cmt
                        if mtd.is_ctor == True:
                            cmt += " (constructor)"
                        if mtd.is_dtor == True:
                            cmt += " (destructor)"
                        idc.set_func_cmt(mtd.start_ea, cmt, 1)

                    for off, v in c.vftables.items():
                        v.cls.ida_name = new_name

                        new_vft_name = generate_vftable_name(
                            self.__ooanalyzer.get_classes(), v, c, off)

                        idc.set_name(v.start_ea, "%s_%x" % (
                            new_vft_name, v.start_ea), ida_name.SN_CHECK)
                        ida_struct.set_struc_name(v.id, new_vft_name)
                        ida_struct.set_struc_cmt(
                            v.id, v.cls.ida_name + " virtual function table", 1)

                    break

        return 0

# =============================================================================
# This is the entry point for the plugin


class PyOOAnalyzer_plugin(idaapi.plugin_t):
    '''
    Main plugin class
    '''

    flags = idaapi.PLUGIN_UNL

    comment = "IDA OOAnalyzer plugin"
    help = "Email jsg@cert.org"
    wanted_name = "OOAnalyzer"
    wanted_hotkey = "F3"

    def init(self):
        '''
        Create the Plugin
        '''

        return idaapi.PLUGIN_OK

    def run(self, arg):
        '''
        Run the plugin
        '''

        try:
            # is the Class Viewer already open?
            class_viewer = ida_kernwin.find_widget("OOAnalyzer Class Viewer")
            if class_viewer != None:
                ida_kernwin.warning(
                    "OOAnalyzer class information has already been imported")
                return

            global clsExpForm
            global ignore_renamed

            ignore_renamed = False
            clsExpForm = None

            print("\n\nOOAnalyzer Plugin version %s loaded" % PLUGIN_VERSION)

            if clsExpForm is None:
                clsExpForm = PyOOAnalyzerExpForm()
                clsExpForm.Show()

            idaapi.set_dock_pos('OOAnalyzer Class Viewer',
                                'Functions window', idaapi.DP_BOTTOM)
        except Exception as e:
            traceback.print_exc()
            print(e)

        return

    def term(self):
        print("OOAnalyzer Plugin version %s done importing" % PLUGIN_VERSION)
        return


def PLUGIN_ENTRY():
    try:
        return PyOOAnalyzer_plugin()
    except Exception as e:
        traceback.print_exc()
        print(e)
    return None

# Local Variables:
# mode: python
# fill-column: 95
# comment-column: 0
# python-indent-offset: 4
# End:


repository: nihilus/idapathfinder
file: plugins/idapathfinder.py
wanted_name: PathFinder
created_at: 2014-08-31 11:46:51+00:00
pushed_at: 2014-08-31 11:47:43+00:00
url:
https://github.com/nihilus/idapathfinder/blob/ab2a87871c8bb57dd98faea1247bdfe512637696/plugins/idapathfinder.py

content:
import idc
import idaapi
import idautils
import pathfinder
import time

class idapathfinder_t(idaapi.plugin_t):

	flags = 0
	comment = ''
	help = ''
	wanted_name = 'PathFinder'
	wanted_hotkey = ''

	def init(self):
		ui_path = "View/Graphs/"
		self.menu_contexts = []
		self.graph = None

		#self.menu_contexts.append(idaapi.add_menu_item(ui_path,
		#						"Find code paths to the current function block",
		#						"Alt-7",
		#						0,
		#						self.FindBlockPaths,
		#						(None,)))
		self.menu_contexts.append(idaapi.add_menu_item(ui_path,
								"Find function path(s) to here",
								"Alt-6",
								0,
								self.FindPathsFromMany,
								(None,)))
		#self.menu_contexts.append(idaapi.add_menu_item(ui_path,
		#						"Find paths to here from a single function",
		#						"Alt-7",
		#						0,
		#						self.FindPathsFromSingle,
		#						(None,)))
		self.menu_contexts.append(idaapi.add_menu_item(ui_path, 
								"Find function path(s) from here", 
								"Alt-5", 
								0, 
								self.FindPathsToMany, 
								(None,)))
		#self.menu_contexts.append(idaapi.add_menu_item(ui_path, 
		#						"Find paths from here to a single function", 
		#						"Alt-5", 
		#						0, 
		#						self.FindPathsToSingle, 
		#						(None,)))
		return idaapi.PLUGIN_KEEP

	def term(self):
		for context in self.menu_contexts:
			idaapi.del_menu_item(context)
		return None
	
	def run(self, arg):
		self.FindPathsToSingle()

	def _current_function(self):
		return idaapi.get_func(ScreenEA()).startEA

	def _find_and_plot_paths(self, sources, targets, pfc=pathfinder.FunctionPathFinder):
		results = []

		for target in targets:
			pf = pfc(target)
			for source in sources:
				s = time.time()
				r = pf.paths_from(source)
				e = time.time()
				#print "paths_from took %f seconds." % (e-s)

				if r:
					results += r
				else:
					name = idc.Name(target)
					if not name:
						name = "0x%X" % target
					print "No paths found to", name

		if results:
			# Be sure to close any previous graph before creating a new one.
			# Failure to do so may crash IDA.
			try:
				self.graph.Close()
			except:
				pass

			self.graph = pathfinder.PathFinderGraph(results, 'Path Graph', pf.colorize)
			self.graph.Show()

	def _get_user_selected_functions(self, many=False):
		functions = []
		ea = idc.ScreenEA()
		try:
			current_function = idc.GetFunctionAttr(ea, idc.FUNCATTR_START)
		except:
			current_function = None

		while True:
			function = idc.ChooseFunction("Select a function and click 'OK' until all functions have been selected. When finished, click 'Cancel' to display the graph.")
			# ChooseFunction automatically jumps to the selected function
			# if the enter key is pressed instead of clicking 'OK'. Annoying.
			if idc.ScreenEA() != ea:
				idc.Jump(ea)

			if not function or function == idc.BADADDR or function == current_function:
				break
			elif function not in functions:
				functions.append(function)

			if not many:
				break

		return functions
			
	def FindPathsToSingle(self, arg):
		source = self._current_function()

		if source:
			targets = self._get_user_selected_functions()
			if targets:
				print source, targets
				self._find_and_plot_paths([source], targets)

	def FindPathsToMany(self, arg):
		source = self._current_function()

		if source:
			targets = self._get_user_selected_functions(many=True)
			if targets:
				self._find_and_plot_paths([source], targets)

	def FindPathsFromSingle(self, arg):
		target = self._current_function()

		if target:
			sources = self._get_user_selected_functions()
			if sources:
				self._find_and_plot_paths(sources, [target])

	def FindPathsFromMany(self, arg):
		target = self._current_function()

		if target:
			sources = self._get_user_selected_functions(many=True)
			if sources:
				self._find_and_plot_paths(sources, [target])

	def FindBlockPaths(self, arg):
		target = idc.ScreenEA()
		source = idaapi.get_func(idc.ScreenEA())

		if source:
			self._find_and_plot_paths([source.startEA], [target], pfc=pathfinder.BlockPathFinder)
		else:
			print "Block graph error: The location must be part of a function!"

def PLUGIN_ENTRY():
	return idapathfinder_t()


repository: xerub/idastuff
file: hnight.py
description: IDA Pro/Hexrays plugins
url: https://github.com/xerub/idastuff/blob/62b720c1b5f60114d89a8a214177716ce3266b07/hnight.py
wanted_name: HexNight
comment: show symbolic names for ARM sysregs in Pseudocode-View
created_at: 2017-04-05 22:24:02+00:00
pushed_at: 2018-10-02 00:43:46+00:00
content:
# Print ARM sysregs using symbolic names.
#
# Copyright (c) 2017 xerub.  All rights reserved

import idautils
import idaapi
import idc

import traceback

hexnight_cb_info = None
hexnight_cb = None

# generated from capstone/arch/AArch64/AArch64BaseInfo.h
regs64 = {
    # Apple specific
    0xc780 : "HID0",
    0xc781 : "EHID0",
    0xc788 : "HID1",
    0xc790 : "HID2",
    0xc798 : "HID3",
    0xc799 : "EHID3",
    0xc7a0 : "HID4",
    0xc7a8 : "HID5",
    0xc7a9 : "EHID5",
    0xc7b0 : "HID6",
    0xc7b8 : "HID7",
    0xc7c0 : "HID8",
    0xc7c8 : "HID9",
    0xc7d0 : "HID10",
    0xc7e8 : "HID11",
    0xc7d8 : "HID11",
    0xc7d9 : "EHID11",
    0xefa0 : "CYC_CFG",
    0xefb0 : "ACC_OVRD",
    0xefa8 : "CYC_OVRD",
    0xdf80 : "LSU_ERR_STS",
    0xdf90 : "E_LSU_ERR_STS",
    0xdf88 : "LSU_ERR_CTL",
    0xf780 : "MMU_ERR_STS",
    0xf790 : "E_MMU_ERR_STS",
    0xdfc0 : "L2C_ERR_STS",
    0xdfc8 : "L2C_ERR_ADR",
    0xdfd0 : "L2C_ERR_INF",
    0xe784 : "MIGSTS_EL1",
    0xe793 : "KTRR_LOWER_EL1",
    0xe794 : "KTRR_UPPER_EL1",
    0xe792 : "KTRR_LOCK_EL1",
    # end of Apple specific
                                  # Op0 Op1  CRn   CRm   Op2
    0x9808 : "MDCCSR_EL0",        # 10  011  0000  0001  000
    0x9828 : "DBGDTRRX_EL0",      # 10  011  0000  0101  000
    0x8080 : "MDRAR_EL1",         # 10  000  0001  0000  000
    0x808c : "OSLSR_EL1",         # 10  000  0001  0001  100
    0x83f6 : "DBGAUTHSTATUS_EL1", # 10  000  0111  1110  110
    0xdce6 : "PMCEID0_EL0",       # 11  011  1001  1100  110
    0xdce7 : "PMCEID1_EL0",       # 11  011  1001  1100  111
    0xc000 : "MIDR_EL1",          # 11  000  0000  0000  000
    0xc800 : "CCSIDR_EL1",        # 11  001  0000  0000  000
    0xc801 : "CLIDR_EL1",         # 11  001  0000  0000  001
    0xd801 : "CTR_EL0",           # 11  011  0000  0000  001
    0xc005 : "MPIDR_EL1",         # 11  000  0000  0000  101
    0xc006 : "REVIDR_EL1",        # 11  000  0000  0000  110
    0xc807 : "AIDR_EL1",          # 11  001  0000  0000  111
    0xd807 : "DCZID_EL0",         # 11  011  0000  0000  111
    0xc008 : "ID_PFR0_EL1",       # 11  000  0000  0001  000
    0xc009 : "ID_PFR1_EL1",       # 11  000  0000  0001  001
    0xc00a : "ID_DFR0_EL1",       # 11  000  0000  0001  010
    0xc00b : "ID_AFR0_EL1",       # 11  000  0000  0001  011
    0xc00c : "ID_MMFR0_EL1",      # 11  000  0000  0001  100
    0xc00d : "ID_MMFR1_EL1",      # 11  000  0000  0001  101
    0xc00e : "ID_MMFR2_EL1",      # 11  000  0000  0001  110
    0xc00f : "ID_MMFR3_EL1",      # 11  000  0000  0001  111
    0xc010 : "ID_ISAR0_EL1",      # 11  000  0000  0010  000
    0xc011 : "ID_ISAR1_EL1",      # 11  000  0000  0010  001
    0xc012 : "ID_ISAR2_EL1",      # 11  000  0000  0010  010
    0xc013 : "ID_ISAR3_EL1",      # 11  000  0000  0010  011
    0xc014 : "ID_ISAR4_EL1",      # 11  000  0000  0010  100
    0xc015 : "ID_ISAR5_EL1",      # 11  000  0000  0010  101
    0xc020 : "ID_A64PFR0_EL1",    # 11  000  0000  0100  000
    0xc021 : "ID_A64PFR1_EL1",    # 11  000  0000  0100  001
    0xc028 : "ID_A64DFR0_EL1",    # 11  000  0000  0101  000
    0xc029 : "ID_A64DFR1_EL1",    # 11  000  0000  0101  001
    0xc02c : "ID_A64AFR0_EL1",    # 11  000  0000  0101  100
    0xc02d : "ID_A64AFR1_EL1",    # 11  000  0000  0101  101
    0xc030 : "ID_A64ISAR0_EL1",   # 11  000  0000  0110  000
    0xc031 : "ID_A64ISAR1_EL1",   # 11  000  0000  0110  001
    0xc038 : "ID_A64MMFR0_EL1",   # 11  000  0000  0111  000
    0xc039 : "ID_A64MMFR1_EL1",   # 11  000  0000  0111  001
    0xc018 : "MVFR0_EL1",         # 11  000  0000  0011  000
    0xc019 : "MVFR1_EL1",         # 11  000  0000  0011  001
    0xc01a : "MVFR2_EL1",         # 11  000  0000  0011  010
    0xc601 : "RVBAR_EL1",         # 11  000  1100  0000  001
    0xe601 : "RVBAR_EL2",         # 11  100  1100  0000  001
    0xf601 : "RVBAR_EL3",         # 11  110  1100  0000  001
    0xc608 : "ISR_EL1",           # 11  000  1100  0001  000
    0xdf01 : "CNTPCT_EL0",        # 11  011  1110  0000  001
    0xdf02 : "CNTVCT_EL0",        # 11  011  1110  0000  010
    0x8818 : "TRCSTATR",          # 10  001  0000  0011  000
    0x8806 : "TRCIDR8",           # 10  001  0000  0000  110
    0x880e : "TRCIDR9",           # 10  001  0000  0001  110
    0x8816 : "TRCIDR10",          # 10  001  0000  0010  110
    0x881e : "TRCIDR11",          # 10  001  0000  0011  110
    0x8826 : "TRCIDR12",          # 10  001  0000  0100  110
    0x882e : "TRCIDR13",          # 10  001  0000  0101  110
    0x8847 : "TRCIDR0",           # 10  001  0000  1000  111
    0x884f : "TRCIDR1",           # 10  001  0000  1001  111
    0x8857 : "TRCIDR2",           # 10  001  0000  1010  111
    0x885f : "TRCIDR3",           # 10  001  0000  1011  111
    0x8867 : "TRCIDR4",           # 10  001  0000  1100  111
    0x886f : "TRCIDR5",           # 10  001  0000  1101  111
    0x8877 : "TRCIDR6",           # 10  001  0000  1110  111
    0x887f : "TRCIDR7",           # 10  001  0000  1111  111
    0x888c : "TRCOSLSR",          # 10  001  0001  0001  100
    0x88ac : "TRCPDSR",           # 10  001  0001  0101  100
    0x8bd6 : "TRCDEVAFF0",        # 10  001  0111  1010  110
    0x8bde : "TRCDEVAFF1",        # 10  001  0111  1011  110
    0x8bee : "TRCLSR",            # 10  001  0111  1101  110
    0x8bf6 : "TRCAUTHSTATUS",     # 10  001  0111  1110  110
    0x8bfe : "TRCDEVARCH",        # 10  001  0111  1111  110
    0x8b97 : "TRCDEVID",          # 10  001  0111  0010  111
    0x8b9f : "TRCDEVTYPE",        # 10  001  0111  0011  111
    0x8ba7 : "TRCPIDR4",          # 10  001  0111  0100  111
    0x8baf : "TRCPIDR5",          # 10  001  0111  0101  111
    0x8bb7 : "TRCPIDR6",          # 10  001  0111  0110  111
    0x8bbf : "TRCPIDR7",          # 10  001  0111  0111  111
    0x8bc7 : "TRCPIDR0",          # 10  001  0111  1000  111
    0x8bcf : "TRCPIDR1",          # 10  001  0111  1001  111
    0x8bd7 : "TRCPIDR2",          # 10  001  0111  1010  111
    0x8bdf : "TRCPIDR3",          # 10  001  0111  1011  111
    0x8be7 : "TRCCIDR0",          # 10  001  0111  1100  111
    0x8bef : "TRCCIDR1",          # 10  001  0111  1101  111
    0x8bf7 : "TRCCIDR2",          # 10  001  0111  1110  111
    0x8bff : "TRCCIDR3",          # 10  001  0111  1111  111
    0xc660 : "ICC_IAR1_EL1",      # 11  000  1100  1100  000
    0xc640 : "ICC_IAR0_EL1",      # 11  000  1100  1000  000
    0xc662 : "ICC_HPPIR1_EL1",    # 11  000  1100  1100  010
    0xc642 : "ICC_HPPIR0_EL1",    # 11  000  1100  1000  010
    0xc65b : "ICC_RPR_EL1",       # 11  000  1100  1011  011
    0xe659 : "ICH_VTR_EL2",       # 11  100  1100  1011  001
    0xe65b : "ICH_EISR_EL2",      # 11  100  1100  1011  011
    0xe65d : "ICH_ELSR_EL2",      # 11  100  1100  1011  101
    0x9828 : "DBGDTRTX_EL0",      # 10  011  0000  0101  000
    0x8084 : "OSLAR_EL1",         # 10  000  0001  0000  100
    0xdce4 : "PMSWINC_EL0",       # 11  011  1001  1100  100
    0x8884 : "TRCOSLAR",          # 10  001  0001  0000  100
    0x8be6 : "TRCLAR",            # 10  001  0111  1100  110
    0xc661 : "ICC_EOIR1_EL1",     # 11  000  1100  1100  001
    0xc641 : "ICC_EOIR0_EL1",     # 11  000  1100  1000  001
    0xc659 : "ICC_DIR_EL1",       # 11  000  1100  1011  001
    0xc65d : "ICC_SGI1R_EL1",     # 11  000  1100  1011  101
    0xc65e : "ICC_ASGI1R_EL1",    # 11  000  1100  1011  110
    0xc65f : "ICC_SGI0R_EL1",     # 11  000  1100  1011  111
    0x8002 : "OSDTRRX_EL1",       # 10  000  0000  0000  010
    0x801a : "OSDTRTX_EL1",       # 10  000  0000  0011  010
    0x9000 : "TEECR32_EL1",       # 10  010  0000  0000  000
    0x8010 : "MDCCINT_EL1",       # 10  000  0000  0010  000
    0x8012 : "MDSCR_EL1",         # 10  000  0000  0010  010
    0x9820 : "DBGDTR_EL0",        # 10  011  0000  0100  000
    0x8032 : "OSECCR_EL1",        # 10  000  0000  0110  010
    0xa038 : "DBGVCR32_EL2",      # 10  100  0000  0111  000
    0x8004 : "DBGBVR0_EL1",       # 10  000  0000  0000  100
    0x800c : "DBGBVR1_EL1",       # 10  000  0000  0001  100
    0x8014 : "DBGBVR2_EL1",       # 10  000  0000  0010  100
    0x801c : "DBGBVR3_EL1",       # 10  000  0000  0011  100
    0x8024 : "DBGBVR4_EL1",       # 10  000  0000  0100  100
    0x802c : "DBGBVR5_EL1",       # 10  000  0000  0101  100
    0x8034 : "DBGBVR6_EL1",       # 10  000  0000  0110  100
    0x803c : "DBGBVR7_EL1",       # 10  000  0000  0111  100
    0x8044 : "DBGBVR8_EL1",       # 10  000  0000  1000  100
    0x804c : "DBGBVR9_EL1",       # 10  000  0000  1001  100
    0x8054 : "DBGBVR10_EL1",      # 10  000  0000  1010  100
    0x805c : "DBGBVR11_EL1",      # 10  000  0000  1011  100
    0x8064 : "DBGBVR12_EL1",      # 10  000  0000  1100  100
    0x806c : "DBGBVR13_EL1",      # 10  000  0000  1101  100
    0x8074 : "DBGBVR14_EL1",      # 10  000  0000  1110  100
    0x807c : "DBGBVR15_EL1",      # 10  000  0000  1111  100
    0x8005 : "DBGBCR0_EL1",       # 10  000  0000  0000  101
    0x800d : "DBGBCR1_EL1",       # 10  000  0000  0001  101
    0x8015 : "DBGBCR2_EL1",       # 10  000  0000  0010  101
    0x801d : "DBGBCR3_EL1",       # 10  000  0000  0011  101
    0x8025 : "DBGBCR4_EL1",       # 10  000  0000  0100  101
    0x802d : "DBGBCR5_EL1",       # 10  000  0000  0101  101
    0x8035 : "DBGBCR6_EL1",       # 10  000  0000  0110  101
    0x803d : "DBGBCR7_EL1",       # 10  000  0000  0111  101
    0x8045 : "DBGBCR8_EL1",       # 10  000  0000  1000  101
    0x804d : "DBGBCR9_EL1",       # 10  000  0000  1001  101
    0x8055 : "DBGBCR10_EL1",      # 10  000  0000  1010  101
    0x805d : "DBGBCR11_EL1",      # 10  000  0000  1011  101
    0x8065 : "DBGBCR12_EL1",      # 10  000  0000  1100  101
    0x806d : "DBGBCR13_EL1",      # 10  000  0000  1101  101
    0x8075 : "DBGBCR14_EL1",      # 10  000  0000  1110  101
    0x807d : "DBGBCR15_EL1",      # 10  000  0000  1111  101
    0x8006 : "DBGWVR0_EL1",       # 10  000  0000  0000  110
    0x800e : "DBGWVR1_EL1",       # 10  000  0000  0001  110
    0x8016 : "DBGWVR2_EL1",       # 10  000  0000  0010  110
    0x801e : "DBGWVR3_EL1",       # 10  000  0000  0011  110
    0x8026 : "DBGWVR4_EL1",       # 10  000  0000  0100  110
    0x802e : "DBGWVR5_EL1",       # 10  000  0000  0101  110
    0x8036 : "DBGWVR6_EL1",       # 10  000  0000  0110  110
    0x803e : "DBGWVR7_EL1",       # 10  000  0000  0111  110
    0x8046 : "DBGWVR8_EL1",       # 10  000  0000  1000  110
    0x804e : "DBGWVR9_EL1",       # 10  000  0000  1001  110
    0x8056 : "DBGWVR10_EL1",      # 10  000  0000  1010  110
    0x805e : "DBGWVR11_EL1",      # 10  000  0000  1011  110
    0x8066 : "DBGWVR12_EL1",      # 10  000  0000  1100  110
    0x806e : "DBGWVR13_EL1",      # 10  000  0000  1101  110
    0x8076 : "DBGWVR14_EL1",      # 10  000  0000  1110  110
    0x807e : "DBGWVR15_EL1",      # 10  000  0000  1111  110
    0x8007 : "DBGWCR0_EL1",       # 10  000  0000  0000  111
    0x800f : "DBGWCR1_EL1",       # 10  000  0000  0001  111
    0x8017 : "DBGWCR2_EL1",       # 10  000  0000  0010  111
    0x801f : "DBGWCR3_EL1",       # 10  000  0000  0011  111
    0x8027 : "DBGWCR4_EL1",       # 10  000  0000  0100  111
    0x802f : "DBGWCR5_EL1",       # 10  000  0000  0101  111
    0x8037 : "DBGWCR6_EL1",       # 10  000  0000  0110  111
    0x803f : "DBGWCR7_EL1",       # 10  000  0000  0111  111
    0x8047 : "DBGWCR8_EL1",       # 10  000  0000  1000  111
    0x804f : "DBGWCR9_EL1",       # 10  000  0000  1001  111
    0x8057 : "DBGWCR10_EL1",      # 10  000  0000  1010  111
    0x805f : "DBGWCR11_EL1",      # 10  000  0000  1011  111
    0x8067 : "DBGWCR12_EL1",      # 10  000  0000  1100  111
    0x806f : "DBGWCR13_EL1",      # 10  000  0000  1101  111
    0x8077 : "DBGWCR14_EL1",      # 10  000  0000  1110  111
    0x807f : "DBGWCR15_EL1",      # 10  000  0000  1111  111
    0x9080 : "TEEHBR32_EL1",      # 10  010  0001  0000  000
    0x809c : "OSDLR_EL1",         # 10  000  0001  0011  100
    0x80a4 : "DBGPRCR_EL1",       # 10  000  0001  0100  100
    0x83c6 : "DBGCLAIMSET_EL1",   # 10  000  0111  1000  110
    0x83ce : "DBGCLAIMCLR_EL1",   # 10  000  0111  1001  110
    0xd000 : "CSSELR_EL1",        # 11  010  0000  0000  000
    0xe000 : "VPIDR_EL2",         # 11  100  0000  0000  000
    0xe005 : "VMPIDR_EL2",        # 11  100  0000  0000  101
    0xc082 : "CPACR_EL1",         # 11  000  0001  0000  010
    0xc080 : "SCTLR_EL1",         # 11  000  0001  0000  000
    0xe080 : "SCTLR_EL2",         # 11  100  0001  0000  000
    0xf080 : "SCTLR_EL3",         # 11  110  0001  0000  000
    0xc081 : "ACTLR_EL1",         # 11  000  0001  0000  001
    0xe081 : "ACTLR_EL2",         # 11  100  0001  0000  001
    0xf081 : "ACTLR_EL3",         # 11  110  0001  0000  001
    0xe088 : "HCR_EL2",           # 11  100  0001  0001  000
    0xf088 : "SCR_EL3",           # 11  110  0001  0001  000
    0xe089 : "MDCR_EL2",          # 11  100  0001  0001  001
    0xf089 : "SDER32_EL3",        # 11  110  0001  0001  001
    0xe08a : "CPTR_EL2",          # 11  100  0001  0001  010
    0xf08a : "CPTR_EL3",          # 11  110  0001  0001  010
    0xe08b : "HSTR_EL2",          # 11  100  0001  0001  011
    0xe08f : "HACR_EL2",          # 11  100  0001  0001  111
    0xf099 : "MDCR_EL3",          # 11  110  0001  0011  001
    0xc100 : "TTBR0_EL1",         # 11  000  0010  0000  000
    0xe100 : "TTBR0_EL2",         # 11  100  0010  0000  000
    0xf100 : "TTBR0_EL3",         # 11  110  0010  0000  000
    0xc101 : "TTBR1_EL1",         # 11  000  0010  0000  001
    0xc102 : "TCR_EL1",           # 11  000  0010  0000  010
    0xe102 : "TCR_EL2",           # 11  100  0010  0000  010
    0xf102 : "TCR_EL3",           # 11  110  0010  0000  010
    0xe108 : "VTTBR_EL2",         # 11  100  0010  0001  000
    0xe10a : "VTCR_EL2",          # 11  100  0010  0001  010
    0xe180 : "DACR32_EL2",        # 11  100  0011  0000  000
    0xc200 : "SPSR_EL1",          # 11  000  0100  0000  000
    0xe200 : "SPSR_EL2",          # 11  100  0100  0000  000
    0xf200 : "SPSR_EL3",          # 11  110  0100  0000  000
    0xc201 : "ELR_EL1",           # 11  000  0100  0000  001
    0xe201 : "ELR_EL2",           # 11  100  0100  0000  001
    0xf201 : "ELR_EL3",           # 11  110  0100  0000  001
    0xc208 : "SP_EL0",            # 11  000  0100  0001  000
    0xe208 : "SP_EL1",            # 11  100  0100  0001  000
    0xf208 : "SP_EL2",            # 11  110  0100  0001  000
    0xc210 : "SPSel",             # 11  000  0100  0010  000
    0xda10 : "NZCV",              # 11  011  0100  0010  000
    0xda11 : "DAIF",              # 11  011  0100  0010  001
    0xc212 : "CurrentEL",         # 11  000  0100  0010  010
    0xe218 : "SPSR_irq",          # 11  100  0100  0011  000
    0xe219 : "SPSR_abt",          # 11  100  0100  0011  001
    0xe21a : "SPSR_und",          # 11  100  0100  0011  010
    0xe21b : "SPSR_fiq",          # 11  100  0100  0011  011
    0xda20 : "FPCR",              # 11  011  0100  0100  000
    0xda21 : "FPSR",              # 11  011  0100  0100  001
    0xda28 : "DSPSR_EL0",         # 11  011  0100  0101  000
    0xda29 : "DLR_EL0",           # 11  011  0100  0101  001
    0xe281 : "IFSR32_EL2",        # 11  100  0101  0000  001
    0xc288 : "AFSR0_EL1",         # 11  000  0101  0001  000
    0xe288 : "AFSR0_EL2",         # 11  100  0101  0001  000
    0xf288 : "AFSR0_EL3",         # 11  110  0101  0001  000
    0xc289 : "AFSR1_EL1",         # 11  000  0101  0001  001
    0xe289 : "AFSR1_EL2",         # 11  100  0101  0001  001
    0xf289 : "AFSR1_EL3",         # 11  110  0101  0001  001
    0xc290 : "ESR_EL1",           # 11  000  0101  0010  000
    0xe290 : "ESR_EL2",           # 11  100  0101  0010  000
    0xf290 : "ESR_EL3",           # 11  110  0101  0010  000
    0xe298 : "FPEXC32_EL2",       # 11  100  0101  0011  000
    0xc300 : "FAR_EL1",           # 11  000  0110  0000  000
    0xe300 : "FAR_EL2",           # 11  100  0110  0000  000
    0xf300 : "FAR_EL3",           # 11  110  0110  0000  000
    0xe304 : "HPFAR_EL2",         # 11  100  0110  0000  100
    0xc3a0 : "PAR_EL1",           # 11  000  0111  0100  000
    0xdce0 : "PMCR_EL0",          # 11  011  1001  1100  000
    0xdce1 : "PMCNTENSET_EL0",    # 11  011  1001  1100  001
    0xdce2 : "PMCNTENCLR_EL0",    # 11  011  1001  1100  010
    0xdce3 : "PMOVSCLR_EL0",      # 11  011  1001  1100  011
    0xdce5 : "PMSELR_EL0",        # 11  011  1001  1100  101
    0xdce8 : "PMCCNTR_EL0",       # 11  011  1001  1101  000
    0xdce9 : "PMXEVTYPER_EL0",    # 11  011  1001  1101  001
    0xdcea : "PMXEVCNTR_EL0",     # 11  011  1001  1101  010
    0xdcf0 : "PMUSERENR_EL0",     # 11  011  1001  1110  000
    0xc4f1 : "PMINTENSET_EL1",    # 11  000  1001  1110  001
    0xc4f2 : "PMINTENCLR_EL1",    # 11  000  1001  1110  010
    0xdcf3 : "PMOVSSET_EL0",      # 11  011  1001  1110  011
    0xc510 : "MAIR_EL1",          # 11  000  1010  0010  000
    0xe510 : "MAIR_EL2",          # 11  100  1010  0010  000
    0xf510 : "MAIR_EL3",          # 11  110  1010  0010  000
    0xc518 : "AMAIR_EL1",         # 11  000  1010  0011  000
    0xe518 : "AMAIR_EL2",         # 11  100  1010  0011  000
    0xf518 : "AMAIR_EL3",         # 11  110  1010  0011  000
    0xc600 : "VBAR_EL1",          # 11  000  1100  0000  000
    0xe600 : "VBAR_EL2",          # 11  100  1100  0000  000
    0xf600 : "VBAR_EL3",          # 11  110  1100  0000  000
    0xc602 : "RMR_EL1",           # 11  000  1100  0000  010
    0xe602 : "RMR_EL2",           # 11  100  1100  0000  010
    0xf602 : "RMR_EL3",           # 11  110  1100  0000  010
    0xc681 : "CONTEXTIDR_EL1",    # 11  000  1101  0000  001
    0xde82 : "TPIDR_EL0",         # 11  011  1101  0000  010
    0xe682 : "TPIDR_EL2",         # 11  100  1101  0000  010
    0xf682 : "TPIDR_EL3",         # 11  110  1101  0000  010
    0xde83 : "TPIDRRO_EL0",       # 11  011  1101  0000  011
    0xc684 : "TPIDR_EL1",         # 11  000  1101  0000  100
    0xdf00 : "CNTFRQ_EL0",        # 11  011  1110  0000  000
    0xe703 : "CNTVOFF_EL2",       # 11  100  1110  0000  011
    0xc708 : "CNTKCTL_EL1",       # 11  000  1110  0001  000
    0xe708 : "CNTHCTL_EL2",       # 11  100  1110  0001  000
    0xdf10 : "CNTP_TVAL_EL0",     # 11  011  1110  0010  000
    0xe710 : "CNTHP_TVAL_EL2",    # 11  100  1110  0010  000
    0xff10 : "CNTPS_TVAL_EL1",    # 11  111  1110  0010  000
    0xdf11 : "CNTP_CTL_EL0",      # 11  011  1110  0010  001
    0xe711 : "CNTHP_CTL_EL2",     # 11  100  1110  0010  001
    0xff11 : "CNTPS_CTL_EL1",     # 11  111  1110  0010  001
    0xdf12 : "CNTP_CVAL_EL0",     # 11  011  1110  0010  010
    0xe712 : "CNTHP_CVAL_EL2",    # 11  100  1110  0010  010
    0xff12 : "CNTPS_CVAL_EL1",    # 11  111  1110  0010  010
    0xdf18 : "CNTV_TVAL_EL0",     # 11  011  1110  0011  000
    0xdf19 : "CNTV_CTL_EL0",      # 11  011  1110  0011  001
    0xdf1a : "CNTV_CVAL_EL0",     # 11  011  1110  0011  010
    0xdf40 : "PMEVCNTR0_EL0",     # 11  011  1110  1000  000
    0xdf41 : "PMEVCNTR1_EL0",     # 11  011  1110  1000  001
    0xdf42 : "PMEVCNTR2_EL0",     # 11  011  1110  1000  010
    0xdf43 : "PMEVCNTR3_EL0",     # 11  011  1110  1000  011
    0xdf44 : "PMEVCNTR4_EL0",     # 11  011  1110  1000  100
    0xdf45 : "PMEVCNTR5_EL0",     # 11  011  1110  1000  101
    0xdf46 : "PMEVCNTR6_EL0",     # 11  011  1110  1000  110
    0xdf47 : "PMEVCNTR7_EL0",     # 11  011  1110  1000  111
    0xdf48 : "PMEVCNTR8_EL0",     # 11  011  1110  1001  000
    0xdf49 : "PMEVCNTR9_EL0",     # 11  011  1110  1001  001
    0xdf4a : "PMEVCNTR10_EL0",    # 11  011  1110  1001  010
    0xdf4b : "PMEVCNTR11_EL0",    # 11  011  1110  1001  011
    0xdf4c : "PMEVCNTR12_EL0",    # 11  011  1110  1001  100
    0xdf4d : "PMEVCNTR13_EL0",    # 11  011  1110  1001  101
    0xdf4e : "PMEVCNTR14_EL0",    # 11  011  1110  1001  110
    0xdf4f : "PMEVCNTR15_EL0",    # 11  011  1110  1001  111
    0xdf50 : "PMEVCNTR16_EL0",    # 11  011  1110  1010  000
    0xdf51 : "PMEVCNTR17_EL0",    # 11  011  1110  1010  001
    0xdf52 : "PMEVCNTR18_EL0",    # 11  011  1110  1010  010
    0xdf53 : "PMEVCNTR19_EL0",    # 11  011  1110  1010  011
    0xdf54 : "PMEVCNTR20_EL0",    # 11  011  1110  1010  100
    0xdf55 : "PMEVCNTR21_EL0",    # 11  011  1110  1010  101
    0xdf56 : "PMEVCNTR22_EL0",    # 11  011  1110  1010  110
    0xdf57 : "PMEVCNTR23_EL0",    # 11  011  1110  1010  111
    0xdf58 : "PMEVCNTR24_EL0",    # 11  011  1110  1011  000
    0xdf59 : "PMEVCNTR25_EL0",    # 11  011  1110  1011  001
    0xdf5a : "PMEVCNTR26_EL0",    # 11  011  1110  1011  010
    0xdf5b : "PMEVCNTR27_EL0",    # 11  011  1110  1011  011
    0xdf5c : "PMEVCNTR28_EL0",    # 11  011  1110  1011  100
    0xdf5d : "PMEVCNTR29_EL0",    # 11  011  1110  1011  101
    0xdf5e : "PMEVCNTR30_EL0",    # 11  011  1110  1011  110
    0xdf7f : "PMCCFILTR_EL0",     # 11  011  1110  1111  111
    0xdf60 : "PMEVTYPER0_EL0",    # 11  011  1110  1100  000
    0xdf61 : "PMEVTYPER1_EL0",    # 11  011  1110  1100  001
    0xdf62 : "PMEVTYPER2_EL0",    # 11  011  1110  1100  010
    0xdf63 : "PMEVTYPER3_EL0",    # 11  011  1110  1100  011
    0xdf64 : "PMEVTYPER4_EL0",    # 11  011  1110  1100  100
    0xdf65 : "PMEVTYPER5_EL0",    # 11  011  1110  1100  101
    0xdf66 : "PMEVTYPER6_EL0",    # 11  011  1110  1100  110
    0xdf67 : "PMEVTYPER7_EL0",    # 11  011  1110  1100  111
    0xdf68 : "PMEVTYPER8_EL0",    # 11  011  1110  1101  000
    0xdf69 : "PMEVTYPER9_EL0",    # 11  011  1110  1101  001
    0xdf6a : "PMEVTYPER10_EL0",   # 11  011  1110  1101  010
    0xdf6b : "PMEVTYPER11_EL0",   # 11  011  1110  1101  011
    0xdf6c : "PMEVTYPER12_EL0",   # 11  011  1110  1101  100
    0xdf6d : "PMEVTYPER13_EL0",   # 11  011  1110  1101  101
    0xdf6e : "PMEVTYPER14_EL0",   # 11  011  1110  1101  110
    0xdf6f : "PMEVTYPER15_EL0",   # 11  011  1110  1101  111
    0xdf70 : "PMEVTYPER16_EL0",   # 11  011  1110  1110  000
    0xdf71 : "PMEVTYPER17_EL0",   # 11  011  1110  1110  001
    0xdf72 : "PMEVTYPER18_EL0",   # 11  011  1110  1110  010
    0xdf73 : "PMEVTYPER19_EL0",   # 11  011  1110  1110  011
    0xdf74 : "PMEVTYPER20_EL0",   # 11  011  1110  1110  100
    0xdf75 : "PMEVTYPER21_EL0",   # 11  011  1110  1110  101
    0xdf76 : "PMEVTYPER22_EL0",   # 11  011  1110  1110  110
    0xdf77 : "PMEVTYPER23_EL0",   # 11  011  1110  1110  111
    0xdf78 : "PMEVTYPER24_EL0",   # 11  011  1110  1111  000
    0xdf79 : "PMEVTYPER25_EL0",   # 11  011  1110  1111  001
    0xdf7a : "PMEVTYPER26_EL0",   # 11  011  1110  1111  010
    0xdf7b : "PMEVTYPER27_EL0",   # 11  011  1110  1111  011
    0xdf7c : "PMEVTYPER28_EL0",   # 11  011  1110  1111  100
    0xdf7d : "PMEVTYPER29_EL0",   # 11  011  1110  1111  101
    0xdf7e : "PMEVTYPER30_EL0",   # 11  011  1110  1111  110
    0x8808 : "TRCPRGCTLR",        # 10  001  0000  0001  000
    0x8810 : "TRCPROCSELR",       # 10  001  0000  0010  000
    0x8820 : "TRCCONFIGR",        # 10  001  0000  0100  000
    0x8830 : "TRCAUXCTLR",        # 10  001  0000  0110  000
    0x8840 : "TRCEVENTCTL0R",     # 10  001  0000  1000  000
    0x8848 : "TRCEVENTCTL1R",     # 10  001  0000  1001  000
    0x8858 : "TRCSTALLCTLR",      # 10  001  0000  1011  000
    0x8860 : "TRCTSCTLR",         # 10  001  0000  1100  000
    0x8868 : "TRCSYNCPR",         # 10  001  0000  1101  000
    0x8870 : "TRCCCCTLR",         # 10  001  0000  1110  000
    0x8878 : "TRCBBCTLR",         # 10  001  0000  1111  000
    0x8801 : "TRCTRACEIDR",       # 10  001  0000  0000  001
    0x8809 : "TRCQCTLR",          # 10  001  0000  0001  001
    0x8802 : "TRCVICTLR",         # 10  001  0000  0000  010
    0x880a : "TRCVIIECTLR",       # 10  001  0000  0001  010
    0x8812 : "TRCVISSCTLR",       # 10  001  0000  0010  010
    0x881a : "TRCVIPCSSCTLR",     # 10  001  0000  0011  010
    0x8842 : "TRCVDCTLR",         # 10  001  0000  1000  010
    0x884a : "TRCVDSACCTLR",      # 10  001  0000  1001  010
    0x8852 : "TRCVDARCCTLR",      # 10  001  0000  1010  010
    0x8804 : "TRCSEQEVR0",        # 10  001  0000  0000  100
    0x880c : "TRCSEQEVR1",        # 10  001  0000  0001  100
    0x8814 : "TRCSEQEVR2",        # 10  001  0000  0010  100
    0x8834 : "TRCSEQRSTEVR",      # 10  001  0000  0110  100
    0x883c : "TRCSEQSTR",         # 10  001  0000  0111  100
    0x8844 : "TRCEXTINSELR",      # 10  001  0000  1000  100
    0x8805 : "TRCCNTRLDVR0",      # 10  001  0000  0000  101
    0x880d : "TRCCNTRLDVR1",      # 10  001  0000  0001  101
    0x8815 : "TRCCNTRLDVR2",      # 10  001  0000  0010  101
    0x881d : "TRCCNTRLDVR3",      # 10  001  0000  0011  101
    0x8825 : "TRCCNTCTLR0",       # 10  001  0000  0100  101
    0x882d : "TRCCNTCTLR1",       # 10  001  0000  0101  101
    0x8835 : "TRCCNTCTLR2",       # 10  001  0000  0110  101
    0x883d : "TRCCNTCTLR3",       # 10  001  0000  0111  101
    0x8845 : "TRCCNTVR0",         # 10  001  0000  1000  101
    0x884d : "TRCCNTVR1",         # 10  001  0000  1001  101
    0x8855 : "TRCCNTVR2",         # 10  001  0000  1010  101
    0x885d : "TRCCNTVR3",         # 10  001  0000  1011  101
    0x8807 : "TRCIMSPEC0",        # 10  001  0000  0000  111
    0x880f : "TRCIMSPEC1",        # 10  001  0000  0001  111
    0x8817 : "TRCIMSPEC2",        # 10  001  0000  0010  111
    0x881f : "TRCIMSPEC3",        # 10  001  0000  0011  111
    0x8827 : "TRCIMSPEC4",        # 10  001  0000  0100  111
    0x882f : "TRCIMSPEC5",        # 10  001  0000  0101  111
    0x8837 : "TRCIMSPEC6",        # 10  001  0000  0110  111
    0x883f : "TRCIMSPEC7",        # 10  001  0000  0111  111
    0x8890 : "TRCRSCTLR2",        # 10  001  0001  0010  000
    0x8898 : "TRCRSCTLR3",        # 10  001  0001  0011  000
    0x88a0 : "TRCRSCTLR4",        # 10  001  0001  0100  000
    0x88a8 : "TRCRSCTLR5",        # 10  001  0001  0101  000
    0x88b0 : "TRCRSCTLR6",        # 10  001  0001  0110  000
    0x88b8 : "TRCRSCTLR7",        # 10  001  0001  0111  000
    0x88c0 : "TRCRSCTLR8",        # 10  001  0001  1000  000
    0x88c8 : "TRCRSCTLR9",        # 10  001  0001  1001  000
    0x88d0 : "TRCRSCTLR10",       # 10  001  0001  1010  000
    0x88d8 : "TRCRSCTLR11",       # 10  001  0001  1011  000
    0x88e0 : "TRCRSCTLR12",       # 10  001  0001  1100  000
    0x88e8 : "TRCRSCTLR13",       # 10  001  0001  1101  000
    0x88f0 : "TRCRSCTLR14",       # 10  001  0001  1110  000
    0x88f8 : "TRCRSCTLR15",       # 10  001  0001  1111  000
    0x8881 : "TRCRSCTLR16",       # 10  001  0001  0000  001
    0x8889 : "TRCRSCTLR17",       # 10  001  0001  0001  001
    0x8891 : "TRCRSCTLR18",       # 10  001  0001  0010  001
    0x8899 : "TRCRSCTLR19",       # 10  001  0001  0011  001
    0x88a1 : "TRCRSCTLR20",       # 10  001  0001  0100  001
    0x88a9 : "TRCRSCTLR21",       # 10  001  0001  0101  001
    0x88b1 : "TRCRSCTLR22",       # 10  001  0001  0110  001
    0x88b9 : "TRCRSCTLR23",       # 10  001  0001  0111  001
    0x88c1 : "TRCRSCTLR24",       # 10  001  0001  1000  001
    0x88c9 : "TRCRSCTLR25",       # 10  001  0001  1001  001
    0x88d1 : "TRCRSCTLR26",       # 10  001  0001  1010  001
    0x88d9 : "TRCRSCTLR27",       # 10  001  0001  1011  001
    0x88e1 : "TRCRSCTLR28",       # 10  001  0001  1100  001
    0x88e9 : "TRCRSCTLR29",       # 10  001  0001  1101  001
    0x88f1 : "TRCRSCTLR30",       # 10  001  0001  1110  001
    0x88f9 : "TRCRSCTLR31",       # 10  001  0001  1111  001
    0x8882 : "TRCSSCCR0",         # 10  001  0001  0000  010
    0x888a : "TRCSSCCR1",         # 10  001  0001  0001  010
    0x8892 : "TRCSSCCR2",         # 10  001  0001  0010  010
    0x889a : "TRCSSCCR3",         # 10  001  0001  0011  010
    0x88a2 : "TRCSSCCR4",         # 10  001  0001  0100  010
    0x88aa : "TRCSSCCR5",         # 10  001  0001  0101  010
    0x88b2 : "TRCSSCCR6",         # 10  001  0001  0110  010
    0x88ba : "TRCSSCCR7",         # 10  001  0001  0111  010
    0x88c2 : "TRCSSCSR0",         # 10  001  0001  1000  010
    0x88ca : "TRCSSCSR1",         # 10  001  0001  1001  010
    0x88d2 : "TRCSSCSR2",         # 10  001  0001  1010  010
    0x88da : "TRCSSCSR3",         # 10  001  0001  1011  010
    0x88e2 : "TRCSSCSR4",         # 10  001  0001  1100  010
    0x88ea : "TRCSSCSR5",         # 10  001  0001  1101  010
    0x88f2 : "TRCSSCSR6",         # 10  001  0001  1110  010
    0x88fa : "TRCSSCSR7",         # 10  001  0001  1111  010
    0x8883 : "TRCSSPCICR0",       # 10  001  0001  0000  011
    0x888b : "TRCSSPCICR1",       # 10  001  0001  0001  011
    0x8893 : "TRCSSPCICR2",       # 10  001  0001  0010  011
    0x889b : "TRCSSPCICR3",       # 10  001  0001  0011  011
    0x88a3 : "TRCSSPCICR4",       # 10  001  0001  0100  011
    0x88ab : "TRCSSPCICR5",       # 10  001  0001  0101  011
    0x88b3 : "TRCSSPCICR6",       # 10  001  0001  0110  011
    0x88bb : "TRCSSPCICR7",       # 10  001  0001  0111  011
    0x88a4 : "TRCPDCR",           # 10  001  0001  0100  100
    0x8900 : "TRCACVR0",          # 10  001  0010  0000  000
    0x8910 : "TRCACVR1",          # 10  001  0010  0010  000
    0x8920 : "TRCACVR2",          # 10  001  0010  0100  000
    0x8930 : "TRCACVR3",          # 10  001  0010  0110  000
    0x8940 : "TRCACVR4",          # 10  001  0010  1000  000
    0x8950 : "TRCACVR5",          # 10  001  0010  1010  000
    0x8960 : "TRCACVR6",          # 10  001  0010  1100  000
    0x8970 : "TRCACVR7",          # 10  001  0010  1110  000
    0x8901 : "TRCACVR8",          # 10  001  0010  0000  001
    0x8911 : "TRCACVR9",          # 10  001  0010  0010  001
    0x8921 : "TRCACVR10",         # 10  001  0010  0100  001
    0x8931 : "TRCACVR11",         # 10  001  0010  0110  001
    0x8941 : "TRCACVR12",         # 10  001  0010  1000  001
    0x8951 : "TRCACVR13",         # 10  001  0010  1010  001
    0x8961 : "TRCACVR14",         # 10  001  0010  1100  001
    0x8971 : "TRCACVR15",         # 10  001  0010  1110  001
    0x8902 : "TRCACATR0",         # 10  001  0010  0000  010
    0x8912 : "TRCACATR1",         # 10  001  0010  0010  010
    0x8922 : "TRCACATR2",         # 10  001  0010  0100  010
    0x8932 : "TRCACATR3",         # 10  001  0010  0110  010
    0x8942 : "TRCACATR4",         # 10  001  0010  1000  010
    0x8952 : "TRCACATR5",         # 10  001  0010  1010  010
    0x8962 : "TRCACATR6",         # 10  001  0010  1100  010
    0x8972 : "TRCACATR7",         # 10  001  0010  1110  010
    0x8903 : "TRCACATR8",         # 10  001  0010  0000  011
    0x8913 : "TRCACATR9",         # 10  001  0010  0010  011
    0x8923 : "TRCACATR10",        # 10  001  0010  0100  011
    0x8933 : "TRCACATR11",        # 10  001  0010  0110  011
    0x8943 : "TRCACATR12",        # 10  001  0010  1000  011
    0x8953 : "TRCACATR13",        # 10  001  0010  1010  011
    0x8963 : "TRCACATR14",        # 10  001  0010  1100  011
    0x8973 : "TRCACATR15",        # 10  001  0010  1110  011
    0x8904 : "TRCDVCVR0",         # 10  001  0010  0000  100
    0x8924 : "TRCDVCVR1",         # 10  001  0010  0100  100
    0x8944 : "TRCDVCVR2",         # 10  001  0010  1000  100
    0x8964 : "TRCDVCVR3",         # 10  001  0010  1100  100
    0x8905 : "TRCDVCVR4",         # 10  001  0010  0000  101
    0x8925 : "TRCDVCVR5",         # 10  001  0010  0100  101
    0x8945 : "TRCDVCVR6",         # 10  001  0010  1000  101
    0x8965 : "TRCDVCVR7",         # 10  001  0010  1100  101
    0x8906 : "TRCDVCMR0",         # 10  001  0010  0000  110
    0x8926 : "TRCDVCMR1",         # 10  001  0010  0100  110
    0x8946 : "TRCDVCMR2",         # 10  001  0010  1000  110
    0x8966 : "TRCDVCMR3",         # 10  001  0010  1100  110
    0x8907 : "TRCDVCMR4",         # 10  001  0010  0000  111
    0x8927 : "TRCDVCMR5",         # 10  001  0010  0100  111
    0x8947 : "TRCDVCMR6",         # 10  001  0010  1000  111
    0x8967 : "TRCDVCMR7",         # 10  001  0010  1100  111
    0x8980 : "TRCCIDCVR0",        # 10  001  0011  0000  000
    0x8990 : "TRCCIDCVR1",        # 10  001  0011  0010  000
    0x89a0 : "TRCCIDCVR2",        # 10  001  0011  0100  000
    0x89b0 : "TRCCIDCVR3",        # 10  001  0011  0110  000
    0x89c0 : "TRCCIDCVR4",        # 10  001  0011  1000  000
    0x89d0 : "TRCCIDCVR5",        # 10  001  0011  1010  000
    0x89e0 : "TRCCIDCVR6",        # 10  001  0011  1100  000
    0x89f0 : "TRCCIDCVR7",        # 10  001  0011  1110  000
    0x8981 : "TRCVMIDCVR0",       # 10  001  0011  0000  001
    0x8991 : "TRCVMIDCVR1",       # 10  001  0011  0010  001
    0x89a1 : "TRCVMIDCVR2",       # 10  001  0011  0100  001
    0x89b1 : "TRCVMIDCVR3",       # 10  001  0011  0110  001
    0x89c1 : "TRCVMIDCVR4",       # 10  001  0011  1000  001
    0x89d1 : "TRCVMIDCVR5",       # 10  001  0011  1010  001
    0x89e1 : "TRCVMIDCVR6",       # 10  001  0011  1100  001
    0x89f1 : "TRCVMIDCVR7",       # 10  001  0011  1110  001
    0x8982 : "TRCCIDCCTLR0",      # 10  001  0011  0000  010
    0x898a : "TRCCIDCCTLR1",      # 10  001  0011  0001  010
    0x8992 : "TRCVMIDCCTLR0",     # 10  001  0011  0010  010
    0x899a : "TRCVMIDCCTLR1",     # 10  001  0011  0011  010
    0x8b84 : "TRCITCTRL",         # 10  001  0111  0000  100
    0x8bc6 : "TRCCLAIMSET",       # 10  001  0111  1000  110
    0x8bce : "TRCCLAIMCLR",       # 10  001  0111  1001  110
    0xc663 : "ICC_BPR1_EL1",      # 11  000  1100  1100  011
    0xc643 : "ICC_BPR0_EL1",      # 11  000  1100  1000  011
    0xc230 : "ICC_PMR_EL1",       # 11  000  0100  0110  000
    0xc664 : "ICC_CTLR_EL1",      # 11  000  1100  1100  100
    0xf664 : "ICC_CTLR_EL3",      # 11  110  1100  1100  100
    0xc665 : "ICC_SRE_EL1",       # 11  000  1100  1100  101
    0xe64d : "ICC_SRE_EL2",       # 11  100  1100  1001  101
    0xf665 : "ICC_SRE_EL3",       # 11  110  1100  1100  101
    0xc666 : "ICC_IGRPEN0_EL1",   # 11  000  1100  1100  110
    0xc667 : "ICC_IGRPEN1_EL1",   # 11  000  1100  1100  111
    0xf667 : "ICC_IGRPEN1_EL3",   # 11  110  1100  1100  111
    0xc668 : "ICC_SEIEN_EL1",     # 11  000  1100  1101  000
    0xc644 : "ICC_AP0R0_EL1",     # 11  000  1100  1000  100
    0xc645 : "ICC_AP0R1_EL1",     # 11  000  1100  1000  101
    0xc646 : "ICC_AP0R2_EL1",     # 11  000  1100  1000  110
    0xc647 : "ICC_AP0R3_EL1",     # 11  000  1100  1000  111
    0xc648 : "ICC_AP1R0_EL1",     # 11  000  1100  1001  000
    0xc649 : "ICC_AP1R1_EL1",     # 11  000  1100  1001  001
    0xc64a : "ICC_AP1R2_EL1",     # 11  000  1100  1001  010
    0xc64b : "ICC_AP1R3_EL1",     # 11  000  1100  1001  011
    0xe640 : "ICH_AP0R0_EL2",     # 11  100  1100  1000  000
    0xe641 : "ICH_AP0R1_EL2",     # 11  100  1100  1000  001
    0xe642 : "ICH_AP0R2_EL2",     # 11  100  1100  1000  010
    0xe643 : "ICH_AP0R3_EL2",     # 11  100  1100  1000  011
    0xe648 : "ICH_AP1R0_EL2",     # 11  100  1100  1001  000
    0xe649 : "ICH_AP1R1_EL2",     # 11  100  1100  1001  001
    0xe64a : "ICH_AP1R2_EL2",     # 11  100  1100  1001  010
    0xe64b : "ICH_AP1R3_EL2",     # 11  100  1100  1001  011
    0xe658 : "ICH_HCR_EL2",       # 11  100  1100  1011  000
    0xe65a : "ICH_MISR_EL2",      # 11  100  1100  1011  010
    0xe65f : "ICH_VMCR_EL2",      # 11  100  1100  1011  111
    0xe64c : "ICH_VSEIR_EL2",     # 11  100  1100  1001  100
    0xe660 : "ICH_LR0_EL2",       # 11  100  1100  1100  000
    0xe661 : "ICH_LR1_EL2",       # 11  100  1100  1100  001
    0xe662 : "ICH_LR2_EL2",       # 11  100  1100  1100  010
    0xe663 : "ICH_LR3_EL2",       # 11  100  1100  1100  011
    0xe664 : "ICH_LR4_EL2",       # 11  100  1100  1100  100
    0xe665 : "ICH_LR5_EL2",       # 11  100  1100  1100  101
    0xe666 : "ICH_LR6_EL2",       # 11  100  1100  1100  110
    0xe667 : "ICH_LR7_EL2",       # 11  100  1100  1100  111
    0xe668 : "ICH_LR8_EL2",       # 11  100  1100  1101  000
    0xe669 : "ICH_LR9_EL2",       # 11  100  1100  1101  001
    0xe66a : "ICH_LR10_EL2",      # 11  100  1100  1101  010
    0xe66b : "ICH_LR11_EL2",      # 11  100  1100  1101  011
    0xe66c : "ICH_LR12_EL2",      # 11  100  1100  1101  100
    0xe66d : "ICH_LR13_EL2",      # 11  100  1100  1101  101
    0xe66e : "ICH_LR14_EL2",      # 11  100  1100  1101  110
    0xe66f : "ICH_LR15_EL2",      # 11  100  1100  1101  111
    0xff90 : "CPM_IOACC_CTL_EL3"
}

# generated from https://github.com/gdelugre/ida-arm-system-highlight
regs32 = {
                                # cpnum Op1  CRn   CRm   Op2
                                # 1111  111  1111  1111  111
    0x38000 : "DBGDIDR",
    0x38002 : "DBGDTRRX",
    0x38004 : "DBGBVR0",
    0x38005 : "DBGBCR0",
    0x38006 : "DBGWVR0",
    0x38007 : "DBGWCR0",
    0x38008 : "DBGDSCR",
    0x3800c : "DBGBVR1",
    0x3800d : "DBGBCR1",
    0x3800e : "DBGWVR1",
    0x3800f : "DBGWCR1",
    0x38012 : "DBGDSCR",
    0x38014 : "DBGBVR2",
    0x38015 : "DBGBCR2",
    0x38016 : "DBGWVR2",
    0x38017 : "DBGWCR2",
    0x3801a : "DBGDTRTX",
    0x3801c : "DBGBVR3",
    0x3801d : "DBGBCR3",
    0x3801e : "DBGWVR3",
    0x3801f : "DBGWCR3",
    0x38021 : "DBGBXVR0",
    0x38024 : "DBGBVR4",
    0x38025 : "DBGBCR4",
    0x38028 : "DBGDTRRX",
    0x38029 : "DBGBXVR1",
    0x3802c : "DBGBVR5",
    0x3802d : "DBGBCR5",
    0x38030 : "DBGWFAR",
    0x38038 : "DBGVCR",
    0x38080 : "DBGDRAR",
    0x38084 : "DBGOSLAR",
    0x38084 : "DBGOSLSR",
    0x3809c : "DBGOSDLR",
    0x380a4 : "DBGPRCR",
    0x38100 : "DBGDSAR",
    0x38387 : "DBGDEVID2",
    0x3838f : "DBGDEVID1",
    0x383f6 : "DBGAUTHSTATUS",
    0x38397 : "DBGDEVID",
    0x3c000 : "MIDR",
    0x3c001 : "CTR",
    0x3c002 : "TCMTR",
    0x3c003 : "TLBTR",
    0x3c004 : "MIDR",
    0x3c005 : "MPIDR",
    0x3c006 : "REVIDR",
    0x3c007 : "MIDR",
    0x3c008 : "ID_PFR0",
    0x3c009 : "ID_PFR1",
    0x3c00a : "ID_DFR0",
    0x3c00b : "ID_AFR0",
    0x3c00c : "ID_MMFR0",
    0x3c00d : "ID_MMFR1",
    0x3c00e : "ID_MMFR2",
    0x3c00f : "ID_MMFR3",
    0x3c010 : "ID_ISAR0",
    0x3c011 : "ID_ISAR1",
    0x3c012 : "ID_ISAR2",
    0x3c013 : "ID_ISAR3",
    0x3c014 : "ID_ISAR4",
    0x3c015 : "ID_ISAR5",
    0x3c800 : "CCSIDR",
    0x3c801 : "CLIDR",
    0x3c807 : "AIDR",
    0x3d000 : "CCSELR",
    0x3e000 : "VPIDR",
    0x3e005 : "VMPIDR",
    0x3c080 : "SCTLR",
    0x3c081 : "ACTLR",
    0x3c082 : "CPACR",
    0x3c088 : "SCR",
    0x3c089 : "SDER",
    0x3c08a : "NSACR",
    0x3e080 : "HSCTLR",
    0x3e081 : "HACTLR",
    0x3e088 : "HCR",
    0x3e089 : "HDCR",
    0x3e08a : "HCPTR",
    0x3e08b : "HSTR",
    0x3e08f : "HACR",
    0x3c510 : "MAIR0",
    0x3c511 : "MAIR1",
    0x3c518 : "AMAIR0",
    0x3c519 : "AMAIR1",
    0x3e510 : "HMAIR0",
    0x3e511 : "HMAIR1",
    0x3e518 : "HAMAIR0",
    0x3e519 : "HAMAIR1",
    0x3c600 : "VBAR",
    0x3c601 : "MVBAR",
    0x3c608 : "ISR",
    0x3e600 : "HVBAR",
    0x3c680 : "FCSEIDR",
    0x3c681 : "CONTEXTIDR",
    0x3c682 : "TPIDRURW",
    0x3c683 : "TPIDRURO",
    0x3c684 : "TPIDRPRW",
    0x3e682 : "HTPIDR",
    0x3c700 : "CNTFRQ",
    0x3c780 : "IL1Data0",
    0x3c781 : "IL1Data1",
    0x3c782 : "IL1Data2",
    0x3c788 : "DL1Data0",
    0x3c789 : "DL1Data1",
    0x3c78a : "DL1Data2",
    0x3c7e1 : "CCNT",
    0x3c7e2 : "PMN0",
    0x3c7e3 : "PMN1",
    0x3c7a0 : "RAMINDEX",
    0x3cf80 : "L2ACTLR",
    0x3cf83 : "L2FPR",
    0x3e780 : "CBAR",
    0x3c100 : "TTBR0",
    0x3c101 : "TTBR1",
    0x3e102 : "HTCR",
    0x3e10a : "VTCR",
    0x3c180 : "DACR",
    0x3c280 : "DFSR",
    0x3c281 : "IFSR",
    0x3c288 : "ADFSR",
    0x3c288 : "AIFSR",
    0x3e288 : "HADFSR",
    0x3e289 : "HAIFSR",
    0x3e290 : "HSR",
    0x3c300 : "DFAR",
    0x3c302 : "IFAR",
    0x3e300 : "HDFAR",
    0x3e302 : "HIFAR",
    0x3e304 : "HPFAR",
    0x3c384 : "NOP",
    0x3c388 : "ICIALLUIS",
    0x3c38e : "BPIALLIS",
    0x3c3d1 : "DCCMVAC",
    0x3c3d2 : "DCCSW",
    0x3c3d4 : "CP15DSB",
    0x3c3d5 : "CP15DMB",
    0x3c3d9 : "DCCMVAU",
    0x3c3e9 : "NOP",
    0x3c3f1 : "DCCIMVAC",
    0x3c3f2 : "DCCISW",
    0x3c3a0 : "PAR",
    0x3c3a8 : "ICIALLU",
    0x3c3a9 : "ICIMVAU",
    0x3c3ac : "CP15ISB",
    0x3c3ae : "BPIALL",
    0x3c3af : "BPIMVA",
    0x3c3b1 : "DCIMVAC",
    0x3c3b2 : "DCISW",
    0x3c3c0 : "ATS1CPR",
    0x3c3c1 : "ATS1CPW",
    0x3c3c2 : "ATS1CUR",
    0x3c3c3 : "ATS1CUW",
    0x3c3c4 : "ATS12NSOPR",
    0x3c3c5 : "ATS12NSOPW",
    0x3c3c6 : "ATS12NSOUR",
    0x3c3c7 : "ATS12NSOUW",
    0x3e3c0 : "ATS1HR",
    0x3e3c1 : "ATS1HR",
    0x3c418 : "TLBIALLIS",
    0x3c419 : "TLBIMVAIS",
    0x3c41a : "TLBIASIDIS",
    0x3c41b : "TLBIMVAAIS",
    0x3c428 : "ITLBIALL",
    0x3c429 : "ITLBIMVA",
    0x3c42a : "ITLBIASID",
    0x3c430 : "DTLBIALL",
    0x3c431 : "DTLBIMVA",
    0x3c432 : "DTLBIASID",
    0x3c438 : "TLBIALL",
    0x3c439 : "TLBIMVA",
    0x3c43a : "TLBIASID",
    0x3c43b : "TLBIMVAA",
    0x3e418 : "TLBIALLHIS",
    0x3e419 : "TLBIMVAHIS",
    0x3e41c : "TLBIALLNSNHIS",
    0x3e438 : "TLBIALLH",
    0x3e439 : "TLBIMVAH",
    0x3e43c : "TLBIALLNSNH",
    0x3c4e0 : "PMCR",
    0x3c4e1 : "PMNCNTENSET",
    0x3c4e2 : "PMNCNTENCLR",
    0x3c4e3 : "PMOVSR",
    0x3c4e4 : "PMSWINC",
    0x3c4e5 : "PMSELR",
    0x3c4e6 : "PMCEID0",
    0x3c4e7 : "PMCEID1",
    0x3c4e8 : "PMCCNTR",
    0x3c4e9 : "PMXEVTYPER",
    0x3c4ea : "PMXEVCNTR",
    0x3c4f0 : "PMUSERENR",
    0x3c4f1 : "PMINTENSET",
    0x3c4f2 : "PMINTENCLR",
    0x3c4f3 : "PMOVSSET",
    0x3cc82 : "L2CTLR",
    0x3cc83 : "L2ECTLR",
}

shifts = [ 14, 11, 7, 3, 0 ]
shiftz = [ 14, 11, -1, 7, 3, 0 ]

inttype = None
myenum64 = None
myenum32 = None

class cblock_visitor_t(idaapi.ctree_visitor_t):
    def __init__(self):
        idaapi.ctree_visitor_t.__init__(self, idaapi.CV_FAST)
        return

    def visit_expr(self, expr):
        try:
            if expr.op == idaapi.cot_call and expr.x.op == idaapi.cot_helper:
                #print idaapi.tag_remove(expr.x.print1(None))
                if expr.x.helper == "ARM64_SYSREG" and len(expr.a) == 5: # and idaapi.getseg(expr.ea).use64():
                    reg = 0
                    for j, i in enumerate(expr.a):
                        if i.type != inttype:
                            break
                        #print i.n.value(i.type)
                        reg = reg | (i.numval() << shifts[j])
                    else:
                        if reg in regs64.keys():
                            # apparently, we cannot change the name
                            # so remove all params but the first one,
                            # make it symbolic, and set its value accordingly
                            while len(expr.a) > 1:
                                expr.a.pop_back()
                            expr.a[0].n.nf.flags = idaapi.enumflag()
                            expr.a[0].n.nf.serial = 0
                            expr.a[0].n.nf.type_name = self.mkenu64()
                            expr.a[0].n._value = reg
                            #print "ok"
                if expr.x.helper == "__mrc" and len(expr.a) == 5: # and not idaapi.getseg(expr.ea).use64():
                    reg = 0
                    for j, i in enumerate(expr.a):
                        if i.type != inttype:
                            break
                        reg = reg | (i.numval() << shifts[j])
                    else:
                        if reg in regs32.keys():
                            while len(expr.a) > 1:
                                expr.a.pop_back()
                            expr.a[0].n.nf.flags = idaapi.enumflag()
                            expr.a[0].n.nf.serial = 0
                            expr.a[0].n.nf.type_name = self.mkenu32()
                            expr.a[0].n._value = reg
                if expr.x.helper == "__mcr" and len(expr.a) == 6: # and not idaapi.getseg(expr.ea).use64():
                    reg = 0
                    for j, i in enumerate(expr.a):
                        if shiftz[j] < 0:
                            continue
                        if i.type != inttype:
                            break
                        reg = reg | (i.numval() << shiftz[j])
                    else:
                        if reg in regs32.keys():
                            while len(expr.a) > 3:
                                expr.a.pop_back()
                            expr.a[1] = expr.a[2]
                            expr.a.pop_back()
                            expr.a[0].n.nf.flags = idaapi.enumflag()
                            expr.a[0].n.nf.serial = 0
                            expr.a[0].n.nf.type_name = self.mkenu32()
                            expr.a[0].n._value = reg
        except:
            traceback.print_exc()
        return 0

    def mkenu64(self):
        global myenum64
        if myenum64 is None:
            myenum64 = "ARM64_SYSREG_aenum"
            enu = idaapi.add_enum(0, myenum64, 0)
            for i in regs64.keys():
                idaapi.add_enum_member(enu, regs64[i], i)
        return myenum64

    def mkenu32(self):
        global myenum32
        if myenum32 is None:
            myenum32 = "ARM32_SYSREG_aenum"
            enu = idaapi.add_enum(0, myenum32, 0)
            for i in regs32.keys():
                idaapi.add_enum_member(enu, regs32[i], i)
        return myenum32

class hexrays_callback_info(object):
    def __init__(self):
        return

    def event_callback(self, event, *args):
        try:
            if event == idaapi.hxe_maturity:
                cfunc, maturity = args
                if maturity == idaapi.CMAT_BUILT:
                    cbv = cblock_visitor_t()
                    cbv.apply_to(cfunc.body, None)
        except:
            traceback.print_exc()
        return 0

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

class HexHNightPlugin_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_HIDE
    comment = "show symbolic names for ARM sysregs in Pseudocode-View"
    help = "Runs transparently"
    wanted_name = "HexNight"
    wanted_hotkey = ""

    def init(self):
        # Some initialization
        global hexnight_cb_info, hexnight_cb, inttype

        if idaapi.init_hexrays_plugin() and idaapi.ph_get_id() == idaapi.PLFM_ARM:
            inttype = idaapi.get_int_type_by_width_and_sign(4, True)
            hexnight_cb_info = hexrays_callback_info()
            hexnight_cb = hexnight_cb_info.event_callback
            if idaapi.install_hexrays_callback(hexnight_cb):
                print "Hexnight plugin installed"
                addon = idaapi.addon_info_t();
                addon.id = "org.xerub.hexnight";
                addon.name = "Hexnight";
                addon.producer = "xerub";
                addon.url = "https://twitter.com/xerub";
                addon.version = "6.95";
                idaapi.register_addon( addon );
                return idaapi.PLUGIN_KEEP
        print "Hexnight plugin failed"
        return idaapi.PLUGIN_SKIP

    def run(self, arg=0):
        return

    def term(self):
        remove()

def PLUGIN_ENTRY():
    return HexHNightPlugin_t()


repository: maskelihileci/IDA_Ollvm-unflattener
file: ollvm-unflattener.py
description: Control Flow Flattening Deobfuscator for Obfuscator-LLVM as a plugin for IDA Pro.
created_at: 2025-04-14 12:09:54+00:00
pushed_at: 2025-04-14 12:41:54+00:00
url:
https://github.com/maskelihileci/IDA_Ollvm-unflattener/blob/84000443d9a9f30a0b4908fc34f2e23b40923e79/ollvm-unflattener.py

content:
# MiasmUnflattener/__init__.py or main.py

import ida_idaapi
import ida_kernwin
import ida_funcs
import ida_bytes
import ida_auto
import ida_nalt    # For get_imagebase
import ida_segment # For segment handling
import idc         # For legacy IDA API fallbacks
import ida_ida
import idaapi

import sys
import os
import io          # For potential memory stream handling later
import logging     # Use standard logging for unflattener core


from miasm.core.locationdb import LocationDB
from miasm.analysis.binary import Container

from MiasmUnflattener.unflattener import Unflattener

# --- Configuration ---
PLUGIN_NAME = "OLLVM Unflattener"
PLUGIN_VERSION = "1.0.0" # Version update for completion
PLUGIN_AUTHOR = "maskelihileci"
UNFLATTENER_AUTHOR = "cdong1012 -> https://github.com/cdong1012"
PLUGIN_COMMENT = "Unflattens OLLVM control flow flattening using Miasm."
PLUGIN_HELP = "Right-click on a function start or inside a function in Disassembly or Functions view."
WANTED_HOTKEY = "" # Assign hotkey later if desired

# Action names (must be unique)
ACTION_UNFLATTEN = "miasm_unflattener:unflatten"
ACTION_UNFLATTEN_RECURSIVE = "miasm_unflattener:unflatten_recursive"

# --- Logging Setup ---
# Configure logging to send messages from the unflattener core to IDA output window
core_logger = logging.getLogger("unflattener") # Match the logger name in unflattener.py
logger = logging.getLogger(__name__)
# Clear existing handlers if any during reload
for handler in core_logger.handlers[:]:
    core_logger.removeHandler(handler)

ida_handler = None # Define globally to remove in term()
class IdaLogHandler(logging.Handler):
    def emit(self, record):
        try:
            log_entry = self.format(record)
            if record.levelno >= logging.ERROR: ida_kernwin.warning(f"[{record.name}] {log_entry}\n")
            elif record.levelno >= logging.WARNING: ida_kernwin.warning(f"[{record.name}] {log_entry}\n")
            elif record.levelno >= logging.INFO: ida_kernwin.msg(f"[{record.name}] {log_entry}\n")
            else: ida_kernwin.msg(f"[{record.name}] DEBUG: {log_entry}\n")
        except Exception as e: print(f"Error in IdaLogHandler: {e}")

log_formatter = logging.Formatter('%(levelname)s: %(message)s')
ida_handler = IdaLogHandler()
ida_handler.setFormatter(log_formatter)
core_logger.addHandler(ida_handler)
core_logger.setLevel(logging.DEBUG) # Set default level to DEBUG for more details
core_logger.propagate = False

# --- Check for Miasm Dependency and Load Unflattener ---
MIASM_AVAILABLE = True


# --- Action Handlers ---
class UnflattenActionHandler(ida_kernwin.action_handler_t):
    def __init__(self, recursive=False):  # -> Sınıfın yapıcı metodu, recursive parametresi alır
        ida_kernwin.action_handler_t.__init__(self)  # -> Üst sınıfın yapıcısını çağırır
        self.recursive = recursive  # -> Recursive işlem yapılıp yapılmayacağını saklar
        self.action_name = ACTION_UNFLATTEN_RECURSIVE if recursive else ACTION_UNFLATTEN  # -> Aksiyon ismini recursive durumuna göre ayarlar


    def set_ida_context(self,unflattener_instance, code_start_va: int, code_size: int, image_base: int):
        """Set IDA Pro segment information and calculate binary_base_va."""
        logger.debug(f"Setting IDA context: start_va=0x{code_start_va:X}, size=0x{code_size:X}, image_base=0x{image_base:X}")
        unflattener_instance.code_start_va = code_start_va
        unflattener_instance.code_size = code_size
        unflattener_instance.image_base = image_base
        # binary_base_va'yı doğru şekilde ayarla
        # IDA VA'sı ile Miasm offset'i arasındaki fark
        if hasattr(unflattener_instance.container, 'bin_stream') and hasattr(unflattener_instance.container.bin_stream, 'offset'):
          unflattener_instance.binary_base_va = unflattener_instance.code_start_va - unflattener_instance.container.bin_stream.offset
        else:
          unflattener_instance.binary_base_va = unflattener_instance.image_base
        logger.debug(f"Set binary_base_va for VA conversion: 0x{unflattener_instance.binary_base_va:X} (code_start_va: 0x{unflattener_instance.code_start_va:X})")
        logger.debug(f"Set binary_base_va for VA conversion: 0x{unflattener_instance.binary_base_va:X}")

        unflattener_instance.text_section_range['lower'] = code_start_va
        unflattener_instance.text_section_range['upper'] = code_start_va + code_size
        logger.debug(f"Text section range (approx): 0x{unflattener_instance.text_section_range['lower']:X} - 0x{unflattener_instance.text_section_range['upper']:X}")



    def get_all_bytes_from_segments(self):  # -> Tüm segmentlerden byte verilerini alan fonksiyon
        """Gets all bytes from all segments."""  # -> Açıklama
        all_bytes = b""  # -> Tüm segmentlerin birleşik byte verisi
        min_ea = idaapi.BADADDR  # -> Segmentlerin en küçük adresi, başlangıçta BADADDR
        max_ea = 0  # -> Segmentlerin en büyük adresi, başlangıçta sıfır
        for i in range(ida_segment.get_segm_qty()):  # -> Segment sayısı kadar döngü
            seg = ida_segment.getnseg(i)  # -> i. segmenti al
            if not seg:  # -> Segment yoksa
                continue  # -> Sonraki segmente geç
            seg_bytes = ida_bytes.get_bytes(seg.start_ea, seg.end_ea - seg.start_ea)  # -> Segmentin tüm byte verisini al
            if seg_bytes is None:  # -> Veri alınamazsa
                continue  # -> Sonraki segmente geç
            all_bytes += seg_bytes  # -> Segment verisini tüm verilere ekle
            if seg.start_ea < min_ea:  # -> Segmentin başlangıcı mevcut minimumdan küçükse
                min_ea = seg.start_ea  # -> Yeni minimum yap
            if seg.end_ea > max_ea:  # -> Segmentin sonu mevcut maksimumdan büyükse
                max_ea = seg.end_ea  # -> Yeni maksimum yap
        return all_bytes, min_ea, max_ea  # -> Tüm veriyi, minimum ve maksimum adresleri döner

    def get_full_input_file_bytes(self):  # -> PE dosyasının tamamını okuyan fonksiyon
        """PE dosyasının tamamını okur."""  # -> Açıklama
        import os  # -> os modülünü içeri aktar
        pe_path = ida_nalt.get_input_file_path()  # -> IDA'da açık olan dosyanın yolunu al
        if not os.path.isfile(pe_path):  # -> Dosya mevcut değilse
            raise RuntimeError(f"PE dosyası bulunamadı: {pe_path}")  # -> Hata fırlat
        with open(pe_path, 'rb') as f:  # -> Dosyayı ikili okuma modunda aç
            data = f.read()  # -> Tüm içeriği oku
        return data  # -> Okunan byte verisini döner

    def activate(self, ctx):
        """Called when the action is triggered."""
        prefix = f"{PLUGIN_NAME} ({'Recursive' if self.recursive else 'Single'}): "
        
        try:
            # Get target function
            target_address = self._get_target_function(ctx)
            if target_address == ida_idaapi.BADADDR:
                return 0

            # Show wait box
            ida_kernwin.show_wait_box(f"HIDECANCEL\n{PLUGIN_NAME}: Processing...")

            # Get all binary data
            ida_kernwin.msg(f"{prefix}Reading binary data...\n")
            binary_data = self.get_full_input_file_bytes()
            _, min_ea, max_ea = self.get_all_bytes_from_segments()
            if not binary_data:
                raise RuntimeError("Failed to read binary data from segments")

            # Create LocationDB and stream
            self.loc_db: LocationDB = LocationDB()
            stream = io.BytesIO(binary_data)
            
            # Get architecture info
            info = ida_ida.inf_get_procname().lower()
            if info.startswith("metapc"):
                arch = "x86_64" if ida_ida.inf_is_64bit() else "x86_32"
            elif info.startswith("arm"):
                arch = "aarch64" if ida_ida.inf_is_64bit() else "arm"
                 
                
            ida_kernwin.msg(f"{prefix}Architecture: {arch}\n")
            
            # Create container 
            container = Container.from_stream(stream, self.loc_db)
            
            # Create unflattener instance
            unflattener_instance = Unflattener(container, binary_data , arch=arch)
            self.set_ida_context(unflattener_instance,min_ea, len(binary_data), ida_nalt.get_imagebase())
            
            # Process function
            if not self.recursive:
                print("single_function_start")
                result = self._process_single_function(unflattener_instance, target_address)
            else:
                print("_process_recursive_start")
                result = self._process_recursive(unflattener_instance, target_address)

            return 1 if result else 0

        except Exception as e:
            core_logger.exception("Error in activate")
            ida_kernwin.warning(f"{prefix}Error: {str(e)}")
            return 0
        finally:
            ida_kernwin.hide_wait_box()

    def update(self, ctx):
        """Enable/disable action based on context."""
        # Enable only if Miasm loaded OK and we are in a suitable view/context
        if MIASM_AVAILABLE:
             widget_type = ida_kernwin.get_widget_type(ctx.widget)
             if widget_type in [ida_kernwin.BWN_DISASM, ida_kernwin.BWN_FUNCS]:
                 ea = ctx.cur_ea;
                 if ea == ida_idaapi.BADADDR: ea = ida_kernwin.get_screen_ea()
                 if ea != ida_idaapi.BADADDR and ida_funcs.get_func(ea): return ida_kernwin.AST_ENABLE_FOR_WIDGET
                 elif widget_type == ida_kernwin.BWN_FUNCS:
                     sel = ida_kernwin.get_highlight(ctx.widget);
                     if sel and sel[0]: return ida_kernwin.AST_ENABLE_FOR_WIDGET
        return ida_kernwin.AST_DISABLE_FOR_WIDGET

    def _get_target_function(self, ctx) -> int:
        """Gets the target function address from context or user selection."""
        try:
            # First try to get address from context
            if ctx.cur_func:
                return ctx.cur_func.start_ea
            
            # Try to get from current screen position
            ea = ida_kernwin.get_screen_ea()
            if ea != ida_idaapi.BADADDR:
                func = ida_funcs.get_func(ea)
                if func:
                    return func.start_ea
            
            # If no valid address found, ask user
            target = ida_kernwin.ask_addr(
                ida_idaapi.BADADDR,
                "Please enter function start address (hex)"
            )
            
            if target != ida_idaapi.BADADDR:
                func = ida_funcs.get_func(target)
                if func:
                    return func.start_ea
                else:
                    ida_kernwin.warning(f"{PLUGIN_NAME}: Address 0x{target:X} is not a function start.")
            
            return ida_idaapi.BADADDR
            
        except Exception as e:
            core_logger.error(f"Error getting target function: {e}")
            return ida_idaapi.BADADDR

    def get_interval_hull(self,func_interval):
        """
        Calculates the hull (min start, max end) from various interval representations.

        Args:
            func_interval: Can be an object with .hull(), a list/tuple of (start, end),
                        or an intervaltree Interval object.

        Returns:
            tuple[int, int] | None: (hull_start, hull_end) or None if invalid.
        """
        if hasattr(func_interval, 'hull') and callable(func_interval.hull):
            # intervaltree.IntervalTree likely
            hull = func_interval.hull()
            if hull:
                return int(hull[0]), int(hull[1]) # hull() returns [begin, end)
            else:
                logger.warning("func_interval.hull() returned empty.")
                return None
        elif hasattr(func_interval, 'begin') and hasattr(func_interval, 'end'):
            # intervaltree.Interval likely
            return int(func_interval.begin), int(func_interval.end)
        elif isinstance(func_interval, (list, tuple)):
            if not func_interval:
                logger.warning("func_interval list/tuple is empty.")
                return None
            # Check if it's a single interval [start, end] or [(start, end)]
            if len(func_interval) == 2 and isinstance(func_interval[0], int) and isinstance(func_interval[1], int):
                # Assume it's a single interval [start, end]
                return int(func_interval[0]), int(func_interval[1])
            elif all(isinstance(item, (list, tuple)) and len(item) == 2 for item in func_interval):
                # List of intervals [(start1, end1), (start2, end2), ...]
                min_start = min(int(iv[0]) for iv in func_interval)
                max_end = max(int(iv[1]) for iv in func_interval)
                return min_start, max_end
            else:
                logger.error(f"Unrecognized format within func_interval list/tuple: {func_interval}")
                return None
        else:
            logger.error(f"Cannot determine hull from func_interval type: {type(func_interval)}")
            return None


    def apply_ida_patches(self, unflattener_obj: Unflattener, target_address: int, mode_all: bool):
        """
        Temizlenmiş fonksiyon kodunu doğrudan IDA'ya yamalar. Tek veya çoklu fonksiyon
        modunu destekler. Yama öncesi temizlik, yama sonrası analiz ve fonksiyon
        tanımlama adımlarını içerir. (ida_auto kullanımı düzeltildi)

        Args:
            unflattener_obj (Unflattener): Unflattener sınıfından nesne.
            target_address (int): Yeniden inşa edilecek ilk fonksiyonun başlangıç adresi.
            mode_all (bool): True ise `unflat_follow_calls` ile erişilebilir tüm
                             flatten edilmiş fonksiyonları bulur ve yamalar. False ise
                             sadece `target_address`'teki fonksiyonu yamalar.
        """
        logger.info(f"Starting rebuild process. Target: 0x{target_address:X}, Mode All: {mode_all}")

        patch_data_to_process = [] # İşlenecek (patch_bytes, func_interval, original_ea) tuple listesi

        try:
            # 1. Yama verisini al (Bu kısım aynı kaldı)
            if not mode_all:
                logger.info(f"Attempting to unflatten single function at 0x{target_address:X}")
                patch_bytes, func_interval = unflattener_obj.unflat(target_address)
                if patch_bytes is not None and func_interval is not None:
                    patch_data_to_process = [(patch_bytes, func_interval, target_address)]
                # ... (hata/atlama logları) ...
            else:
                logger.info(f"Attempting to unflatten target 0x{target_address:X} and follow calls.")
                patch_data_list_raw = unflattener_obj.unflat_follow_calls(target_address, None)
                if patch_data_list_raw:
                    logger.info(f"Found {len(patch_data_list_raw)} potential functions via follow_calls.")
                    for item in patch_data_list_raw:
                        if isinstance(item, tuple) and len(item) == 3:
                            patch_data_to_process.append(item)
                        elif isinstance(item, tuple) and len(item) == 2:
                            logger.warning(f"Patch data missing original_ea, using None placeholder.")
                            patch_data_to_process.append((item[0], item[1], None))
                        else:
                            logger.error(f"Invalid item format from unflat_follow_calls: {item}. Skipping.")
                else:
                    logger.info(f"No functions found via follow_calls from 0x{target_address:X}.")
                    print(f"No functions found via follow_calls from 0x{target_address:X}.")

        except Exception as e:
            logger.error(f"Error during unflattening phase (target 0x{target_address:X}): {e}", exc_info=True)
            print(f"Error during unflattening phase for 0x{target_address:X}: {e}")
            return

        if not patch_data_to_process:
            logger.info("No valid patch data found. Exiting.")
            print("No patches to apply.")
            return
        
        num_functions_to_patch = len(patch_data_to_process)

        confirmation_message = (f"Found {num_functions_to_patch} function(s) to patch.\n"
                                    f"This will modify the database by undefining code, "
                                    f"filling with NOPs, applying patches, and re-analyzing.\n\n"
                                    f"Proceed with patching?")

        logger.info(f"Requesting user confirmation to patch {num_functions_to_patch} function(s).")
        # ask_yn: 1 = Yes, 0 = No, -1 = Cancel/Close
        # Varsayılan olarak 'No' seçili olsun (0)
        user_choice = ida_kernwin.ask_yn(0, confirmation_message)

        if user_choice != 1: # Eğer kullanıcı 'Yes' demediyse
            logger.info("User cancelled the patching operation.")
            print("Patching operation cancelled by user.")
            return # Fonksiyondan çık


        # --- Her fonksiyonun yama verisini işle ---
        total_patched_bytes = 0
        functions_patched_count = 0
        functions_failed_count = 0

        for patch_bytes, func_interval, original_ea in patch_data_to_process:

            current_target_desc = f"function (start unknown)"
            effective_start_addr = None
            if original_ea:
                current_target_desc = f"function originally at 0x{original_ea:X}"
                effective_start_addr = original_ea

            logger.info(f"--- Processing patch for {current_target_desc} ---")

            if not patch_bytes or not func_interval:
                 logger.warning(f"Skipping invalid data for {current_target_desc}.")
                 functions_failed_count += 1
                 continue

            # 2. Hull belirle (Bu kısım aynı kaldı)
            hull = self.get_interval_hull(func_interval) # Placeholder - Gerçek kodu kullanın
            if hull is None:
                 logger.error(f"Could not determine hull for {current_target_desc}. Trying fallback.")
                 # ... (Fallback logic) ...
                 func = None
                 if original_ea: func = ida_funcs.get_func(original_ea)
                 if func:
                     hull = (func.start_ea, func.end_ea)
                     if effective_start_addr is None: effective_start_addr = func.start_ea
                     logger.warning(f"Using IDA bounds fallback: 0x{hull[0]:X}-0x{hull[1]:X}")
                     print(f"Warning: Using IDA bounds fallback for {current_target_desc}")
                 else:
                     logger.error(f"Fallback failed for {current_target_desc}. Cannot patch.")
                     print(f"Error: Fallback failed for {current_target_desc}. Cannot patch.")
                     functions_failed_count += 1
                     continue

            hull_start, hull_end = hull
            hull_length = hull_end - hull_start

            if effective_start_addr is None: # Hala başlangıç adresi yoksa hull başını kullan
                effective_start_addr = hull_start
                logger.warning(f"Using hull start 0x{hull_start:X} as effective start for {current_target_desc}")

            if hull_length <= 0:
                logger.error(f"Invalid hull range for {current_target_desc}. Skipping.")
                print(f"Error: Invalid hull range for {current_target_desc}. Skipping.")
                functions_failed_count += 1
                continue

            logger.info(f"Original hull: 0x{hull_start:X}-0x{hull_end:X}. Effective start: 0x{effective_start_addr:X}")


            # 3. Tanımları Kaldır, NOP'la ve Yama Yap (Bu kısım aynı kaldı)
            patch_len = len(patch_bytes)
            patching_succeeded = False
            try:
                logger.info(f"Undefining items in hull 0x{hull_start:X}-0x{hull_end:X}")
                ida_bytes.del_items(hull_start, ida_bytes.DELIT_SIMPLE, hull_length)

                nop_fill = b"\xCC" * hull_length
                ida_bytes.patch_bytes(hull_start, nop_fill)
                logger.info(f"Filled hull 0x{hull_start:X}-0x{hull_end:X} with 0xCC.")

                ida_bytes.patch_bytes(effective_start_addr, patch_bytes)
                hex_addr = f"0x{effective_start_addr:X}"
                # ... (Patch byte loglama) ...
                logger.info(f"Patched {patch_len} bytes at {hex_addr} for {current_target_desc}.")
                print(f"Patched {patch_len} bytes at {hex_addr} for {current_target_desc}")
                patching_succeeded = True

            except Exception as e:
                 logger.error(f"Error during patching phase for {current_target_desc}: {e}", exc_info=True)
                 print(f"Error during patching phase for {current_target_desc}: {e}")
                 functions_failed_count += 1
                 continue

            # 4. Yamalanan Alanı Yeniden Analiz Et ve Fonksiyonu Tanımla (DÜZELTİLDİ)
            analysis_succeeded = False
            if patching_succeeded:
                try:
                    # Adım 4a: Yamalanan alanı analiz kuyruğuna ekle
                    analysis_start = effective_start_addr
                    analysis_end = effective_start_addr + patch_len
                    logger.info(f"Requesting analysis for range 0x{analysis_start:X} - 0x{analysis_end:X}")
                    ida_auto.plan_range(analysis_start, analysis_end)

                    # Adım 4b: Analiz kuyruğunun işlenmesini bekle
                    logger.info("Waiting for autoanalysis to complete...")
                    ida_auto.auto_wait() # Bu fonksiyon analiz bitene kadar bekler
                    # auto_wait doğrudan başarı durumu döndürmez,
                    # eğer hata olmadan biterse başarılı varsayılır.
                    logger.info("Autoanalysis finished.")
                    analysis_succeeded = True # Hata fırlatmadıysa başarılı sayalım

                    # Adım 4c: Fonksiyonu yeniden tanımlamayı dene (analiz başarılıysa)
                    if analysis_succeeded:
                        # Eski tanımı sil (varsa ve farklıysa)
                        if original_ea and original_ea != analysis_start:
                            try:
                                if ida_funcs.del_func(original_ea): logger.info(f"Deleted old func at 0x{original_ea:X}")
                            except: pass # Hata olsa da devam et
                        elif ida_funcs.get_func(analysis_start): # Yeni adreste zaten varsa
                             try:
                                if ida_funcs.del_func(analysis_start): logger.info(f"Deleted existing func at 0x{analysis_start:X}")
                             except: pass

                        # Yeni fonksiyonu ekle
                        logger.info(f"Attempting to define function at 0x{analysis_start:X}")
                        try:
                            if ida_funcs.add_func(analysis_start, idc.BADADDR):
                                logger.info(f"Successfully defined function start at 0x{analysis_start:X}")
                            else:
                                logger.warning(f"Could not add function definition at 0x{analysis_start:X}")
                        except Exception as add_e:
                             logger.error(f"Error adding function at 0x{analysis_start:X}: {add_e}")
                    else:
                         # Bu durum auto_wait hata verirse veya ileride kontrol eklenirse olabilir.
                         logger.warning(f"Analysis was not considered successful for 0x{analysis_start:X}, skipping function definition.")


                except Exception as e:
                    logger.error(f"Error during analysis/definition phase for {current_target_desc}: {e}", exc_info=True)
                    print(f"Error during analysis/definition for {current_target_desc}: {e}")
                    # Analiz/tanımlama başarısız olsa bile patch uygulanmış olabilir.

            # 5. Bu fonksiyon için özeti logla (Bu kısım aynı kaldı)
            if patching_succeeded:
                if analysis_succeeded:
                    logger.info(f"Successfully patched and analyzed {current_target_desc}. Patch size: {patch_len} bytes.")
                    functions_patched_count += 1
                else:
                    logger.warning(f"Patched {current_target_desc} (size {patch_len}), but analysis failed/skipped.")
                    functions_patched_count += 1 # Yama yapıldıysa yine de başarılı sayabiliriz

                total_patched_bytes += patch_len
                # ... (Patch boyutu ve hull karşılaştırma logları) ...

            # --- Tek bir fonksiyon için döngü sonu ---

        # --- Genel Özet ---
        logger.info("="*40)
        logger.info(f"Patching process finished.")
        logger.info(f"Successfully patched functions: {functions_patched_count}")
        if functions_failed_count > 0:
             logger.warning(f"Failed/Skipped functions: {functions_failed_count}")
        logger.info(f"Total bytes written in patches: {total_patched_bytes}")
        print("="*40)
        print(f"Patching process finished.")
        print(f"Successfully patched: {functions_patched_count}")
        if functions_failed_count > 0:
             print(f"Failed/Skipped: {functions_failed_count}")
        print(f"Total patch bytes written: {total_patched_bytes}")
        logger.info("="*40)


    def _process_single_function(self, unflattener, target_address: int) -> bool:
        """Process a single function."""
        try:
            self.apply_ida_patches(unflattener, target_address,False)
            return True
        except Exception as e:
            core_logger.error(f"Error processing single function at 0x{target_address:X}: {e}")
            return False

    def _process_recursive(self, unflattener, target_address: int) -> bool:
        """Process function recursively including called functions."""
        try:
            # Özyinelemeli işlem için doğru patch toplama fonksiyonunu çağır
            # unflat_recursive_collect_patches, (patch_dict, interval) formatında liste döndürür.
            self.apply_ida_patches(unflattener,target_address,True)
            return True
        except Exception as e:
            core_logger.error(f"Error processing function recursively at 0x{target_address:X}: {e}")
            return False





# --- Hooks for Pop-up Menu ---
class GuiHooks(ida_kernwin.UI_Hooks):
    # ... (Function content identical to previous correct version) ...
    def finish_populating_widget_popup(self, widget, popup_handle):
        widget_type = ida_kernwin.get_widget_type(widget); menu_path = f"{PLUGIN_NAME}/"
        if widget_type in [ida_kernwin.BWN_DISASM, ida_kernwin.BWN_FUNCS]:
            if MIASM_AVAILABLE:
                ida_kernwin.attach_action_to_popup(widget, popup_handle, ACTION_UNFLATTEN, menu_path)
                ida_kernwin.attach_action_to_popup(widget, popup_handle, ACTION_UNFLATTEN_RECURSIVE, menu_path)


# --- Plugin Class ---
class MiasmUnflattenerPlugin(ida_idaapi.plugin_t):
    # ... (flags, comment, help, wanted_name, wanted_hotkey same) ...
    flags = ida_idaapi.PLUGIN_MOD | ida_idaapi.PLUGIN_PROC
    comment = PLUGIN_COMMENT; help = PLUGIN_HELP; wanted_name = PLUGIN_NAME; wanted_hotkey = WANTED_HOTKEY
    def __init__(self): self.hooks = None

    def init(self):
        """Called when IDA loads the plugin."""
        # ... (Function content identical to previous correct version) ...
        prefix = f"{PLUGIN_NAME} init v{PLUGIN_VERSION}: "; ida_kernwin.msg(f"\n{prefix}Initializing...\n")
        log_level_str = os.environ.get("MIASM_UNFLATTENER_LOGLEVEL", "INFO").upper(); log_level = getattr(logging, log_level_str, logging.INFO)
        core_logger.setLevel(logging.DEBUG)  # Set to DEBUG for more detailed output
        ida_kernwin.msg(f"  Core log level: {logging.getLevelName(core_logger.level)}\n")
        if not MIASM_AVAILABLE: ida_kernwin.warning(f"{prefix}Plugin disabled (missing dependencies).\n"); return ida_idaapi.PLUGIN_OK
        actions = [ ida_kernwin.action_desc_t(ACTION_UNFLATTEN, 'Unflatten Function', UnflattenActionHandler(False), WANTED_HOTKEY or None, 'Unflatten function using Miasm', 199),
                    ida_kernwin.action_desc_t(ACTION_UNFLATTEN_RECURSIVE, 'Unflatten Function (Recursive)', UnflattenActionHandler(True), None, 'Unflatten function and calls using Miasm', 199) ]
        register_ok = all(ida_kernwin.register_action(action) for action in actions)
        if not register_ok: ida_kernwin.warning(f"{prefix}Action registration failed. Skipping load.\n"); self.term(); return ida_idaapi.PLUGIN_SKIP
        try:
            self.hooks = GuiHooks()
            if not self.hooks.hook(): raise RuntimeError("UI Hook installation failed.")
            ida_kernwin.msg(f"{prefix}UI hooks installed.\n")
        except Exception as hook_ex:
             ida_kernwin.warning(f"{prefix}UI hook exception: {hook_ex}. Cleaning up.\n"); core_logger.exception("UI Hook failed")
             if self.hooks: self.hooks.unhook(); self.hooks = None
             self.term(); return ida_idaapi.PLUGIN_SKIP
        ida_kernwin.msg(f"{prefix}Initialization complete.\n"); return ida_idaapi.PLUGIN_KEEP

    def run(self, arg):
        """Called from Edit/Plugins menu."""
        # ... (Function content identical to previous correct version) ...
        ida_kernwin.info(f"{PLUGIN_NAME} v{PLUGIN_VERSION}\nUse right-click menus.")

    def term(self):
        """Called when IDA unloads the plugin."""
        # ... (Function content identical to previous correct version) ...
        prefix = f"{PLUGIN_NAME} term: "; ida_kernwin.msg(f"\n{prefix}Unloading...\n")
        if self.hooks:
            try: self.hooks.unhook(); # ida_kernwin.msg(f"{prefix}UI hooks removed.\n")
            except Exception: pass
            self.hooks = None
        for action_name in [ACTION_UNFLATTEN, ACTION_UNFLATTEN_RECURSIVE]: ida_kernwin.unregister_action(action_name)
        plugin_dir = os.path.dirname(__file__)
        if plugin_dir in sys.path:
            try: sys.path.remove(plugin_dir)
            except ValueError: pass
        global core_logger, ida_handler
        if core_logger and ida_handler:
            try: core_logger.removeHandler(ida_handler); ida_handler = None
            except Exception: pass
        ida_kernwin.msg(f"{prefix}Unloaded.\n")

# --- Plugin Entry Point ---
def PLUGIN_ENTRY():
    return MiasmUnflattenerPlugin()


repository: Ga-ryo/IDAFuzzy
file: ida_fuzzy.py
description: Fuzzy search tool for IDA Pro.
url: https://github.com/Ga-ryo/IDAFuzzy/blob/afd3b34d1fbd2a389f9975de83d5ab46f78aedb6/ida_fuzzy.py
wanted_name: fuzzy search
comment: Fuzzy search everything for IDA
created_at: 2018-04-10 08:16:58+00:00
pushed_at: 2019-12-16 09:03:50+00:00
content:
from idaapi import Form
from idaapi import Choose
import idaapi
from ida_kernwin import *
from fuzzywuzzy import process, fuzz
from idautils import *
import gc
import functools
from PyQt5 import QtCore
import threading

"""
Fuzzy Search v1.0
Goal
1. Search and execute IDA Pro's feature by name(ex: file,next code, run, attach to process ... )
2. Search and goto Function, string, struct,...
3. Automatically update. (when user rename function, hook and refresh)

Choose.CH_QFTYP_FUZZY is not so usable.
1. Not so fuzzy.
2. In the first place, fuzzy choose isn't applied to Functions Window or other embedded chooser.

@TODO
1. Installation
 - install idapython
 - pip install fuzzywuzzy
 - put this file to plugins directory.

2. Usage
3. Implement
 - All feature
 - Functions (hook rename and reload automatically)
 - Strings (symbol and Contents)
 - Structures
 - etc...

4. Show hint?
 - Name = "strings windows", Hint = "Open strings subview in current context."
  -- but add column affects number of pushing tab.
"""

LISTLEN = 10

class Commands(object):
    """
    Command execution proxy.
    """

    def __init__(self, **kwargs):
        self.kwargs = kwargs
        assert (callable(kwargs['fptr']))
        # assert(kwargs.get('description') != None)

    @property
    def description(self):
        return self.kwargs.get('description')

    def execute(self):
        # ea = get_screen_ea()
        # open_strings_window(ea)
        if self.kwargs.get('args') is not None:
            self.kwargs.get('fptr')(*self.kwargs.get('args'))
        else:
            self.kwargs.get('fptr')()

    def get_icon(self):
        if self.kwargs.get('icon') is None:
            return 0
        return self.kwargs.get('icon')


choices = {}
names = []


class EmbeddedChooserClass(Choose):
    """
    A simple chooser to be used as an embedded chooser
    """

    def __init__(self, title, nb=5, flags=0):
        Choose.__init__(self,
                        title,
                        [["Action", 30 | Choose.CHCOL_PLAIN]],
                        embedded=True, height=10, flags=flags)
        # embedded=True, width=30, height=20, flags=flags)

        self.n = 0
        self.items = []
        self.icon = 0

    def OnGetIcon(self, n):
        # print("get icon %d" % n)
        return choices[self.items[n][0]].get_icon()

    def OnSelectionChange(self, n):
        pass
        # print("selection change %d" % n)

    def OnClose(self):
        pass

    def OnGetLine(self, n):
        # print("getline %d" % n)
        return self.items[n]

    def OnGetSize(self):
        n = len(self.items)
        # print("getsize -> %d" % n)
        return n

"""
def hooked_scorer(*args,**kwargs):
    # watch event every time.
    pass

def hook(function, prefunction):
    @functools.wraps(function)
    def run(*args, **kwargs):
        prefunction(*args, **kwargs)
        return function(*args, **kwargs)
    return run
fuzz.WRatio = hook(fuzz.WRatio, hooked_scorer)

"""

class TerminateException(Exception):
    pass

def hooked_scorer(*args, **kwargs):
    if kwargs.pop('terminate_event').is_set():
        raise TerminateException
    return fuzz.WRatio(*args, **kwargs)

class FuzzySearchThread(QtCore.QThread):
    refresh_list = QtCore.pyqtSignal([str]*LISTLEN)

    def __init__(self, parent=None):
        super(FuzzySearchThread, self).__init__(parent)
        self.stopped = False
        self.mutex = QtCore.QMutex()
        self.terminate_event = threading.Event()

    def setup(self,  s):
        self.stoppped = False
        self.s = s

    def stop( self ):
        with QtCore.QMutexLocker( self.mutex ):
            self.stopped = True

    def run(self):
        f = functools.partial(hooked_scorer, terminate_event=self.terminate_event)
        try:
            res = process.extract(self.s, names, limit=LISTLEN, scorer=f)  # f.iStr1.value won't change until Form.Execute() returns.
            extracts = []
            for i in res:
                extracts.append(i[0])
            for i in range(10-len(res)):
                extracts.append("")
            self.refresh_list.emit(*extracts)  # call main Thread's UI function.
        except TerminateException:
            pass
        self.stop()
        self.finished.emit()


# --------------------------------------------------------------------------
class FuzzySearchForm(Form):
    def __init__(self):
        self.invert = False
        self.EChooser = EmbeddedChooserClass("Title", flags=Choose.CH_MODAL)
        self.selected_id = 0
        self.s = ""
        self.fst = FuzzySearchThread()
        self.fst.refresh_list.connect(self.refresh_list)
        self.fst.finished.connect(self.finished)
        # self.EChooser = EmbeddedChooserClass("Title", flags=Choose.CH_CAN_REFRESH)

	# Portability fix from Python2 to Python3.
        try:
            self.cEChooser = super().cEChooser #super() will raise exception in python2
        except:
            pass

        Form.__init__(self, r"""STARTITEM 
        IDA Fuzzy Search
        {FormChangeCb}
        <:{iStr1}>

        <Results:{cEChooser}>
""", {
            'iStr1': Form.StringInput(),
            'cEChooser': Form.EmbeddedChooserControl(self.EChooser),
            'FormChangeCb': Form.FormChangeCb(self.OnFormChange),
        })
        # self.modal = False

    def OnFormChange(self, fid):
        if fid == -1:
            # initialize
            pass
        elif fid == -2:
            # terminate
            pass
        elif fid == self.cEChooser.id:
            self.selected_id = self.GetControlValue(self.cEChooser)[0]
        elif fid == self.iStr1.id:
            self.s = self.GetControlValue(self.iStr1)
            self.EChooser.items = []
            if self.s == '':
                self.RefreshField(self.cEChooser)
                return 1
            self.fst.stop()
            self.fst.quit()  #  if you type speedy, FuzzySearch which executed before is not finished here.
            self.fst.terminate_event.set()
            self.fst.wait()
            #self.fst.terminate()  # but last time's FuzzySearch is meaningless, so terminate this. <- little dangerous?

            #stop and quit take time.(and maybe non-blocking)
            #So if you type speedy, some start() call will be ignored.
            #re-create thread solve this.
            self.fst = FuzzySearchThread()
            self.fst.refresh_list.connect(self.refresh_list)
            self.fst.finished.connect(self.finished)
            self.fst.setup(self.s)
            self.fst.start()

            # extracts = process.extract(s, names, limit=10)  # f.iStr1.value won't change until Form.Execute() returns.
        else:
            pass
        return 1

    def refresh_list(self, *extracts):
        for ex in extracts:
            # self.EChooser.items.append([ex[0], choices[ex[0]].description])
            self.EChooser.items.append([ex])
        self.RefreshField(self.cEChooser)
        self.SetControlValue(self.cEChooser, [0])  # set cursor top

    def finished(self):
        pass

    def get_selected_item(self):
        if self.selected_id == -1:
            return None
        item_name = self.EChooser.items[self.selected_id][0]
        return choices[item_name]


# --------------------------------------------------------------------------
def fuzzy_search_main():
    # Create form
    global f, choices, names
    choices = {}
    names = []
    gc.collect()

    # Runtime collector.
    # Pros
    # 1. No need to refresh automatically.(When GDB start, libc symbol,PIE's symbol,etc... address will change.When user rename symbol, also.)
    # 1.1. If you want to search library's function, view module list and right-click onto target library. Then click "Analyze module".
    # 2. Action's state is collect (When user start typing, active window is FuzzySearchForm. So filter doesn't works correctly. ex: OpHex is active at Disas view but not active at FuzzySearchForm.)
    # Cons
    # 1. Become slow in case large file.
    # 1.1. Re-generate dictionary isn't matter.(But scoring time will be bigger.)
    # func ptr and icon id
    registered_actions = get_registered_actions()
    for action in registered_actions:
        # IDA's bug? tilde exists many times in label. ex) Abort -> ~A~bort
        # So fix it.
        label = get_action_label(action).replace('~', '')
        icon = get_action_icon(action)[1]
        desctription = get_action_tooltip(action)
        if get_action_state(action)[1] > idaapi.AST_ENABLE:
            continue
        choices[label] = Commands(fptr=process_ui_action, args=[action], description=desctription, icon=icon)

    # Functions()
    # Heads()
    for n in Names():
        demangled = idc.demangle_name(n[1], idc.get_inf_attr(idc.INF_SHORT_DN))
        name = demangled if demangled else n[1]
        # jump to addr
        choices[name] = Commands(fptr=jumpto, args=[n[0]], description="Jump to " + name, icon=124)

    for n in Structs():
        choices[n[2]] = Commands(fptr=open_structs_window, args=[n[1]],
                                 description="Jump to Structure definition of " + n[2], icon=52)

    for k, v in choices.items():
        names.append(k)

    f = FuzzySearchForm()

    # Compile (in order to populate the controls)
    f.Compile()
    f.iStr1.value = ""
    # Execute the form
    ok = f.Execute()

    if ok == 1 and len(f.EChooser.items) > 0:
        f.get_selected_item().execute()
    # Dispose the form
    f.Free()


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

    def activate(self, ctx):
        action = fuzzy_search_main()

        if action:
            idaapi.process_ui_action(action)
        return 1

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


class FuzzySearchPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_FIX | idaapi.PLUGIN_HIDE
    comment = "Fuzzy search everything for IDA"
    help = "Fuzzy search everything"
    wanted_name = "fuzzy search"
    wanted_hotkey = ""

    def init(self):
        print("Fuzzy Search Plugin loaded.")
        idaapi.register_action(
            idaapi.action_desc_t("fz:fuzzysearch", "Fuzzy Search", fuzzy_search_handler(), "Shift+SPACE", "", -1))

        return idaapi.PLUGIN_KEEP

    def term(self):
        idaapi.unregister_action("fz:fuzzysearch")
        pass

    def run(self, arg):
        pass


def PLUGIN_ENTRY():
    return FuzzySearchPlugin()


repository: wlingze/golang_rev
file: golang_rev.py
description: A simple golang reverse engineering ida plugin
url: https://github.com/wlingze/golang_rev/blob/d5d6c16086642581d1f3890dd42efe81588f8a68/golang_rev.py
wanted_name: Golang_Rev
comment: Golang_Rev
created_at: 2020-07-22 09:55:01+00:00
pushed_at: 2020-12-16 11:40:56+00:00
content:
#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2020 wlz <wlz@kyria>
#
# Distributed under terms of the MIT license.

import idc
import idaapi
import idautils
import ida_bytes
import ida_name

GOLANG_FUNC = "golang_rev:golang_func"
GOLANG_STRING = "golang_rev:golang_string"
RENAME_POINTER = 'golang_rev:rename_pointer'

if idaapi.IDA_SDK_VERSION < 740:
    Qword = idc.Qword
    MakeQword = idc.MakeQword
    MakeStr = idc.MakeStr
    MakeName = idc.MakeName
    Name = idc.Name
    Dword = idc.Dword
    GetString = idc.GetString
else:
    Qword = ida_bytes.get_qword
    Dword = ida_bytes.get_dword
    MakeQword = idc.create_qword
    MakeStr = idc.create_strlit

    def MakeName(ea, name):
        return idc.set_name(ea, name, ida_name.SN_CHECK)

    def Name(ea):
        return idc.get_name(ea, ida_name.GN_VISIBLE)

    def GetString(ea, length=-1, strtype=0):
        return idc.get_strlit_contents(ea, length, strtype).decode()


class menu_action_handler_t(idaapi.action_handler_t):
    """
    Action handler for menu actions
    """

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

    def activate(self, ctx):
        if self.action == GOLANG_STRING:
            self.set_string()
        elif self.action == GOLANG_FUNC:
            self.set_function()
        elif self.action == RENAME_POINTER:
            self.rename_pointer()
        else:
            return 0
        return 1

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

    def set_string(self):
        start = idc.here()
        str_len = idaapi.ask_long(0, "string len...")
        if str_len:
            MakeStr(start, start + str_len)
        return 1

    def rename_pointer(self):
        name = 'p_{}'.format(Name(Qword(idc.here())))
        MakeName(idc.here(), name)

    def set_function(self):
        try:
            base = idautils.ida_segment.get_segm_by_name('.gopclntab').start_ea
        except:
            print("Cannot get segment .gopclntab!")
            return 0
        ea = base + 8
        len = Qword(ea) * 8 * 2
        ptr = base + 8 + 8
        end = base + len
        while ptr <= end:
            MakeQword(ptr)
            func_addr = Qword(ptr)
            MakeQword(ptr + 8)
            name_offset = Qword(ptr + 8)
            name_addr = Dword(base + 8 + name_offset) + base
            name = GetString(name_addr)
            name = name.replace('.', '_') \
                .replace("<-", '_chan_left_') \
                .replace('*', '_ptr_') \
                .replace('-', '_') \
                .replace(';', '') \
                .replace('"', '') \
                .replace('\\\\', '')
            name = name.replace('(', '') \
                .replace(')', '') \
                .replace('/', '_') \
                .replace(' ', '_') \
                .replace(',', 'comma') \
                .replace('{', '') \
                .replace('}', '')
            MakeName(func_addr, name)
            print(name)
            print(ptr)
            ptr += 16
        idc.jumpto(idc.get_name_ea_simple('main_main'))
        return 1


class UI_Hook(idaapi.UI_Hooks):
    def __init__(self):
        idaapi.UI_Hooks.__init__(self)

    def finish_populating_widget_popup(self, form, popup):
        form_type = idaapi.get_widget_type(form)

        if form_type == idaapi.BWN_DISASM or form_type == idaapi.BWN_DUMP:
            t0, t1, view = idaapi.twinpos_t(), idaapi.twinpos_t(
            ), idaapi.get_current_viewer()
            if idaapi.read_selection(view, t0, t1) \
                    or idc.get_item_size(idc.get_screen_ea()) > 1:
                idaapi.attach_action_to_popup(form, popup, GOLANG_FUNC, None)
                idaapi.attach_action_to_popup(form, popup, GOLANG_STRING, None)
                idaapi.attach_action_to_popup(form, popup, RENAME_POINTER,
                                              None)


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

    help = ""
    wanted_name = "Golang_Rev"
    wanted_hotkey = ""

    def init(self):
        self.hexrays_inited = False
        self.registered_actions = []
        self.registered_hx_actions = []

        global ARCH
        global BITS
        ARCH = idaapi.ph_get_id()
        info = idaapi.get_inf_structure()
        if info.is_64bit():
            BITS = 64
        elif info.is_32bit():
            BITS = 32
        else:
            BITS = 16

        print("Golang_Rev plugin has been loaded.")

        # Register menu actions
        menu_actions = (
            idaapi.action_desc_t(GOLANG_STRING, "Set golang string",
                                 menu_action_handler_t(GOLANG_STRING), None,
                                 None, 9),
            idaapi.action_desc_t(GOLANG_FUNC, "Set golang function",
                                 menu_action_handler_t(GOLANG_FUNC), None,
                                 None, 9),
            idaapi.action_desc_t(RENAME_POINTER, "Set pointer",
                                 menu_action_handler_t(RENAME_POINTER), None,
                                 None, 9),
        )
        for action in menu_actions:
            idaapi.register_action(action)
            self.registered_actions.append(action.name)

        # Add ui hook
        self.ui_hook = UI_Hook()
        self.ui_hook.hook()

        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        pass

    def term(self):
        if hasattr(self, "ui_hook"):
            self.ui_hook.unhook()

        # Unregister actions
        for action in self.registered_actions:
            idaapi.unregister_action(action)


def PLUGIN_ENTRY():
    return Golang_Rev()


repository: harelon/GoFastAnalyzer
file: go_analyzer.py
description: Go fastcall analysis for ida decompiler
wanted_name: GoAnalyzer
created_at: 2022-09-17 20:04:03+00:00
pushed_at: 2025-02-26 22:44:57+00:00
url:
https://github.com/harelon/GoFastAnalyzer/blob/4766e86d0bcead38491be3b88d5e4de20e72e2d0/go_analyzer.py

content:
import re
import json
import itertools
from dataclasses import dataclass
from collections import defaultdict

import ida_ua
import ida_ida
import ida_idp
import ida_gdl
import ida_nalt
import ida_name
import ida_xref
import ida_bytes
import ida_funcs
import ida_allins
import ida_idaapi
import ida_struct
import ida_hexrays
import ida_netnode
import ida_segment
import ida_segregs
import ida_typeinf
from idc import BADADDR

from GoAnalyzer.UniqueGoCalls import known_functions
from GoAnalyzer.GoCallinfo import GoCall, get_sized_register_by_name
from GoAnalyzer.utils import (
    BYTE_SIZE,
    go_fast_convention,
    GO_SUPPORTED,
    runtime_morestack_functions,
    remove_tinfo_zeroes,
    go_calling_convention,
)

go_version_regex = re.compile("go\\d\\.(\\d{1,2})(\\.\\d{1,2})?$")
concat_string_number = re.compile("runtime\\.concatstring\\d")


@dataclass
class GoString:
    pointer: int = 0
    length: int = 0


class Xmm15Optimizer(ida_hexrays.microcode_filter_t):
    """
    According to the go internal abi xmm15 stores 16 bytes of zero,
    we want to optimize the decompiler output to show 16 zeros instead of saying the value of xmm15 is undefined
    """

    xmm15_code = ida_idp.str2reg("xmm15")

    def apply(self, cdg: ida_hexrays.codegen_t) -> int:
        """
        Define moving of xmm15 into the stack as a 16 bytes zero
        as specified in the go abi
        """
        l_reg = cdg.load_effective_address(0)

        # initialize micro operands
        off = ida_hexrays.mop_t(l_reg, 8)
        sel_reg = ida_hexrays.reg2mreg(ida_segregs.R_ss)
        sel = ida_hexrays.mop_t(sel_reg, 2)
        zero = ida_hexrays.mop_t()
        zero.make_number(0, 8)

        # this puts zero at the first 8 bytes of the memory
        cdg.emit(ida_hexrays.m_stx, zero, sel, off)

        # add 8 to the offset
        offset_8 = ida_hexrays.mop_t()
        offset_8.make_number(8, 8)
        cdg.emit(ida_hexrays.m_add, offset_8, off, off)

        # put zero at the second 8 bytes of the memory
        cdg.emit(ida_hexrays.m_stx, zero, sel, off)
        return ida_hexrays.MERR_OK

    def match(self, cdg: ida_hexrays.codegen_t) -> bool:
        """check if the second operand is xmm15"""
        if cdg.insn.itype in [ida_allins.NN_movups, ida_allins.NN_movdqu]:
            return cdg.insn.Op2.reg == self.xmm15_code


class FunctionStringVisitor(ida_hexrays.ctree_visitor_t):
    def __init__(
        self,
        cfunc: ida_hexrays.citem_t,
        func: ida_funcs.func_t,
        string_tinfo: ida_typeinf.tinfo_t,
    ) -> None:
        ida_hexrays.ctree_visitor_t.__init__(self, ida_hexrays.CV_PARENTS)
        self.cfunc = cfunc
        self.func = func
        self.string_tinfo = string_tinfo
        self.string_assignments: dict[str, GoString] = defaultdict(GoString)

    def define_string(self, string_ref: int, string_size: int) -> None:
        """Define correctly a string constant"""
        # string is longer than the segment it is contained in, probably false
        if ida_segment.getseg(string_ref).size() < string_size:
            return
        # no need to define a string, the reference is to a struct already
        if ida_bytes.get_flags(string_ref) & ida_bytes.FF_STRUCT == ida_bytes.FF_STRUCT:
            return

        iterator = ida_bytes.get_bytes(string_ref, string_size)
        if iterator is None:
            return

        # if the string is not decodable don't define it
        try:
            decoded = iterator.decode("utf-8")
        except UnicodeError:
            return

        # undefine the string
        ida_bytes.del_items(string_ref, 0, string_size)
        # define the string correctly
        ida_bytes.create_strlit(string_ref, string_size, ida_nalt.STRTYPE_C)

    def instruction_stores_valid_size(
        self, insn: ida_ua.insn_t, size_reg: int
    ) -> bool | None:
        """
        Check if the instruction stores a valid size for our searched string,
        if data is stored to our wanted register not in the expected format we know it's ruined,
        if the instruction is for call, all registers are changed
        """
        # the string size is stored as a dword in the register, and for some reason its phrase is void and not imm
        if (
            insn.Op2.dtype == ida_ua.dt_dword
            and insn.Op2.phrase == ida_ua.o_void
            and insn.Op1.reg == size_reg
        ):
            return True
        # if the string is used and not for our purpose then it probably won't be used for the string size
        elif insn.Op1.reg == size_reg:
            return None
        # if we arrive at a call everything changes, we can't trust registers after it
        elif insn.get_canon_feature() & ida_idp.CF_CALL != 0:
            return None
        return False

    def instruction_stores_valid_obj(
        self, insn: ida_ua.insn_t, obj_reg: int
    ) -> bool | None:
        """
        Check if the instruction stores a valid size for our searched string,
        if data is stored to our wanted register not in the expected format we know it's ruined,
        if the instruction is for call, all registers are changed
        """
        # the string size is stored as a dword in the register, and for some reason its phrase is void and not imm
        if (
            insn.Op2.dtype == ida_ua.dt_qword
            and insn.Op2.phrase == ida_ua.o_imm
            and insn.Op1.reg == obj_reg
        ):
            return True
        # if the string is used and not for our purpose then it probably won't be used for the string size
        elif insn.Op1.reg == obj_reg:
            return None
        # if we arrive at a call everything changes, we can't trust registers after it
        elif insn.get_canon_feature() & ida_idp.CF_CALL != 0:
            return None
        return False

    def visit_insn(self, e: ida_hexrays.cinsn_t) -> int:
        if e.op != ida_hexrays.cit_block:
            return 0

        for expr in e.cblock:
            if expr.op != ida_hexrays.cit_expr:
                continue
            current_expr = expr.cexpr
            if current_expr.op != ida_hexrays.cot_asg:
                continue

            # we are referencing a member of a type
            if (
                current_expr.x.op != ida_hexrays.cot_memref
                or (
                    current_expr.x.x.op != ida_hexrays.cot_var
                    and current_expr.x.x.op != ida_hexrays.cot_memptr
                    and current_expr.x.x.op != ida_hexrays.cot_idx
                )
                or current_expr.x.x.type.compare(self.string_tinfo)
            ):
                continue

            member_ref = current_expr.x
            member_parent_ref = current_expr.x.x

            string_key = member_parent_ref.dstr()

            # get member offset check if it is the ptr or the length
            member = ida_typeinf.udt_member_t()
            member.offset = member_ref.m * BYTE_SIZE
            self.string_tinfo.find_udt_member(member, ida_typeinf.STRMEM_OFFSET)

            # skip cast
            skip_cast_y = (
                current_expr.y.x
                if current_expr.y.op == ida_hexrays.cot_cast
                else current_expr.y
            )

            should_continue = False
            # save the offset of the member to be able to restore our strings
            if member.name == "ptr":
                obj_ea = None
                skip_ref_y = (
                    skip_cast_y.x
                    if skip_cast_y.op == ida_hexrays.cot_ref
                    else skip_cast_y
                )

                if skip_ref_y.op == ida_hexrays.cot_idx:
                    e = skip_ref_y
                    while e.ea == BADADDR:
                        e = self.cfunc.find_parent_of(e)

                    # we found the parent instruction
                    suspected_string_ref_ea = e.ea
                    insn = ida_ua.insn_t()

                    result = ida_xref.get_first_dref_from(suspected_string_ref_ea)

                    # find the string size if it is mentioned before the instruction
                    # same logic as previous loop
                    prev_insn = suspected_string_ref_ea
                    # if we found the size no need to search for it again
                    while result == BADADDR:
                        prev_insn = ida_ua.decode_prev_insn(insn, prev_insn)
                        if (
                            prev_insn < self.func.start_ea
                            or prev_insn == BADADDR
                            or insn.get_canon_feature() & ida_idp.CF_CALL != 0
                        ):
                            should_continue = True
                            break
                        else:
                            result = ida_xref.get_first_dref_from(prev_insn)

                    if should_continue:
                        continue

                    obj_ea = result

                elif skip_ref_y.op == ida_hexrays.cot_obj:
                    obj_ea = skip_ref_y.obj_ea

                if obj_ea:
                    self.string_assignments[string_key].pointer = obj_ea

            elif member.name == "len":
                if skip_cast_y.op == ida_hexrays.cot_num:
                    self.string_assignments[string_key].length = skip_cast_y.n._value

        return 0

    def finalize(self):
        for go_string in self.string_assignments.values():
            if go_string.pointer != 0 and go_string.length != 0:
                self.define_string(go_string.pointer, go_string.length)

    def visit_expr(self, e: ida_hexrays.cexpr_t) -> int:
        """
        Iterate all objects in decompiler and correct the size of strings
        """
        # in the case that inside the decompiler an object's type is a string
        # we then define the string param correctly
        if e.op == ida_hexrays.cot_call:
            for item in e.a:
                orig_item: ida_hexrays.carg_t = item
                if item.op == ida_hexrays.cot_cast:
                    item = item.x
                if item.op == ida_hexrays.cot_call:
                    if (
                        orig_item.formal_type.compare(self.string_tinfo) == 0
                        and item.a[0].op == ida_hexrays.cot_num
                        and item.a[1].op == ida_hexrays.cot_obj
                    ):
                        self.define_string(item.a[1].obj_ea, item.a[0].n._value)

        if e.op == ida_hexrays.cot_obj:
            # sometimes the actual ea of the instruction is not saved in the decompiler
            # try to find it in the parent of the item
            while e.ea == BADADDR:
                e = self.cfunc.find_parent_of(e)

            # we found the parent instruction
            suspected_string_ref_ea = e.ea
            string_size = 0

            # get the string referenced by the instruction
            insn = ida_ua.insn_t()
            # move the pointer so next instruction we parse is after the one which points to the string
            next_insn_ea = (
                ida_ua.decode_insn(insn, suspected_string_ref_ea)
                + suspected_string_ref_ea
            )

            # get the register that references the string
            reg_name = ida_idp.get_reg_name(
                insn.Op1.reg, ida_ua.get_dtype_size(ida_ua.dt_qword)
            )

            # object referenced not by a param passing register, we can't know its size
            if reg_name not in go_fast_convention:
                return 0

            reg_index = go_fast_convention.index(reg_name)
            string_ref = ida_xref.get_first_dref_from(suspected_string_ref_ea)
            if string_ref == BADADDR:
                # get the next reg in the fastcall, the next register where the string size will be stored
                size_reg = ida_idp.str2reg(reg_name)
                result = self.instruction_stores_valid_size(insn, size_reg)
                if not result or reg_index == 0:
                    return 0
                lookup_reg = ida_idp.str2reg(go_fast_convention[reg_index - 1])
                string_size = insn.Op2.value
                check_reg = self.instruction_stores_valid_obj
            else:
                if reg_index == len(go_fast_convention) - 1:
                    return 0
                lookup_reg = ida_idp.str2reg(go_fast_convention[reg_index + 1])
                check_reg = self.instruction_stores_valid_size

            found = False
            while not found:
                next_insn_ea = ida_ua.decode_insn(insn, next_insn_ea) + next_insn_ea
                # string size instruction must be in the function
                if next_insn_ea > self.func.end_ea:
                    break
                else:
                    result = check_reg(insn, lookup_reg)
                    if result is None:
                        break
                    else:
                        found = result

            # find the string size if it is mentioned before the instruction
            # same logic as previous loop
            prev_insn = suspected_string_ref_ea
            # if we found the size no need to search for it again
            while not found:
                prev_insn = ida_ua.decode_prev_insn(insn, prev_insn)
                if prev_insn < self.func.start_ea or prev_insn == BADADDR:
                    return 0
                else:
                    result = check_reg(insn, lookup_reg)
                    if result is None:
                        return 0
                    else:
                        found = result

            if string_ref == BADADDR:
                suspected_string_ref_ea = insn.ea
                string_ref = ida_xref.get_first_dref_from(suspected_string_ref_ea)
            else:
                # if the Op2 of the instruction is referencing an address it means it is not the constant length we are looking for
                if insn.Op2.addr != 0:
                    return 0
                string_size = insn.Op2.value

            # if we can't extract the string then we have a problem
            self.define_string(string_ref, string_size)

        return 0


class CallAnalysisHooks(ida_hexrays.Hexrays_Hooks):
    # 8 bit registers have different ids so we need a list of all 64 bit registers and all 8 bit registers
    valid_reg_ids = set(
        itertools.chain.from_iterable(
            [
                (
                    ida_idp.str2reg(item),
                    ida_idp.str2reg(get_sized_register_by_name(item, 1)),
                )
                for item in go_fast_convention
            ]
        )
    )

    def __init__(self, node: ida_netnode.netnode, detected_go: bool) -> None:
        super().__init__()
        self.marked_eas_set = set()
        self.marked_eas_list = list()

        self.tag_name = "x"
        self.node = node
        self.node_index = self.node.index()

        self.detected_go = detected_go

        # load previously fixed function prototypes
        blob = self.node.getblob(self.node_index, self.tag_name)
        if blob:
            decompiled_eas = json.loads(blob.decode("ascii"))
            self.marked_eas_set = set(decompiled_eas)
            self.marked_eas_list = decompiled_eas

        self.known_functions = known_functions

    def save_eas(self) -> None:
        self.node.setblob(
            json.dumps(self.marked_eas_list).encode("ascii"),
            self.node_index,
            self.tag_name,
        )

    def guess_func_return_type(self, ea: int) -> str:
        used_reg = -1
        func = ida_funcs.get_func(ea)
        flow_chart = ida_gdl.qflow_chart_t("", func, func.start_ea, func.end_ea, 0)
        while flow_chart is not None:
            should_continue = False
            insn = ida_ua.insn_t()
            funcs = set()
            for i in range(flow_chart.node_qty()):

                # iterate blocks that return from the function they might contain the result registers
                if flow_chart.is_ret_block(i):

                    # skip the stack extension block, it is considered as a return block for some reason
                    current_ea = flow_chart[i].end_ea
                    if ida_xref.get_first_cref_from(current_ea) == func.start_ea:
                        continue

                    # iterate from the return backwards until we find a call which spoils all registers
                    while (
                        current_ea > flow_chart[i].start_ea
                        and not insn.get_canon_feature() & ida_idp.CF_CALL
                    ):
                        current_ea = ida_ua.decode_prev_insn(insn, current_ea)

                        # if the register we currently change is a return register we check it
                        if (
                            insn.Op1.type == ida_ua.o_reg
                            and insn.Op1.reg in self.valid_reg_ids
                            # in order to assign a value to a register we need to have a second operand
                            and insn.Op2.dtype != ida_ua.o_void
                        ):
                            # each time we take the max register which is changed, we can assume the function
                            # wouldn't change it if it wasn't part of its result
                            used_reg = max(
                                used_reg,
                                go_fast_convention.index(
                                    ida_idp.get_reg_name(insn.Op1.reg, 8)
                                ),
                            )

                    # gather all call instructions which are called before returning from the function
                    if insn.get_canon_feature() & ida_idp.CF_CALL:
                        funcs.add(insn.ea)

            # the function might return the result of its last called function
            if used_reg == -1:

                # check our gathered calls for their prototypes
                for ea in funcs:
                    # get the callee ea
                    xrefed_func = ida_xref.get_first_cref_from(ea)
                    while func.start_ea <= xrefed_func <= func.end_ea:
                        xrefed_func = ida_xref.get_next_cref_from(ea, xrefed_func)

                    # ignore calls to registers, we can't follow those
                    if xrefed_func == BADADDR:
                        continue

                    func = ida_funcs.get_func(xrefed_func)
                    # the called function is already defined, take its return type
                    if ida_nalt.get_aflags(xrefed_func) & ida_nalt.AFL_USERTI:
                        tinfo = ida_typeinf.tinfo_t()
                        ida_hexrays.get_type(
                            xrefed_func, tinfo, ida_hexrays.GUESSED_FUNC
                        )
                        func_details = ida_typeinf.func_type_data_t()
                        tinfo.get_func_details(func_details)
                        return func_details.rettype.dstr()
                    # recurse the inner function and try to find its return registers
                    else:
                        flow_chart = ida_gdl.qflow_chart_t(
                            "", func, func.start_ea, func.end_ea, 0
                        )
                        should_continue = True
                        break

            if should_continue:
                continue
            flow_chart = None

        # we couldn't find any registers used in the return blocks
        if used_reg == -1:
            return "void"
        elif used_reg == 0:
            return "void *"
        else:
            # sometimes functions' parents create their return type for them
            struc_name = f"retval_{func.start_ea:X}"
            if ida_struct.get_struc_id(struc_name) == BADADDR:
                # create the result struct if it doesn't exist
                retval_struc = ida_struct.add_struc(BADADDR, struc_name)
                for i in range(used_reg + 1):
                    ida_struct.add_struc_member(
                        ida_struct.get_struc(retval_struc),
                        f"part_{i}",
                        BADADDR,
                        ida_bytes.qword_flag(),
                        None,
                        8,
                    )
            return struc_name

    def fix_call_by_ea(self, mba: ida_hexrays.mba_t, callee_ea: int) -> None:
        # check if we already fixed the call
        if callee_ea in self.marked_eas_set:
            return

        # mark call as fixed one
        self.marked_eas_set.add(callee_ea)
        self.marked_eas_list.append(callee_ea)

        # save the current information
        self.save_eas()

        if ida_name.get_name(callee_ea) in runtime_morestack_functions:
            func_declaration = "void func();"

        # if the name of the function is runtime.concatstring{number} we know its prototype
        elif concat_string_number.match(ida_funcs.get_func_name(callee_ea)):

            string_count = int(ida_funcs.get_func_name(callee_ea)[-1:])

            string_concat_declaration = f"string (*)(void * tmpbuf, {', '.join([f'string a{d}' for d in range(string_count)])});"

            string_concat_tinfo = ida_typeinf.tinfo_t()
            ida_typeinf.parse_decl(
                string_concat_tinfo, None, string_concat_declaration, 0
            )

            # calculate calling convention by the string concatenation count
            func_declaration = GoCall(
                mba, callee_ea, string_concat_tinfo, False
            ).get_decl_string()

        # check if the prototype of the function is already known by ida
        elif ida_nalt.get_aflags(callee_ea) & ida_nalt.AFL_USERTI:
            tinfo = ida_typeinf.tinfo_t()
            ida_hexrays.get_type(callee_ea, tinfo, ida_hexrays.GUESSED_FUNC)
            func_declaration = GoCall(
                mba, callee_ea, tinfo, self.detected_go
            ).get_decl_string()

        # we can only guess the calling convention manually by the assembly
        else:
            func = ida_funcs.get_func(callee_ea)
            end_ea = func.end_ea
            current_ea = func.start_ea

            # find the first jump
            next_ea = ida_bytes.next_head(current_ea, end_ea)
            while ida_xref.get_next_cref_from(current_ea, next_ea) == BADADDR:
                current_ea = next_ea
                next_ea = ida_bytes.next_head(current_ea, end_ea)
                if current_ea == BADADDR:
                    return

            # start iterating the stack extension block
            insn = ida_ua.insn_t()
            current_insn_ea = current_ea + ida_ua.decode_insn(insn, current_ea)
            current_insn_ea = ida_xref.get_next_cref_from(current_ea, current_insn_ea)

            # if the first instruction in the block is not one of storing registers than it means we are not in a go fastcall
            current_insn_ea = current_insn_ea + ida_ua.decode_insn(
                insn, current_insn_ea
            )
            if (
                insn.Op1.type != ida_ua.o_displ
                or insn.Op1.addr == 0
                or insn.Op2.reg not in self.valid_reg_ids
            ):
                xrefed_func = ida_xref.get_first_cref_from(insn.ea)
                while func.start_ea <= xrefed_func <= func.end_ea:
                    xrefed_func = ida_xref.get_next_cref_from(insn.ea, xrefed_func)
                if ida_name.get_name(xrefed_func) not in runtime_morestack_functions:
                    return

            func_declaration = ""

            param_regs_list = []
            # iterate until the stack extension call
            while insn.get_canon_feature() & ida_idp.CF_CALL == 0:
                # extract stored register into func declaration
                if insn.Op1.type == ida_ua.o_displ and insn.Op1.addr != 0:
                    op2_dtype_size = ida_ua.get_dtype_size(insn.Op2.dtype)
                    param_regs_list.append(f"__int{op2_dtype_size * BYTE_SIZE}")
                    if not self.detected_go:
                        param_regs_list[
                            -1
                        ] += f"@<{ida_idp.get_reg_name(insn.Op2.reg, op2_dtype_size)}>"
                current_insn_ea = current_insn_ea + ida_ua.decode_insn(
                    insn, current_insn_ea
                )

            # finalize func declaration
            func_declaration = self.guess_func_return_type(callee_ea)
            if not self.detected_go:
                if func_declaration != "void":
                    call = GoCall(None, 0, detected_go=self.detected_go)
                    call_ret = ida_typeinf.tinfo_t()
                    ida_typeinf.parse_decl(call_ret, None, f"{func_declaration};", 0)
                    call.add_ret(call_ret)
                    func_declaration = (
                        f"{call.ret_type} __usercall func@<{call.ret_loc}>"
                    )
                else:
                    func_declaration = f"void __usercall func"
            else:
                func_declaration = f"{func_declaration} __golang func"
            func_declaration += f"({','.join(param_regs_list)});"

        # set function typeinfo
        call_info = ida_typeinf.tinfo_t()
        ida_typeinf.parse_decl(call_info, None, func_declaration, 0)
        ida_typeinf.apply_tinfo(callee_ea, call_info, ida_typeinf.TINFO_DEFINITE)

    # fix the current call information
    def maturity(self, cfunc: ida_hexrays.cfunc_t, new_maturity: int) -> int:
        if new_maturity == ida_hexrays.CMAT_FINAL:
            self.fix_call_by_ea(cfunc.mba, cfunc.entry_ea)
            # fix strings in decompiler
            func = ida_funcs.get_func(cfunc.entry_ea)
            tinfo = ida_typeinf.tinfo_t()
            ida_typeinf.parse_decl(tinfo, None, "string;", ida_typeinf.PT_SIL)
            function_string_visitor = FunctionStringVisitor(cfunc.body, func, tinfo)
            function_string_visitor.apply_to(cfunc.body, None)
            function_string_visitor.finalize()
        return 0

    def build_callinfo(
        self, blk: ida_hexrays.mblock_t, call_type: ida_typeinf.tinfo_t
    ) -> ida_hexrays.mcallinfo_t | None:
        """
        Fix the calls our decompiled function calls to, resulting in a satisfying decompilation
        Also fix casts of known functions
        """
        call_insn = blk.tail

        # in the hexrays microcode the callee is at the left operand
        # we check here that the type of call is global, which means it is a call to a specific function
        # and not to an interface
        result = None
        if call_insn.l.t == ida_hexrays.mop_v:
            callee_ea = call_insn.l.g
            func_name = ida_funcs.get_func_name(callee_ea)
            if func_name in self.known_functions.keys():
                result = self.known_functions[func_name].get_decl_mi(
                    blk.mba, call_insn.ea, callee_ea
                )
            else:
                self.fix_call_by_ea(blk.mba, callee_ea)
        return result


class R14Optimizer(ida_hexrays.optinsn_t):

    def __init__(self, runtime_g: ida_typeinf.tinfo_t) -> None:
        super().__init__()
        self.r14_code = ida_hexrays.reg2mreg(ida_idp.str2reg("r14"))
        self.runtime_g = runtime_g

    def func(
        self, blk: ida_hexrays.mblock_t, ins: ida_hexrays.minsn_t, optflags: int
    ) -> int:
        """
        Find uses of r14 in the microcode, replace them with calls to an helper named CurrentGoroutine
        This helper is meant to mimic how ida converts fs/gs access to the NtCurrentTeb helper
        """
        current_insn = blk.head

        while current_insn != blk.tail and current_insn is not None:
            left_op = current_insn.l
            if (
                current_insn.opcode == ida_hexrays.m_mov
                and left_op.t == ida_hexrays.mop_r
                and left_op.r == self.r14_code
            ):
                # take the operand and the type from the original micro instruction
                call_insn = ida_hexrays.minsn_t(current_insn.ea)
                call_insn.opcode = ida_hexrays.m_call

                # initialize the helper part
                helper_mop = ida_hexrays.mop_t()
                helper_mop.make_helper("CurrentGoroutine")
                call_insn.l = helper_mop

                # initialize the callinfo part
                callinfo_mop = ida_hexrays.mop_t()
                callinfo = ida_hexrays.mcallinfo_t()

                # set function return regs
                regs_vec = ida_hexrays.mopvec_t()
                return_mop = ida_hexrays.mop_t()
                return_mop.make_reg(current_insn.d.r, 8)
                regs_vec.push_back(return_mop)

                # set the function return locations
                regs_list = ida_hexrays.mlist_t()
                regs_list.add(current_insn.d.r, 8)

                # set the location of the returned value
                ret_loc = ida_typeinf.argloc_t()
                ret_loc.set_reg1(current_insn.d.r)

                # set the registers spoiled by this function
                spoiled_list = ida_hexrays.mlist_t()
                spoiled_list.add(current_insn.d.r, 8)

                # set all of the callinfo information
                # set the return type to be the runtime_g (the goroutine struct)
                callinfo.return_type = self.runtime_g.copy()
                callinfo.retregs = regs_vec
                callinfo.return_regs = regs_list
                callinfo.spoiled = spoiled_list
                callinfo.return_argloc = ret_loc
                callinfo.flags = ida_hexrays.FCI_NOSIDE
                callinfo.cc = go_calling_convention

                # set the actual micro instruction
                callinfo_mop.t = ida_hexrays.mop_f
                callinfo_mop.f = callinfo
                callinfo_mop.size = 8
                call_insn.d = callinfo_mop

                # put the call to the helper as an operand
                blk.insert_into_block(call_insn, current_insn.prev)
                blk.make_nop(current_insn)
            current_insn = current_insn.next

        return 0


class GoAnalyzer(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_FIX
    wanted_name = "GoAnalyzer"
    GOANALYZER_NODE = "GoAnalyzerNode"

    def init(self) -> int:
        self.enabled = False
        self.initialized = False

        if not ida_hexrays.init_hexrays_plugin():
            return ida_idaapi.PLUGIN_SKIP

        # only 64 bit go has the abi we fix
        if not ida_ida.inf_is_64bit() or ida_idp.ph_get_id() != ida_idp.PLFM_386:
            return ida_idaapi.PLUGIN_SKIP

        # go version is specified in the data/go.buildinfo sections, check it
        data_segment: ida_segment.segment_t
        if ida_ida.inf_get_filetype() == ida_ida.f_PE:
            data_segment = ida_segment.get_segm_by_name(".data")
        elif ida_ida.inf_get_filetype() == ida_ida.f_ELF:
            data_segment = ida_segment.get_segm_by_name(".go.buildinfo")
        else:
            return ida_idaapi.PLUGIN_SKIP

        # if there is no data segment means it can't be in go
        if data_segment is None:
            return ida_idaapi.PLUGIN_SKIP

        # in newer versions of go they had some data put before the magic string we search for
        seg_start = data_segment.start_ea
        binpat_vec = ida_bytes.compiled_binpat_vec_t()
        go_buildinf_pattern = " ".join(hex(char)[2:] for char in b"\xff Go buildinf:")

        ida_bytes.parse_binpat_str(binpat_vec, seg_start, go_buildinf_pattern, 0x10)
        version_string_start = ida_bytes.bin_search(
            seg_start, seg_start + 0x1000, binpat_vec, ida_bytes.BIN_SEARCH_FORWARD
        )

        if version_string_start == BADADDR:
            return ida_idaapi.PLUGIN_SKIP

        # check the data segment if the version string length is reasonable
        VERSION_STRING_OFFSET = 0x20
        version_string_len = ida_bytes.get_db_byte(
            version_string_start + VERSION_STRING_OFFSET
        )
        if (
            version_string_len is None
            or version_string_len > 15
            or version_string_len < 5
        ):
            return ida_idaapi.PLUGIN_SKIP

        # we probably have the go version string
        supposed_version_string = ida_bytes.get_bytes(
            version_string_start + VERSION_STRING_OFFSET + 1, version_string_len
        ).decode("ascii")
        # if it is unknown allow manual activation
        if "unknown" in supposed_version_string.lower():
            print("Didn't detect go version, try manually in the plugins menu")
            return ida_idaapi.PLUGIN_KEEP

        match = go_version_regex.match(supposed_version_string)
        MINIMAL_REGISTER_ABI_GO_VERSION = 17
        # check if the go version is higher than 1.17 which means that the register abi is on
        if not match or int(match.group(1)) < MINIMAL_REGISTER_ABI_GO_VERSION:
            return ida_idaapi.PLUGIN_SKIP

        print(f"Detected {supposed_version_string}")
        self.run(None)
        return ida_idaapi.PLUGIN_KEEP

    def run(self, _) -> int:
        """Install all the hooks we need and do our initialization"""
        if not self.initialized:
            self.detected_go = False
            if (
                GO_SUPPORTED
                and ida_idaapi.get_inf_structure().cc.cm & ida_typeinf.CM_CC_MASK
                == ida_typeinf.CM_CC_GOLANG
            ):
                self.detected_go = True

            # initialize our netnode, we don't want to recompile everything all the time
            self.node = ida_netnode.netnode()
            self.node.create(GoAnalyzer.GOANALYZER_NODE)

            self.hooks = CallAnalysisHooks(self.node, self.detected_go)
            self.filter = Xmm15Optimizer()
            self.runtime_g = ida_typeinf.tinfo_t()
            # create the goroutine struct if it doesn't already exist
            if (
                ida_typeinf.get_named_type(None, "runtime_g", ida_typeinf.NTF_TYPE)
                is None
            ):
                runtime_g_struct = ida_struct.add_struc(BADADDR, "runtime_g")
                # newer versions of ida seem to need the structure to have members in order to parse it correctly
                ida_struct.add_struc_member(
                    ida_struct.get_struc(runtime_g_struct),
                    "data",
                    BADADDR,
                    0,
                    None,
                    8,
                )
                ida_typeinf.parse_decl(
                    self.runtime_g, None, "runtime_g;", ida_typeinf.PT_SIL
                )
            else:
                tinfo = ida_typeinf.tinfo_t()
                ida_typeinf.parse_decl(tinfo, None, "runtime_g;", ida_typeinf.PT_SIL)
                remove_tinfo_zeroes(tinfo)
                ida_typeinf.parse_decl(
                    self.runtime_g, None, "runtime_g_nozeroes *;", ida_typeinf.PT_SIL
                )

            # create the string struct if it doesn't already exist
            if ida_typeinf.get_named_type(None, "string", ida_typeinf.NTF_TYPE) is None:
                string_struc = ida_struct.add_struc(BADADDR, "string")
                ida_struct.add_struc_member(
                    ida_struct.get_struc(string_struc),
                    "ptr",
                    BADADDR,
                    ida_bytes.qword_flag(),
                    None,
                    8,
                )
                ida_struct.add_struc_member(
                    ida_struct.get_struc(string_struc),
                    "len",
                    BADADDR,
                    ida_bytes.qword_flag(),
                    None,
                    8,
                )

            self.optimizer = R14Optimizer(self.runtime_g)
            self.initialized = True
            print("GoAnalyzer Initialized")

        if not self.enabled:
            self.hooks.hook()

            # install the xmm15 zero operand optimizer
            ida_hexrays.install_microcode_filter(self.filter, True)

            # install the r14 current goroutine optimizer
            self.optimizer.install()
            self.enabled = True
            print("GoAnalyzer Enabled")
        else:
            self.hooks.unhook()

            # uninstall the xmm15 zero operand optimizer
            ida_hexrays.install_microcode_filter(self.filter, False)

            # uninstall the r14 current goroutine optimizer
            self.optimizer.remove()

            self.enabled = False
            print("GoAnalyzer Disabled")
        return 0

    def term(self) -> None:
        pass


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


repository: hoangcuongflp/Stuff
file: ida-master/plugins/leafblower/leafblower.py
description: because I've orderliness obsession :(
wanted_name: leafblower
comment: Assists in identifying common POSIX functions in RISC architectures
created_at: 2015-10-12 16:35:32+00:00
pushed_at: 2014-12-22 00:32:06+00:00
url:
https://github.com/hoangcuongflp/Stuff/blob/ee100760a05f524ff73837222d80f1dd12da3cee/ida-master/plugins/leafblower/leafblower.py

content:
#################################################################################################
#
# A plugin to help identify common POSIX functions such as printf, sprintf, memcmp, strcpy, etc.
#
# This plugin will really only work with RISC architectures, as it assumes a fixed instruction
# size, and that function arguments are passed via registers.
#
#################################################################################################
import idc
import idaapi
import idautils

class LeafBlowerFunctionChooser(idaapi.Choose2):

    MIN_XREFS = 25
    MUST_HAVE_LOOP = True

    def __init__(self, lbobj):
        self.lb = lbobj
        self.min_xrefs = self.MIN_XREFS
        self.must_have_loop = self.MUST_HAVE_LOOP

        idaapi.Choose2.__init__(self, self.lb.TITLE, self.lb.COLUMNS)
        self.icon = 41

        self.populate_items()

    def OnSelectLine(self, n):
        idc.Jump(idc.LocByName(self.items[n][0]))

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

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

    def OnRefresh(self, n):
        self.populate_items()
        return len(self.items)

    def OnClose(self):
        pass

    def OnCommand(self, n, cmd):
        if cmd == self.show_all_toggle_cmd:

            if self.min_xrefs == self.MIN_XREFS:
                self.min_xrefs = 0
            if self.min_xrefs != self.MIN_XREFS:
                self.min_xrefs = self.MIN_XREFS

            if self.must_have_loop == self.MUST_HAVE_LOOP:
                self.must_have_loop = False
            else:
                self.must_have_loop = self.MUST_HAVE_LOOP

        elif cmd == self.rename_cmd:

            if idc.AskYN(0, "Are you sure you want to rename all 'sub_XXXXXX' functions to 'leaf_XXXXXX'?") == 1:
                for item in self.items:
                    # Is this a leaf function?
                    if item[-1] == True:
                        current_name = item[0]
                        if current_name.startswith('sub_'):
                            new_name = current_name.replace('sub_', 'leaf_')
                            idc.MakeName(idc.LocByName(current_name), new_name)

        self.populate_items()
        return 0

    def populate_items(self):
        self.items = []

        for function in self.lb.functions:
            candidates = []

            if function.leaf:
                if function.xrefs < self.min_xrefs:
                    continue
                if function.loop == False and self.must_have_loop == True:
                    continue

            for candidate in function.candidates:
                candidates.append(candidate)

            if function.xrefs:
                xrefs = str(function.xrefs)
            else:
                xrefs = "*"

            if function.argc is not None:
                argc = str(function.argc)
            else:
                argc = "*"

            if function.leaf:
                loops = str(function.loop)
            else:
                loops = "*"

            name = idc.Name(function.start)

            self.items.append([name, xrefs, argc, loops, ', '.join(candidates), function.leaf])

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

        self.show_all_toggle_cmd = self.AddCommand("Toggle 'show all leaf functions'")
        self.rename_cmd = self.AddCommand("Rename all sub_XXXXXX leaf functions to 'leaf_XXXXXX'")

class ArchitectureSettings(object):

    def __init__(self, **kwargs):
        for (k,v) in kwargs.iteritems():
            setattr(self, k, v)

class ArchitectureSpecific(object):
    '''
    Architecture specific configuration / code.
    '''

    # Architecture specific definitions go here.
    #
    #@name  - An arbitrary name describing the architecture. Currently unused.
    #@argv  - A list of register names used to pass function arguments.
    #@delay - The number of jump/branch delay slots used by the architecture (typcally 0).
    #@size  - The architecture's insturction size (assumes a fixed-width instruction set).
    ARCHES = [
                ArchitectureSettings(name="MIPS", argv=['$a0', '$a1', '$a2', '$a3'], delay=1, size=4),
                ArchitectureSettings(name="ARM", argv=['R0', 'R1', 'R2', 'R3'], delay=0, size=4),
             ]

    def __init__(self):
        # Get a list of the current processor module's registers
        self.registers = idaapi.ph_get_regnames()

        # Find the list of function argument registers that fit the current processor module
        # TODO: Probably better to look this up based on the processor module name, as it can't
        #       distinguish between ARM and Thumb.
        self.argv = None
        for arch in self.ARCHES:
            if not (set(arch.argv) - set(self.registers)):
                self.argv = list(arch.argv)
                self.delay_slot = arch.delay
                self.insn_size = arch.size
                break

        if self.argv is None:
            raise Exception("Unknown/unsupported architecture!")

class Prototype(object):

    def __init__(self, name, argc, argv=[], loop=True, fmtarg=None):
        self.name = name
        self.argc = argc
        self.loop = loop
        self.argv = argv
        self.fmtarg = fmtarg

class Function(object):
    '''
    A wrapper class for storing function related info.
    '''

    PROTOTYPES = [
                    Prototype(name="atoi", argc=1),
                    Prototype(name="atol", argc=1),
                    Prototype(name="strlen", argc=1),
                    Prototype(name="strcpy", argc=2),
                    Prototype(name="strcat", argc=2),
                    Prototype(name="strcmp", argc=2),
                    Prototype(name="strstr", argc=2),
                    Prototype(name="strchr", argc=2),
                    Prototype(name="strrchr", argc=2),
                    Prototype(name="bzero", argc=2),
                    Prototype(name="strtol", argc=3),
                    Prototype(name="strncpy", argc=3),
                    Prototype(name="strncat", argc=3),
                    Prototype(name="strncmp", argc=3),
                    Prototype(name="memcpy", argc=3),
                    Prototype(name="memmove", argc=3),
                    Prototype(name="bcopy", argc=3),
                    Prototype(name="memcmp", argc=3),
                    Prototype(name="memset", argc=3),

                    Prototype(name="printf", argc=1, fmtarg=0),
                    Prototype(name="sprintf", argc=2, fmtarg=1),
                    Prototype(name="snprintf", argc=3, fmtarg=2),
                    Prototype(name="fprintf", argc=2, fmtarg=1),
                    Prototype(name="fscanf", argc=2, fmtarg=1),
                    Prototype(name="sscanf", argc=2, fmtarg=1),
                 ]

    def __init__(self, **kwargs):
        self.argc = None
        self.loop = False
        self.leaf = False
        self.xrefs = None
        self.fmtarg = None
        self.start = idc.BADADDR
        self.end = idc.BADADDR
        self.candidates = {}
        self.argp = ArgParser()

        for (k,v) in kwargs.iteritems():
            setattr(self, k, v)

        self.name = idc.Name(self.start)

        if self.xrefs is None:
            self.xrefs = len([x for x in idautils.XrefsTo(self.start)])

        if not self.candidates:
            for prototype in self.PROTOTYPES:
                if self.leaf and prototype.fmtarg is None and prototype.argc == self.argc and prototype.loop == self.loop:
                    if self.candidates.has_key(prototype.name):
                        self.candidates[prototype.name] += 1
                    else:
                        self.candidates[prototype.name] = 1
                elif not self.leaf and self.fmtarg is not None and prototype.fmtarg is not None and self.fmtarg == prototype.fmtarg:
                    if self.candidates.has_key(prototype.name):
                        self.candidates[prototype.name] += 1
                    else:
                        self.candidates[prototype.name] = 1

class ArgParser(object):
    '''
    Attempts to identify the number of arguments a function takes as well as
    what type of arguments a function takes.
    '''

    # An iterable list of canonical flags indicating an operand has been changed
    CHANGE_OPND = [
                    idaapi.CF_CHG1,
                    idaapi.CF_CHG2,
                    idaapi.CF_CHG3,
                    idaapi.CF_CHG4,
                    idaapi.CF_CHG5,
                    idaapi.CF_CHG6,
                  ]

    # An iterable list of canonical flags indicating an operand has been used
    USE_OPND = [
                    idaapi.CF_USE1,
                    idaapi.CF_USE2,
                    idaapi.CF_USE3,
                    idaapi.CF_USE4,
                    idaapi.CF_USE5,
                    idaapi.CF_USE6,
               ]

    def __init__(self):
        self.arch = ArchitectureSpecific()

    def argc(self, function):
        '''
        Counts the number of arguments used by the specified function.
        '''
        argv = set()
        notargv = set()
        ea = function.startEA

        while ea < function.endEA:
            idaapi.decode_insn(ea)
            features = idaapi.cmd.get_canon_feature()

            for n in range(0, len(self.USE_OPND)):
                if idaapi.cmd.Operands[n].type in [idaapi.o_reg, idaapi.o_displ, idaapi.o_phrase]:
                    try:
                        regname = self.arch.registers[idaapi.cmd.Operands[n].reg]
                        index = self.arch.argv.index(regname)
                    except ValueError:
                        continue

                    if features & self.USE_OPND[n] and regname not in notargv:
                        argv.update(self.arch.argv[:index+1])

            for n in range(0, len(self.CHANGE_OPND)):
                if idaapi.cmd.Operands[n].type in [idaapi.o_reg, idaapi.o_displ, idaapi.o_phrase]:
                    try:
                        regname = self.arch.registers[idaapi.cmd.Operands[n].reg]
                        index = self.arch.argv.index(regname)
                    except ValueError:
                        continue

                    if features & self.CHANGE_OPND[n] and regname not in argv:
                        notargv.update(self.arch.argv[index:])

            if argv.union(notargv) == set(self.arch.argv):
                break

            # TODO: Use idc.NextHead(ea) instead...
            ea += self.arch.insn_size

        return len(argv)

    def trace(self, ea):
        '''
        Given an EA where an argument register is set, attempt to trace what
        function call that argument is passed to.

        @ea - The address of an instruction that modifies a function argument register.

        Returns a tuple of (function EA, argv index, argument register name) on success.
        Returns None on failure.
        '''
        idaapi.decode_insn(ea)
        features = idaapi.cmd.get_canon_feature()

        for n in range(0, len(self.CHANGE_OPND)):
            if idaapi.cmd.Operands[n].type in [idaapi.o_reg, idaapi.o_displ, idaapi.o_phrase]:
                try:
                    regname = self.arch.registers[idaapi.cmd.Operands[n].reg]
                    index = self.arch.argv.index(regname)
                except ValueError:
                    continue

                if features & self.CHANGE_OPND[n]:
                    ea = ea - (self.arch.delay_slot * self.arch.insn_size)

                    while True:
                        idaapi.decode_insn(ea)

                        if idaapi.is_call_insn(ea):
                            for xref in idautils.XrefsFrom(ea):
                                if xref.type in [idaapi.fl_CF, idaapi.fl_CN]:
                                    return (xref.to, index, regname)
                            # If we couldn't figure out where the function call was going to, just quit
                            break

                        if idaapi.is_basic_block_end(ea):
                            break

                        ea += self.arch.insn_size

        return (None, None, None)


    def argv(self, func):
        '''
        Attempts to identify what types of arguments are passed to a given function.
        Currently unused.
        '''
        args = [None for x in self.arch.argv]

        for xref in idautils.XrefsTo(func.startEA):
            if idaapi.is_call_insn(xref.frm):
                idaapi.decode_insn(xref.frm)

                ea = xref.frm + (self.arch.delay_slot * self.arch.insn_size)
                end_ea = (xref.frm - (self.arch.insn_size * 10))

                while ea >= end_ea:
                    # Stop searching if we've reached a conditional block or another call
                    if idaapi.is_basic_block_end(ea) or (ea != xref.frm and idaapi.is_call_insn(ea)):
                        break

                    idaapi.decode_insn(ea)
                    features = idaapi.cmd.get_canon_feature()

                    for n in range(0, len(self.CHANGE_OPND)):
                        if idaapi.cmd.Operands[n].type in [idaapi.o_reg, idaapi.o_displ, idaapi.o_phrase]:
                            try:
                                regname = self.arch.registers[idaapi.cmd.Operands[n].reg]
                                index = self.arch.argv.index(regname)
                            except ValueError:
                                continue

                            if features & self.CHANGE_OPND[n]:
                                for xref in idautils.XrefsFrom(ea):
                                    # TODO: Where is this xref type defined?
                                    if xref.type == 1:
                                        string = idc.GetString(xref.to)
                                        if string and len(string) > 4:
                                            args[index] = str
                                        break

                    ea -= self.arch.insn_size

                yield args

class LeafFunctionFinder(object):
    '''
    Class that searches for functions that do not call any other functions.
    '''

    TITLE = "Leaf functions"

    COLUMNS = [
                ["Function", 25 | idaapi.Choose2.CHCOL_PLAIN],
                ["Xrefs", 8 | idaapi.Choose2.CHCOL_PLAIN],
                ["argc", 8 | idaapi.Choose2.CHCOL_PLAIN],
                ["Has Loop(s)", 8 | idaapi.Choose2.CHCOL_PLAIN],
                ["Possible candidate(s)", 50 | idaapi.Choose2.CHCOL_PLAIN],
              ]

    def __init__(self):
        self.functions = []
        self.arch = ArchitectureSpecific()
        self.argp = ArgParser()
        self._find_leafs()

    def _find_leafs(self):
        # Loop through every function
        for func_ea in idautils.Functions():
            # Count the number of xrefs to this function
            func = idaapi.get_func(func_ea)
            if func:
                leaf_function = True
                ea = func.startEA

                # Loop through all instructions in this function looking
                # for call instructions; if found, then this is not a leaf.
                while ea <= func.endEA:
                    idaapi.decode_insn(ea)
                    if idaapi.is_call_insn(ea):
                        leaf_function = False
                        break

                    ea += self.arch.insn_size

                if leaf_function:
                    self.functions.append(Function(start=func.startEA,
                                                   end=func.endEA,
                                                   leaf=True,
                                                   loop=self.has_loop(func),
                                                   argc=self.argp.argc(func)))

        # Sort leafs by xref count, largest first
        self.functions.sort(key=lambda f: f.xrefs, reverse=True)

    def has_loop(self, func):
        '''
        A nieve method for checking to see if a function contains a loop.
        Works pretty well for simple functions though.
        '''
        blocks = [func.startEA] + [block.endEA for block in idaapi.FlowChart(func)]
        for block in blocks:
            for xref in idautils.XrefsTo(block):
                xref_func = idaapi.get_func(xref.frm)
                if xref_func and xref_func.startEA == func.startEA:
                    if xref.frm >= block:
                        return True
        return False

class FormatStringFunctionFinder(object):

    TITLE = "Format string functions"

    COLUMNS = [
                ["Function", 25 | idaapi.Choose2.CHCOL_PLAIN],
                ["Xrefs", 8 | idaapi.Choose2.CHCOL_PLAIN],
                ["Format string argv index", 15 | idaapi.Choose2.CHCOL_PLAIN],
                ["Has Loop(s)", 8 | idaapi.Choose2.CHCOL_PLAIN],
                ["Possible candidate(s)", 50 | idaapi.Choose2.CHCOL_PLAIN],
              ]
    def __init__(self):
        self.functions = []
        self.argp = ArgParser()
        self._find_format_string_functions()

    def _find_format_string_functions(self):
        processed_func_eas = set()

        for string in idautils.Strings():
            if '%' in str(string):
                for xref in idautils.XrefsTo(string.ea):
                    (func_ea, argn, regname) = self.argp.trace(xref.frm)
                    if func_ea is not None and func_ea not in processed_func_eas:
                        # Abuse argc here to show the argv index of the format string
                        self.functions.append(Function(start=func_ea, argc=argn, fmtarg=argn))
                        processed_func_eas.add(func_ea)

        # Sort format string functions by xref count, largest first
        self.functions.sort(key=lambda f: f.xrefs, reverse=True)

class leaf_blower_t(idaapi.plugin_t):

    flags = 0
    comment = "Assists in identifying common POSIX functions in RISC architectures"
    help = ''
    wanted_name = 'leafblower'
    wanted_hotkey = ''

    def init(self):
        self.leaf_context_menu = idaapi.add_menu_item("Search/", "leaf functions", "", 0, self.LeafFromMenu, (None,))
        self.stdio_context_menu = idaapi.add_menu_item("Search/", "format string functions", "", 0, self.FormatFromMenu, (None,))
        return idaapi.PLUGIN_KEEP

    def term(self):
        idaapi.del_menu_item(self.leaf_context_menu)
        idaapi.del_menu_item(self.stdio_context_menu)
        return None

    def run(self):
        pass

    def LeafFromMenu(self, arg):
        LeafBlowerFunctionChooser(LeafFunctionFinder()).show()

    def FormatFromMenu(self, arg):
        LeafBlowerFunctionChooser(FormatStringFunctionFinder()).show()

def PLUGIN_ENTRY():
    return leaf_blower_t()


repository: patois/vds5plugin
file: vds5.py
description: vds5plugin - vds5.py script taken from IDAPython and turned into a plugin
url: https://github.com/patois/vds5plugin/blob/a6f77c0ca8d090c87bbe6f978e54b9acf6359cc0/vds5.py
wanted_name: hrgraph
created_at: 2019-05-20 17:21:36+00:00
pushed_at: 2020-06-04 15:52:48+00:00
content:
import idaapi
import ida_pro
import ida_hexrays
import ida_kernwin
import ida_gdl
import ida_lines

ACTION_NAME = "vds5.py:displaygraph"
ACTION_SHORTCUT = "Ctrl+Shift+G"

CL_WHITE            = ((255)+  (255<<8)+  (255<<16)) #   0
CL_BLUE             = ((0  )+  (0  <<8)+  (255<<16)) #   1
CL_RED              = ((255)+  (0  <<8)+  (0  <<16)) #   2
CL_GREEN            = ((0  )+  (255<<8)+  (0  <<16)) #   3
CL_YELLOW           = ((255)+  (255<<8)+  (0  <<16)) #   4
CL_MAGENTA          = ((255)+  (0  <<8)+  (255<<16)) #   5
CL_CYAN             = ((0  )+  (255<<8)+  (255<<16)) #   6
CL_DARKGREY         = ((85 )+  (85 <<8)+  (85 <<16)) #   7
CL_DARKBLUE         = ((0  )+  (0  <<8)+  (128<<16)) #   8
CL_DARKRED          = ((128)+  (0  <<8)+  (0  <<16)) #   9
CL_DARKGREEN        = ((0  )+  (128<<8)+  (0  <<16)) #  10
CL_DARKYELLOW       = ((128)+  (128<<8)+  (0  <<16)) #  11
CL_DARKMAGENTA      = ((128)+  (0  <<8)+  (128<<16)) #  12
CL_DARKCYAN         = ((0  )+  (128<<8)+  (128<<16)) #  13
CL_GOLD             = ((255)+  (215<<8)+  (0  <<16)) #  14
CL_LIGHTGREY        = ((170)+  (170<<8)+  (170<<16)) #  15
CL_LIGHTBLUE        = ((128)+  (128<<8)+  (255<<16)) #  16
CL_LIGHTRED         = ((255)+  (128<<8)+  (128<<16)) #  17
CL_LIGHTGREEN       = ((128)+  (255<<8)+  (128<<16)) #  18
CL_LIGHTYELLOW      = ((255)+  (255<<8)+  (128<<16)) #  19
CL_LIGHTMAGENTA     = ((255)+  (128<<8)+  (255<<16)) #  20
CL_LIGHTCYAN        = ((128)+  (255<<8)+  (255<<16)) #  21
CL_LILAC            = ((238)+  (130<<8)+  (238<<16)) #  22
CL_TURQUOISE        = ((64 )+  (224<<8)+  (208<<16)) #  23
CL_AQUAMARINE       = ((127)+  (255<<8)+  (212<<16)) #  24
CL_KHAKI            = ((240)+  (230<<8)+  (140<<16)) #  25
CL_PURPLE           = ((160)+  (32 <<8)+  (240<<16)) #  26
CL_YELLOWGREEN      = ((154)+  (205<<8)+  (50 <<16)) #  27
CL_PINK             = ((255)+  (192<<8)+  (203<<16)) #  28
CL_ORANGE           = ((255)+  (165<<8)+  (0  <<16)) #  29
CL_ORCHID           = ((218)+  (112<<8)+  (214<<16)) #  30
CL_BLACK            = ((0  )+  (0  <<8)+  (0  <<16)) #  31

COLORS_LUT = {
    CL_WHITE             : "white",
    CL_BLUE              : "blue",
    CL_RED               : "red",
    CL_GREEN             : "green",
    CL_YELLOW            : "yellow",
    CL_MAGENTA           : "magenta",
    CL_CYAN              : "cyan",
    CL_DARKGREY          : "darkgrey",
    CL_DARKBLUE          : "darkblue",
    CL_DARKRED           : "darkred",
    CL_DARKGREEN         : "darkgreen",
    CL_DARKYELLOW        : "darkyellow",
    CL_DARKMAGENTA       : "darkmagenta",
    CL_DARKCYAN          : "darkcyan",
    CL_GOLD              : "gold",
    CL_LIGHTGREY         : "lightgrey",
    CL_LIGHTBLUE         : "lightblue",
    CL_LIGHTRED          : "lightred",
    CL_LIGHTGREEN        : "lightgreen",
    CL_LIGHTYELLOW       : "lightyellow",
    CL_LIGHTMAGENTA      : "lightmagenta",
    CL_LIGHTCYAN         : "lightcyan",
    CL_LILAC             : "lilac",
    CL_TURQUOISE         : "turquoise",
    CL_AQUAMARINE        : "aquamarine",
    CL_KHAKI             : "khaki",
    CL_PURPLE            : "purple",
    CL_YELLOWGREEN       : "yellowgreen",
    CL_PINK              : "pink",
    CL_ORANGE            : "orange",
    CL_ORCHID            : "orchid",
    CL_BLACK             : "black",
}

def get_color_name(c):
    return COLORS_LUT[c] if c in COLORS_LUT.keys() else "?"

class cfunc_graph_t: # alas we can't inherit gdl_graph_t
    def __init__(self, highlight):
        self.items = [] # list of citem_t
        self.highlight = highlight
        self.succs = [] # list of lists of next nodes
        self.preds = [] # list of lists of previous nodes

    def nsucc(self, n):
        return len(self.succs[n]) if self.size() else 0

    def npred(self, n):
        return len(self.preds[n]) if self.size() else 0

    def succ(self, n, i):
        return self.succs[n][i]

    def pred(self, n, i):
        return self.preds[n][i]

    def size(self):
        return len(self.preds)

    def add_node(self):
        n = self.size()

        def resize(array, new_size):
            if new_size > len(array):
                while len(array) < new_size:
                    array.append([])
            else:
                array = array[:new_size]
            return array

        self.preds = resize(self.preds, n+1)
        self.succs = resize(self.succs, n+1)
        return n

    def add_edge(self, x, y):
        self.preds[y].append(x)
        self.succs[x].append(y)

    def get_expr_name(self, expr):
        name = expr.print1(None)
        name = ida_lines.tag_remove(name)
        name = ida_pro.str2user(name)
        return name

    def get_node_label(self, n):
        item = self.items[n]
        op = item.op
        insn = item.cinsn
        expr = item.cexpr
        parts = [ida_hexrays.get_ctype_name(op)]
        if op == ida_hexrays.cot_ptr:
            parts.append(".%d" % expr.ptrsize)
        elif op == ida_hexrays.cot_memptr:
            parts.append(".%d (m=%d)" % (expr.ptrsize, expr.m))
        elif op == ida_hexrays.cot_memref:
            parts.append(" (m=%d)" % (expr.m,))
        elif op in [
                ida_hexrays.cot_obj,
                ida_hexrays.cot_var]:
            name = self.get_expr_name(expr)
            parts.append(".%d %s" % (expr.refwidth, name))
        elif op in [
                ida_hexrays.cot_num,
                ida_hexrays.cot_helper,
                ida_hexrays.cot_str]:
            name = self.get_expr_name(expr)
            parts.append(" %s" % (name,))
        elif op == ida_hexrays.cit_goto:
            parts.append(" LABEL_%d" % insn.cgoto.label_num)
        elif op == ida_hexrays.cit_asm:
            parts.append("<asm statements; unsupported ATM>")
            # parts.append(" %a.%d" % ())
        parts.append("\\n")
        parts.append("ea: %08X" % item.ea)
        if item.is_expr() and not expr.type.empty():
            parts.append("\\n")
            tstr = expr.type._print()
            parts.append(tstr if tstr else "?")
        return "".join(parts)

    def get_node_color(self, n):
        item = self.items[n]
        if self.highlight is not None and item.obj_id == self.highlight.obj_id:
            return CL_GREEN
        return None

    def gen_gdl(self, fname):
        with open(fname, "wb") as out:
            out.write(b"graph: {\n")

            out.write(b"// *** nodes\n")
            for n in range(len(self.items)):
                item = self.items[n]
                node_label = self.get_node_label(n)
                node_props = [""]
                if n == 0:
                    node_props.append("vertical_order: 0")
                color = self.get_node_color(n)
                if color is not None:
                    node_props.append("color: %s" % get_color_name(color))
                #properties = " ".join(node_props)
                node = "node: { title: \"%d\" label: \"%d: %s\" %s}\n" % (
                    n,
                    n,
                    node_label,
                    " ".join(node_props))
                out.write(bytes(node, encoding="utf-8"))

            out.write(b"// *** edges\n")
            for n in range(len(self.items)):
                item = self.items[n]
                out.write(b"// edges %d -> ?\n" % n)
                for i in range(self.nsucc(n)):
                    t = self.succ(n, i)
                    label = ""
                    if item.is_expr():
                        target = self.items[t]
                        if item.x is not None and item.x == target:
                            label = "x"
                        elif item.y is not None and item.y == target:
                            label = "y"
                        elif item.z is not None and item.z == target:
                            label = "z"
                        if label:
                            label = """ label: "%s" """ % label
                    edge = """edge: { sourcename: "%s" targetname: "%s"%s}\n""" % (
                        str(n), str(t), label)
                    out.write(bytes(edge, encoding="utf-8"))

            out.write(b"}\n")

    def dump(self):
        idaapi.msg("%d items:" % len(self.items))
        for i in self.items:
            idaapi.msg("\t%s (%08x)" % (i, i.ea))

        idaapi.msg("succs:")
        for s in self.succs:
            idaapi.msg("\t%s" % s)

        idaapi.msg("preds:")
        for p in self.preds:
            idaapi.msg("\t%s" % p)


class graph_builder_t(ida_hexrays.ctree_parentee_t):

    def __init__(self, cg):
        ida_hexrays.ctree_parentee_t.__init__(self)
        self.cg = cg
        self.reverse = {} # citem_t -> node#

    def add_node(self, i):
        for k_obj_id in self.reverse.keys():
            if i.obj_id == k_obj_id:
                ida_kernwin.warning("bad ctree - duplicate nodes! (i.ea=%x)" % i.ea)
                return -1

        n = self.cg.add_node()
        if n <= len(self.cg.items):
            self.cg.items.append(i)
        self.cg.items[n] = i
        self.reverse[i.obj_id] = n
        return n

    def process(self, i):
        n = self.add_node(i)
        if n < 0:
            return n
        if len(self.parents) > 1:
            lp = self.parents.back().obj_id
            for k_obj_id, v in self.reverse.items():
                if k_obj_id == lp:
                    p = v
                    break
            self.cg.add_edge(p, n)
        return 0

    def visit_insn(self, i):
        return self.process(i)

    def visit_expr(self, e):
        return self.process(e)


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

    def activate(self, ctx):
        vu = ida_hexrays.get_widget_vdui(ctx.widget)
        vu.get_current_item(ida_hexrays.USE_KEYBOARD)
        highlight = vu.item.e if vu.item.is_citem() else None

        cg = cfunc_graph_t(highlight)
        gb = graph_builder_t(cg)
        gb.apply_to(vu.cfunc.body, None)

        import tempfile
        fname = tempfile.mktemp(suffix=".gdl")
        cg.gen_gdl(fname)
        ida_gdl.display_gdl(fname)
        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 vds5_hooks_t(ida_hexrays.Hexrays_Hooks):
    def populating_popup(self, widget, handle, vu):
        idaapi.attach_action_to_popup(vu.ct, None, ACTION_NAME)
        return 0

class Hrgraph(idaapi.plugin_t):
    comment = ''
    help = ''
    wanted_name = 'hrgraph'
    wanted_hotkey = ''
    hxehook = None
    flags = idaapi.PLUGIN_HIDE

    def init(self):
        self.vds5_hooks = None
        if not ida_hexrays.init_hexrays_plugin():
            idaapi.msg("hexrays-graph: hexrays is not available.")
            return idaapi.PLUGIN_SKIP
            
        ida_kernwin.register_action(
            ida_kernwin.action_desc_t(
                ACTION_NAME,
                "Hex-Rays show C graph (IDAPython)",
                display_graph_ah_t(),
                ACTION_SHORTCUT))
        self.vds5_hooks = vds5_hooks_t()
        self.vds5_hooks.hook()
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        pass

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


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


repository: tmr232/ida-plugins
file: autoenum/autoenum.py
description: IDA Pro Plugins
wanted_name: AutoEnum
comment: Automatic Enum Generation
created_at: 2015-05-27 12:35:29+00:00
pushed_at: 2015-05-31 13:05:09+00:00
url:
https://github.com/tmr232/ida-plugins/blob/c333ebecc9b3d641cdfdc9ab0d66f73e0b551fc3/autoenum/autoenum.py

content:
from collections import defaultdict
import idaapi
import sark
import idc
import sark.exceptions


def apply_enum_by_name(enum, member_name):
    member_value = enum.members[member_name].value
    for line in sark.lines(*sark.get_selection()):
        for operand in line.insn.operands:
            if operand.type.is_imm:
                if operand.imm == member_value:
                    idc.OpEnumEx(line.ea, operand.n, enum.eid, enum.members[member_name].serial)

            elif operand.type.is_displ or operand.type.is_phrase:
                if operand.addr == member_value:
                    idc.OpEnumEx(line.ea, operand.n, enum.eid, enum.members[member_name].serial)


def get_common_value():
    values = defaultdict(int)
    for line in sark.lines(*sark.get_selection()):
        for operand in line.insn.operands:
            if operand.type.is_imm:
                values[operand.imm] += 1

            elif operand.type.is_displ or operand.type.is_phrase:
                values[operand.addr] += 1

    # Ignore 0 as it is usually not interesting
    values[0] = 0
    # Get the most common value
    common_value = max(values.iteritems(), key=lambda x: x[1])[0]
    return common_value


def const_name(enum, value):
    return "{}_{:X}h".format(enum.name.upper(), value)


class AutoEnum(idaapi.plugin_t):
    flags = 0
    comment = "Automatic Enum Generation"
    help = "Automatic Enum Generation"
    wanted_name = "AutoEnum"
    wanted_hotkey = "Shift+M"

    def init(self):
        self._last_enum = ""
        return idaapi.PLUGIN_KEEP

    def term(self):
        pass

    def run(self, arg):
        common_value = get_common_value()

        enum_name = idc.AskStr(self._last_enum, "Enum Name")
        if enum_name is None:
            return

        if not enum_name:
            enum_name = None

        self._last_enum = enum_name

        const_value = idc.AskLong(common_value, "Const Value")
        if const_value is None:
            return

        modify = True

        try:
            enum = sark.add_enum(enum_name)

        except sark.exceptions.EnumAlreadyExists:
            enum = sark.Enum(enum_name)
            yes_no_cancel = idc.AskYN(idaapi.ASKBTN_NO,
                                      "Enum already exists. Modify?\n")
            if yes_no_cancel == idaapi.ASKBTN_CANCEL:
                return

            elif yes_no_cancel == idaapi.ASKBTN_YES:
                modify = True

            else:  # yes_no_cancel == idaapi.ASKBTN_NO:
                modify = False


        member_name = const_name(enum, const_value)

        if modify:

            try:
                enum.members.add(member_name, const_value)
            except sark.exceptions.SarkErrorAddEnumMemeberFailed as ex:
                idaapi.msg("[AutoEnum] Adding struct member failed: {}.".format(ex.message))


        else:
            for member in enum.members:
                if member.value == const_value:
                    member_name = member.name
                    break
                else:
                    return


        # Apply the enum
        apply_enum_by_name(enum, member_name)


def PLUGIN_ENTRY():
    return AutoEnum()


repository: clovme/WTools
file: software/Reverse Cracking/IDA Professional 9.0/python/examples/hexrays/vds19.py
wanted_name: optimize x|~x
created_at: 2025-03-27 00:14:54+00:00
pushed_at: 2025-05-11 10:18:21+00:00
url:
https://github.com/clovme/WTools/blob/f12f01c833a6a6241799a9d01e9117cc28a65dbc/software/Reverse%20Cracking/IDA%20Professional%209.0/python/examples/hexrays/vds19.py

content:
"""
summary: a custom microcode instruction optimization rule (`x | ~x => -1`)

description:
  Installs a custom microcode instruction optimization rule,
  to transform:

      x | ~x

  into

      -1

  To see this plugin in action please use be_ornot_be.idb
"""

import ida_hexrays
import ida_idaapi

# recognize "x | ~x" and replace by -1
class subinsn_optimizer_t(ida_hexrays.minsn_visitor_t):
    cnt = 0
    def visit_minsn(self):      # for each instruction...
        ins = self.curins       # take a reference to the current instruction

        # THE CORE OF THE PLUGIN IS HERE:
        # check the pattern "x | ~x"
        if ins.opcode == ida_hexrays.m_or and ins.r.is_insn(ida_hexrays.m_bnot) and ins.l == ins.r.d.l:
            if not ins.l.has_side_effects(): # avoid destroying side effects
                # pattern matched, convert to "mov -1, ..."
                ins.opcode = ida_hexrays.m_mov
                ins.l.make_number(-1, ins.r.size)
                ins.r = ida_hexrays.mop_t()
                self.cnt = self.cnt + 1 # number of changes we made
        return 0 # continue traversal

# a custom instruction optimizer, boilerplate code
class sample_optimizer_t(ida_hexrays.optinsn_t):
    def func(self, blk, ins, optflags):
        opt = subinsn_optimizer_t()
        ins.for_all_insns(opt)
        if opt.cnt != 0:                # if we modified microcode,
            blk.mba.verify(True)        # run the verifier
        return opt.cnt                  # report the number of changes

# a plugin interface, boilerplate code
class my_plugin_t(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_HIDE
    wanted_name = "optimize x|~x"
    wanted_hotkey = ""
    comment = ""
    help = ""
    def init(self):
        if ida_hexrays.init_hexrays_plugin():
            self.optimizer = sample_optimizer_t()
            self.optimizer.install()
            print("Installed sample optimizer for 'x | ~x'")
            return ida_idaapi.PLUGIN_KEEP # keep us in the memory
    def term(self):
        self.optimizer.remove()
    def run(self, arg):
        if arg == 1:
            return self.optimizer.remove()
        elif arg == 2:
            return self.optimizer.install()

def PLUGIN_ENTRY():
    return my_plugin_t()


repository: zmrbak/IDAPython
file: 配套代码/ida10/python/examples/hexrays/vds19.py
description: 在线课程《IDAPython脚本化软件逆向分析入门》课程源码
wanted_name: optimize x|~x
created_at: 2022-05-02 02:13:34+00:00
pushed_at: 2022-05-14 10:41:43+00:00
url:
https://github.com/zmrbak/IDAPython/blob/478278339ba0c150b3ace657296233d8a89b26ea/%E9%85%8D%E5%A5%97%E4%BB%A3%E7%A0%81/ida10/python/examples/hexrays/vds19.py

content:
"""
summary: a custom microcode instruction optimization rule (`x | ~x => -1`)

description:
  Installs a custom microcode instruction optimization rule,
  to transform:

      x | ~x

  into

      -1

  To see this plugin in action please use be_ornot_be.idb
"""

import ida_hexrays
import ida_idaapi

# recognize "x | ~x" and replace by -1
class subinsn_optimizer_t(ida_hexrays.minsn_visitor_t):
    cnt = 0
    def visit_minsn(self):      # for each instruction...
        ins = self.curins       # take a reference to the current instruction

        # THE CORE OF THE PLUGIN IS HERE:
        # check the pattern "x | ~x"
        if ins.opcode == ida_hexrays.m_or and ins.r.is_insn(ida_hexrays.m_bnot) and ins.l == ins.r.d.l:
            if not ins.l.has_side_effects(): # avoid destroying side effects
                # pattern matched, convert to "mov -1, ..."
                ins.opcode = ida_hexrays.m_mov
                ins.l.make_number(-1, ins.r.size)
                ins.r = ida_hexrays.mop_t()
                self.cnt = self.cnt + 1 # number of changes we made
        return 0 # continue traversal

# a custom instruction optimizer, boilerplate code
class sample_optimizer_t(ida_hexrays.optinsn_t):
    def func(self, blk, ins, optflags):
        opt = subinsn_optimizer_t()
        ins.for_all_insns(opt)
        if opt.cnt != 0:                # if we modified microcode,
            blk.mba.verify(True)        # run the verifier
        return opt.cnt                  # report the number of changes

# a plugin interface, boilerplate code
class my_plugin_t(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_HIDE
    wanted_name = "optimize x|~x"
    wanted_hotkey = ""
    comment = ""
    help = ""
    def init(self):
        if ida_hexrays.init_hexrays_plugin():
            self.optimizer = sample_optimizer_t()
            self.optimizer.install()
            print("Installed sample optimizer for 'x | ~x'")
            return ida_idaapi.PLUGIN_KEEP # keep us in the memory
    def term(self):
        self.optimizer.remove()
    def run(self, arg):
        if arg == 1:
            return self.optimizer.remove()
        elif arg == 2:
            return self.optimizer.install()

def PLUGIN_ENTRY():
    return my_plugin_t()


repository: pfalcon/idapython
file: Scripts/AsmViewer.py
description: Automatically exported from code.google.com/p/idapython
wanted_name: ASM file viewer
comment: ASM viewer
created_at: 2015-04-14 07:52:06+00:00
pushed_at: 2015-04-14 07:54:41+00:00
url:
https://github.com/pfalcon/idapython/blob/081b988a03b88867786ad4131269db6930637a5b/Scripts/AsmViewer.py

content:
# -----------------------------------------------------------------------
# This is an example illustrating how to use customview in Python
# The sample will allow you to open an assembly file and display it in color
# (c) Hex-Rays
#
import idaapi
import idautils
import idc
import os

# ----------------------------------------------------------------------
class asm_colorizer_t(object):
    def is_id(self, ch):
        return ch == '_' or ch.isalpha() or '0' <= ch <= '9'

    def get_identifier(self, line, x, e):
        i = x
        is_digit = line[i].isdigit()
        while i < e:
            ch = line[i]
            if not self.is_id(ch):
                if ch != '.' or not is_digit:
                    break
            i += 1
        return (i, line[x:i])

    def get_quoted_string(self, line, x, e):
        quote = line[x]
        i = x + 1
        while i < e:
            ch = line[i]
            if ch == '\\' and line[i+1] == quote:
                i += 1
            elif ch == quote:
                i += 1 # also take the quote
                break
            i += 1
        return (i, line[x:i])

    def colorize(self, lines):
        for line in lines:
            line = line.rstrip()
            if not line:
                self.add_line()
                continue
            x = 0
            e = len(line)
            s = ""
            while x < e:
                ch = line[x]
                # String?
                if ch == '"' or ch == "'":
                    x, w = self.get_quoted_string(line, x, e)
                    s += self.as_string(w)
                # Tab?
                elif ch == '\t':
                    s += ' ' * 4
                    x += 1
                # Comment?
                elif ch == ';':
                    s += self.as_comment(line[x:])
                    # Done with this line
                    break
                elif ch == '.' and x + 1 < e:
                    x, w = self.get_identifier(line, x + 1, e)
                    s += self.as_directive(ch + w)
                # Identifiers?
                elif self.is_id(ch):
                    x, w = self.get_identifier(line, x, e)
                    # Number?
                    if ch.isdigit():
                        s += self.as_num(w)
                    # Other identifier
                    else:
                        s += self.as_id(w)
                # Output as is
                else:
                    s += ch
                    x += 1
            self.add_line(s)

# -----------------------------------------------------------------------
class asmview_t(idaapi.simplecustviewer_t, asm_colorizer_t):
    def Create(self, fn):
        # Create the customview
        if not idaapi.simplecustviewer_t.Create(self, "Viewing file - %s" % os.path.basename(fn)):
            return False

        self.instruction_list = idautils.GetInstructionList()
        self.instruction_list.extend(["ret"])
        self.register_list    = idautils.GetRegisterList()
        self.register_list.extend(["eax", "ebx", "ecx", "edx", "edi", "esi", "ebp", "esp"])

        self.fn = fn
        if not self.reload_file():
            return False

        self.id_refresh = self.AddPopupMenu("Refresh")
        self.id_close   = self.AddPopupMenu("Close")

        return True

    def reload_file(self):
        if not self.colorize_file(self.fn):
            self.Close()
            return False
        return True

    def colorize_file(self, fn):
        try:
            f = open(fn, "r")
            lines = f.readlines()
            f.close()
            self.ClearLines()
            self.colorize(lines)
            return True
        except:
            return False

    def add_line(self, s=None):
        if not s:
            s = ""
        self.AddLine(s)

    def as_comment(self, s):
        return idaapi.COLSTR(s, idaapi.SCOLOR_RPTCMT)

    def as_id(self, s):
        t = s.lower()
        if t in self.register_list:
            return idaapi.COLSTR(s, idaapi.SCOLOR_REG)
        elif t in self.instruction_list:
            return idaapi.COLSTR(s, idaapi.SCOLOR_INSN)
        else:
            return s

    def as_string(self, s):
        return idaapi.COLSTR(s, idaapi.SCOLOR_STRING)

    def as_num(self, s):
        return idaapi.COLSTR(s, idaapi.SCOLOR_NUMBER)

    def as_directive(self, s):
        return idaapi.COLSTR(s, idaapi.SCOLOR_KEYWORD)

    def OnPopupMenu(self, menu_id):
        """
        A context (or popup) menu item was executed.
        @param menu_id: ID previously registered with AddPopupMenu()
        @return: Boolean
        """
        if self.id_refresh == menu_id:
            return self.reload_file()
        elif self.id_close == menu_id:
            self.Close()
            return True
        return False

    def OnKeydown(self, vkey, shift):
        """
        User pressed a key
        @param vkey: Virtual key code
        @param shift: Shift flag
        @return Boolean. True if you handled the event
        """
        # ESCAPE
        if vkey == 27:
            self.Close()
        elif vkey == ord('H'):
            lineno = self.GetLineNo()
            if lineno is not None:
                line, fg, bg = self.GetLine(lineno)
                if line and line[0] != idaapi.SCOLOR_INV:
                    s = idaapi.SCOLOR_INV + line + idaapi.SCOLOR_INV
                    self.EditLine(lineno, s, fg, bg)
                    self.Refresh()
        elif vkey == ord('C'):
            self.ClearLines()
            self.Refresh()
        elif vkey == ord('S'):
            print "Selection (x1, y1, x2, y2) = ", self.GetSelection()
        elif vkey == ord('I'):
            print "Position (line, x, y) = ", self.GetPos(mouse = 0)
        else:
            return False
        return True

# -----------------------------------------------------------------------
class asmviewplg(idaapi.plugin_t):
    flags = idaapi.PLUGIN_KEEP
    comment = "ASM viewer"
    help = "This is help"
    wanted_name = "ASM file viewer"
    wanted_hotkey = "Alt-F8"
    def __init__(self):
        self.view = None

    def init(self):
        return idaapi.PLUGIN_KEEP
    def run(self, arg):
        if self.view:
            self.Close()
        fn = idc.AskFile(0, "*.asm", "Select ASM file to view")
        if not fn:
            return
        self.view = asmview_t()
        if not self.view.Create(fn):
            return
        self.view.Show()

    def term(self):
        if self.view:
            self.view.Close()

def PLUGIN_ENTRY():
    return asmviewplg()


repository: FBLeee/IDA_stalker_trace_so_
file: stalker_trace_so.py
wanted_name: stalker_trace_so
comment: stalker_trace_so
created_at: 2024-04-19 01:05:26+00:00
pushed_at: 2025-02-24 01:47:21+00:00
description:
可以一键生成对指定函数hook的初步脚本 or trace，根据oacia / stalker_trace_so添加一些过滤规则，可以选择一键生成【trace所有函数的脚本】 or 【trace指定函数的脚本】

url:
https://github.com/FBLeee/IDA_stalker_trace_so_/blob/dac1a56adb83d4ebe43792cb3758a23ae715e277/stalker_trace_so.py

content:
# -*- coding:utf-8 -*-
import os
from idaapi import plugin_t
from idaapi import PLUGIN_PROC
from idaapi import PLUGIN_OK
import ida_nalt
import idautils
import idc
import ida_kernwin


class stalker_trace_so(plugin_t):
    flags = PLUGIN_PROC
    comment = "stalker_trace_so"
    help = ""
    wanted_name = "stalker_trace_so"
    wanted_hotkey = ""

    def init(self):
        print("stalker_trace_so plugin has been loaded.")
        return PLUGIN_OK

    def run(self, arg):
        
        func_addr = []
        func_name = []
        filter_string = ida_kernwin.ask_str("Default value", 0, "输入你需要过滤的特征串：")
        print("filter_string：",filter_string)
        for func_ea in idautils.Functions():
            
            #过滤某些字符串
            all_func_name = idc.get_func_name(func_ea);
            if(filter_string in all_func_name):
                func_name.append('"{}"'.format(idc.get_func_name(func_ea)))
                # thumb mode
                if idc.get_sreg(func_ea, "T"):
                    func_addr.append(hex(func_ea + 1))
                else:
                    func_addr.append(hex(func_ea))
            
            
            
        template_js = 'var func_addr = [{}];\nvar func_name = [{}];\nvar so_name = "{}";\nfunction hook_dlopen() {{\n    Interceptor.attach(Module.findExportByName(null, "android_dlopen_ext"),\n        {{\n            onEnter: function (args) {{\n                var pathptr = args[0];\n                if (pathptr !== undefined && pathptr != null) {{\n                    var path = ptr(pathptr).readCString();\n                    //console.log(path);\n                    if (path.indexOf(so_name) >= 0) {{\n                        this.is_can_hook = true;\n                    }}\n                }}\n            }},\n            onLeave: function (retval) {{\n                if (this.is_can_hook) {{\n                    //you can do any thing before stalker trace so\n                    \n                    trace_so();\n                }}\n            }}\n        }}\n    );\n}}\n\nfunction trace_so(){{\n    var times = 1;\n    var module = Process.getModuleByName(so_name);\n    var pid = Process.getCurrentThreadId();\n    console.log("start Stalker!");\n    Stalker.exclude({{\n        "base": Process.getModuleByName("libc.so").base,\n        "size": Process.getModuleByName("libc.so").size\n    }})\n    Stalker.follow(pid,{{\n        events:{{\n            call:false,\n            ret:false,\n            exec:false,\n            block:false,\n            compile:false\n        }},\n        onReceive:function(events){{\n        }},\n        transform: function (iterator) {{\n            var instruction = iterator.next();\n            do{{\n                if (func_addr.indexOf(instruction.address - module.base) != -1){{\n                    console.log("call" + times+ ":" + func_name[func_addr.indexOf(instruction.address - module.base)])\n                    times=times+1\n                }}\n                iterator.keep();\n            }} while ((instruction = iterator.next()) !== null);\n        }},\n\n        onCallSummary:function(summary){{\n\n        }}\n    }});\n    console.log("Stalker end!");\n}}\n\nsetImmediate(hook_dlopen);\n'
        so_path, so_name = os.path.split(ida_nalt.get_input_file_path())
        hook_code = template_js.format(', '.join(func_addr), ', '.join(func_name), so_name)

        script_name = "trace_" + so_name.split(".")[0] + ".js"
        save_path = os.path.join(so_path, script_name)
        with open(save_path, "w", encoding="utf-8") as f:
            f.write(hook_code)

        print("usage:")
        print(f'frida -U -l "{save_path}" -f [package name]')

    def term(self):
        pass


def PLUGIN_ENTRY():
    return stalker_trace_so()


repository: benwoo1110/PMA-Notes
file: Plugins/hexcopy.py
comment: Copy Hex Bytes
created_at: 2022-09-27 01:21:48+00:00
pushed_at: 2022-12-10 06:36:38+00:00
url:
https://github.com/benwoo1110/PMA-Notes/blob/9d322ef0cbb99709f9227e322846b58dd1b1c9ab/Plugins/hexcopy.py

content:
"""
Adapted from https://github.com/OALabs/hexcopy-ida
"""

############################################################################################
##
## One-Click Hex Copy!
##
## Updated for IDA 7.xx and Python 3
##
## To install:
##      Copy script into plugins directory, i.e: C:\Program Files\<ida version>\plugins
##
## To run:
##      Highlight disassembly instructions and right click and select "Hex Copy" 
##      The hex encoded bytes will be added to your clipboard
##
############################################################################################


__AUTHOR__ = '@herrcore'

PLUGIN_NAME = "Hex Copy"
PLUGIN_HOTKEY = 'Ctrl+H'
VERSION = '3.2.0'


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)

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 copy_to_clip(data):
    QApplication.clipboard().setText(data)


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

class hex_copy(idaapi.plugin_t):
    """
    The IDA Plugin for hex copy.
    """

    flags = idaapi.PLUGIN_PROC | idaapi.PLUGIN_HIDE
    comment = "Copy Hex Bytes"
    help = "Highlight Assembly and right-click 'Copy Hex'"
    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_copy_bytes()

        # 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_copy_bytes()


        # 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_COPY_BYTES  = "prefix:copy_bytes"


    def _init_action_copy_bytes(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.ACTION_COPY_BYTES,         # The action name.
                "Copy Hex",                     # The action text.
                IDACtxEntry(copy_bytes_py3),        # The action handler.
                PLUGIN_HOTKEY,                  # Optional: action shortcut
                "Copy selected bytes as hex",   # Optional: tooltip
                31                              # Copy icon
            )
        else:
            # Describe the action using python2 copy
            action_desc = idaapi.action_desc_t(
                self.ACTION_COPY_BYTES,         # The action name.
                "Copy Hex",                     # The action text.
                IDACtxEntry(copy_bytes_py2),        # The action handler.
                PLUGIN_HOTKEY,                  # Optional: action shortcut
                "Copy selected bytes as hex",   # Optional: tooltip
                31                              # Copy icon
            )


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


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




#------------------------------------------------------------------------------
# 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_hex_copy_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_hex_copy_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,
                hex_copy.ACTION_COPY_BYTES,
                "Copy Hex",
                idaapi.SETMENU_APP,
            )

        # done
        return 0

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


def inject_hex_copy_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,
            hex_copy.ACTION_COPY_BYTES,
            "Copy Hex",
            idaapi.SETMENU_APP
        )

    # done
    return 0

#------------------------------------------------------------------------------
# Byte copy 
#------------------------------------------------------------------------------

def copy_bytes_py2():
    """
    Copy selected bytes to clipboard
    """
    if using_ida7api:
        start = idc.read_selection_start()
        end = idc.read_selection_end()
        if idaapi.BADADDR in (start, end):
            ea = idc.here()
            start = idaapi.get_item_head(ea)
            end = idaapi.get_item_end(ea)
        data = idc.get_bytes(start, end - start).encode('hex')
        print("Bytes copied: %s" % data)
        copy_to_clip(data)
    else:
        start = idc.SelStart()
        end = idc.SelEnd()
        if idaapi.BADADDR in (start, end):
            ea = idc.here()
            start = idaapi.get_item_head(ea)
            end = idaapi.get_item_end(ea)
        data = idc.GetManyBytes(start, end-start).encode('hex')
        print("Bytes copied: %s" % data)
        copy_to_clip(data)
    return 


def copy_bytes_py3():
    """
    Copy selected bytes to clipboard
    """
    if using_ida7api:
        start = idc.read_selection_start()
        end = idc.read_selection_end()
        if idaapi.BADADDR in (start, end):
            ea = idc.here()
            start = idaapi.get_item_head(ea)
            end = idaapi.get_item_end(ea)
        # fix encode bug reference 
        # https://stackoverflow.com/questions/6624453/whats-the-correct-way-to-convert-bytes-to-a-hex-string-in-python-3
        data = idc.get_bytes(start, end - start).hex()
        print ("Bytes copied: %s" % data)
        copy_to_clip(data)
    else:
        start = idc.SelStart()
        end = idc.SelEnd()
        if idaapi.BADADDR in (start, end):
            ea = idc.here()
            start = idaapi.get_item_head(ea)
            end = idaapi.get_item_end(ea)
        data = idc.GetManyBytes(start, end-start).hex()
        print( "Bytes copied: %s" % data)
        copy_to_clip(data)
    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: scottmudge/DebugAutoPatch
file: DebugAutoPatch.py
description: Patching system improvement plugin for IDA.
wanted_name: DebugAutoPatch
comment: Plugin for automatic byte patch injection - no binary-file patching needed!
created_at: 2019-08-29 22:19:34+00:00
pushed_at: 2019-09-07 21:12:02+00:00
url:
https://github.com/scottmudge/DebugAutoPatch/blob/b86d784ab1c99337e8196693fef80e42ba2200cf/DebugAutoPatch.py

content:
# DebugAutoPatch IDA Plugin - Release 0.2
# By Scott Mudge, 2019 -- https://scottmudge.com.
#
# NOTE: This has been tested with IDA 7.0 - I have made attempts at backward/forward compatibility,
# but please report bugs for other versions.
#
# DebugAutoPatch is released under the GNU GPLv3 license. See LICENSE for more information.
# Find information and latest version at https://github.com/scottmudge/DebugAutoPatch
#
# This IDA plugin automatically applies byte patches stored in the NON-debug IDA "Patched bytes" database
# to the debugged process at runtime. It does this at (by default) the entry-point of the application (or DLL).
# The process will then automatically resume with the patched bytes set in memory.
#
# Why? Making modifications to application/.rdata code can be tedious, IDA in particular. First the patches must be
# made with the clunky patching tools, and then the binary must be patched on-disk, followed by re-executing the
# application. Furthermore, patching the actual binary introduces a number of potential issues which could be
# mitigated by leaving it untouched. For instance, if the module or application performs hash checks to ensure it
# has not been modified.
#
# Developer Notes:
# --------------------
# Change Log:
#   * Just see the commit logs.
#

from threading import Thread, Lock, Event
import logging
import idaapi
import os
import idc
import json


#  ----------------------------------------- Globals -----------------------------------------
DAP_VERSION = "0.2"
DAP_NAME = "DebugAutoPatch"
DAP_CONFIG_FILE_PATH = os.path.join(idc.GetIdaDirectory(), 'cfg', 'DebugAutoPatch.cfg')
DAP_WEBSITE = "https://github.com/scottmudge/DebugAutoPatch"
DEBUG_MESSAGE_LEVEL = logging.INFO
DAP_INITIALIZED = False
DAP_INSTANCE = None
#  ---------------------------------------------------------------------------------------------


#  ----------------------------------------- Utilities -----------------------------------------
def dap_msg(string):
    print("[{}]: {}".format(DAP_NAME, string))


def dap_warn(string, details = None):
    if details:
        print("[{} | WARNING]: {}\n\t> Details: {}".format(DAP_NAME, string, details))
    else:
        print("[{} | WARNING]: {}".format(DAP_NAME, string))


def dap_err(string, details = None):
    if details:
        print("[{} | ERROR]: {}\n\t> Details: {}".format(DAP_NAME, string, details))
    else:
        print("[{} | ERROR]: {}".format(DAP_NAME, string))


class KillableThread(Thread):
    """Wraps a killable thread that loops at a preset interval. Runs supplied
    target function.
    """
    def __del__(self):
        self.kill()

    def __init__(self, name, target, sleep_interval):
        """
        Args:
            name: Name of the thread, used for logging.
            target (function): Target function
            sleep_interval (float): Sleep interval seconds between loops.
        """
        super(KillableThread, self).__init__(group=None, target=target, name=name)
        self._trigger = Event()
        self._interval = sleep_interval
        self._target = target
        self._name = name
        self._kill = False
        self.setDaemon(True)

    def trigger(self):
        """Triggers loop, but does not kill it."""
        self._kill = False
        self._trigger.set()

    def run(self):
        """Runs the thread."""
        while True:
            try:
                self._target()
                # If no kill signal is set, sleep for the interval,
                # If kill signal comes in while sleeping, immediately
                #  wake up and handle
                is_triggerer = self._trigger.wait(timeout=self._interval)
                if is_triggerer:
                    if self._kill:
                        break
                    else:
                        self._trigger.clear()
            except(KeyboardInterrupt, SystemExit):
                self.kill()
                continue
        dap_msg("Thread killed! [name={}]".format(self._name))

    def kill(self):
        """Kills the thread."""
        dap_msg("Killing thread... [name={}]".format(self._name))
        self._kill = True
        self._trigger.set()
#  ---------------------------------------------------------------------------------------------


class DapCfg:
    def __init__(self):
        pass

    Enabled = "enabled"
    PrimaryPatchAddr = "primary_patch_addr"


# About form
class DAPAboutForm(idaapi.Form):
    def __init__(self):
        # create About form
        super(DAPAboutForm, self).__init__(
            r"""STARTITEM 0
BUTTON YES* Open DebugAutoPatch Website
DebugAutoPatch - About

           {FormChangeCb}
           DebugAutoPatch IDA plugin v%s.
           (c) Scott Mudge, 2019.

           DebugAutoPatch is released under the GPL v3.
           Find more info at %s
           """ % (DAP_VERSION, DAP_WEBSITE), {
                'FormChangeCb': self.FormChangeCb(self.OnFormChange),
            })

        self.Compile()

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

        return 1


# Create menu handlers for IDA >= 700
try:
    # noinspection PyBroadException
    class DapMenuContext(idaapi.action_handler_t):
        label = None

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

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

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

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

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

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


    class DapMCEnable(DapMenuContext):
        def activate(self, ctx):
            self.plugin.enable_patching()
            return 1


    class DapMCDisable(DapMenuContext):
        def activate(self, ctx):
            self.plugin.disable_patching()
            return 1


    class DapMCCheckUpdate(DapMenuContext):
        def activate(self, ctx):
            self.plugin.check_update()
            return 1


    class DapMCAbout(DapMenuContext):
        def activate(self, ctx):
            self.plugin.about()
            return 1


    class DapMCApplyPatch(DapMenuContext):
        def activate(self, ctx):
            self.plugin.apply_patch_to_memory()
            return 1


    class DapMCApplyPatchesToProc(DapMenuContext):
        def activate(self, ctx):
            self.plugin.apply_patches_to_current_proc()
            return 1


    class DapMCNull(DapMenuContext):
        def activate(self, ctx):
            self.plugin.menu_null()
            return 1


    class DapMCNull2(DapMenuContext):
        def activate(self, ctx):
            self.plugin.menu_null()
            return 1
except:
    pass


class DebugAutoPatchPlugin(idaapi.plugin_t):
    # This keeps the plugin in memory, important for hooking callbacks
    flags = idaapi.PLUGIN_KEEP
    comment = "Plugin for automatic byte patch injection - no binary-file patching needed!"
    help = "See https://github.com/scottmudge/IDA_DebugAutoPatch/blob/master/readme.md"
    wanted_name = "DebugAutoPatch"
    wanted_hotkey = ""

    def __del__(self):
        self.term()

    def __init__(self):
        self.old_ida = False
        self.cfg = None
        self.debug_hook = None
        self.patched_bytes_db = []
        self.patched_bytes_db_lock = Lock()
        self.monitor_thread = None

    class PatchedByte:
        """Container for patched byte type."""
        def __init__(self, addr, orig, patched):
            self.addr = addr
            self.orig = orig
            self.patched = patched

    class PatchVisitor(object):
        """Used for visiting patched bytes when debugger is not active. These patches are then stored in a buffer,
        and are applied when debugger activates."""
        def __init__(self):
            self.skipped = 0
            self.patched = 0
            self.patched_bytes = []

        def __call__(self, ea, fpos, orig, patch_val, cnt=()):
            try:
                if fpos == -1:
                    self.skipped += 1
                    dap_msg("fpos invalid ({}) -- patch skipped".format(fpos))
                else:
                    self.patched += 1
                    # dap_msg(" ea: %x \\ fpos: %x \\ o: %x \\ v: %x" % (ea, fpos, orig, patch_val))
                    self.patched_bytes.append(DebugAutoPatchPlugin.PatchedByte(ea, orig, patch_val))
                return 0
            except:
                return

    class DebugHook(idaapi.DBG_Hooks):
        def __init__(self, *args):
            super(DebugAutoPatchPlugin.DebugHook, self).__init__(*args)
            self.steps = 0

        def dbg_process_start(self, pid, tid, ea, name, base, size):
            dap_msg("Process start hook snagged -- applying patches...")
            result = DAP_INSTANCE.apply_patches_to_current_proc()
            if result >= 0:
                dap_msg("Success!")

    def init(self):
        """Initialization routine."""
        global DAP_INITIALIZED

        if idaapi.IDA_SDK_VERSION < 700:
            self.old_ida = True

        # register menu handlers
        try:
            DapMCNull.register(self, "_________________________")
            DapMCNull2.register(self, "_________________________")
            DapMCEnable.register(self, "Enable Auto-Patching")
            DapMCDisable.register(self, "Disable Auto-Patching")
            DapMCApplyPatch.register(self, "Apply Patch to Memory")
            DapMCApplyPatchesToProc.register(self, "Apply Patches to Current Process")
            DapMCCheckUpdate.register(self, "Check for DebugAutoPatch Update")
            DapMCAbout.register(self, "About DebugAutoPatch")
        except:
            pass

        self.cfg = None

        if not DAP_INITIALIZED:
            DAP_INITIALIZED = True

            if not self.old_ida:
                # Add menu IDA >= 7.0
                idaapi.attach_action_to_menu("Edit/Patch program/Null Menu", DapMCNull.get_name(), idaapi.SETMENU_APP)
                idaapi.attach_action_to_menu("Edit/Patch program/Enable Auto-Patching", DapMCEnable.get_name(),
                                             idaapi.SETMENU_APP)
                idaapi.attach_action_to_menu("Edit/Patch program/Disable Auto-Patching", DapMCDisable.get_name(),
                                             idaapi.SETMENU_APP)
                # idaapi.attach_action_to_menu("Edit/Patch program/Apply Patch to Memory", DapMCApplyPatch.get_name(),
                #                              idaapi.SETMENU_APP)
                # idaapi.attach_action_to_menu("Edit/Patch program/Apply Patches to Current Process",
                #                              DapMCApplyPatchesToProc.get_name(), idaapi.SETMENU_APP)
                idaapi.attach_action_to_menu("Edit/Patch program/Null Menu 2", DapMCNull2.get_name(),
                                             idaapi.SETMENU_APP)
                # idaapi.attach_action_to_menu("Edit/Patch program/Check for DebugAutoPatch Update",
                #                              DapMCCheckUpdate.get_name(), idaapi.SETMENU_APP)
                idaapi.attach_action_to_menu("Edit/Patch program/About DebugAutoPatch",
                                             DapMCAbout.get_name(), idaapi.SETMENU_APP)
            else:
                # Older versions
                idaapi.add_menu_item("Edit/Patch program/", "-", "", 1, self.menu_null, None)
                idaapi.add_menu_item("Edit/Patch program/", "Enable Auto-Patching", "", 1,
                                     self.enable_patching, None)
                idaapi.add_menu_item("Edit/Patch program/", "Disable Auto-Patching", "", 1,
                                     self.disable_patching, None)
                # idaapi.add_menu_item("Edit/Patch program/", "Apply Patch to Memory", "", 1,
                #                      self.apply_patch_to_memory, None)
                # idaapi.add_menu_item("Edit/Patch program/", "Apply Patches to Current Process", "", 1,
                #                      self.apply_patches_to_current_proc, None)
                idaapi.add_menu_item("Edit/Patch program/", "-", "", 1, self.menu_null, None)
                # idaapi.add_menu_item("Edit/Patch program/", "Check for DebugAutoPatch Update", "", 1,
                #                      self.check_update, None)
                idaapi.add_menu_item("Edit/Patch program/", "About DebugAutoPatch", "", 1, self.about, None)

            print("=" * 80)
            print("DebugAutoPatch v{0} Copyright (c) Scott Mudge 2019".format(DAP_VERSION))
            print("DebugAutoPatch is available from menu \"Edit > Patch program\"")
            print("Find more information about DebugAutoPatch at the project github repository")

            self.load_configuration()
            self.set_debug_hooks()

            # Update patch database first
            self.patch_monitor_func()

            dap_msg("Starting patch monitoring thread...")
            self.monitor_thread = KillableThread(name="PatchMonitoring", target=self.patch_monitor_func,
                                                 sleep_interval=0.7)
            self.monitor_thread.start()

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

    def patch_monitor_func(self):
        """Monitors patches and caches patch DB, since IDA has separate DBs for debugged processes and non-debugged
        processes."""
        # Don't collect patches if debugger is on
        try:
            if idaapi.is_debugger_on() or idaapi.is_debugger_busy():
                return

            if not self.patched_bytes_db_lock.acquire(False):
                return
            else:
                try:
                    was_empty = False
                    if len(self.patched_bytes_db) < 1:
                        was_empty = True
                    patches = self.visit_patched_bytes()
                    self.patched_bytes_db = patches
                    if len(patches) > 0 and was_empty:
                        dap_msg("Byte patch buffer populated!")
                finally:
                    self.patched_bytes_db_lock.release()
        except:
            pass

    def enable_patching(self):
        """Enables automatic patching."""
        self.cfg[DapCfg.Enabled] = True
        dap_msg("Automatic patching enabled.")
        pass

    def disable_patching(self):
        """Disables automatic patching."""
        self.cfg[DapCfg.Enabled] = False
        dap_msg("Automatic patching disabled.")
        pass

    def apply_patch_to_memory(self):
        """Adds a new patch to database."""
        # TODO -- Implement
        pass

    def apply_patches_to_current_proc(self):
        """Applies patches to current process. Must first suspend process, check debugger is not active, then
        apply them."""
        if not self.cfg[DapCfg.Enabled]:
            dap_msg("Not applying patches to current process - patching currently disabled.")
            return

        total_applied = 0
        if idaapi.suspend_process():
            self.patched_bytes_db_lock.acquire()
            try:
                if len(self.patched_bytes_db) < 1:
                    dap_msg("No patched bytes currently in database, nothing to do!")
                else:
                    for patch in self.patched_bytes_db:
                        total_applied += self.apply_byte_patch(patch)
                    dap_msg("[{}] total patches applied!".format(total_applied))
            except Exception as e:
                dap_err("Error encountered while applying patches to current debugged process.", str(e))
            except:
                dap_err("Unknown error encountered while applying patches to current debugged process.")
            finally:
                self.patched_bytes_db_lock.release()
        else:
            dap_err("Could not apply patches, could not suspend process!")
        idc.resume_process()
        return total_applied

    @staticmethod
    def about():
        """About window."""
        f = DAPAboutForm()
        f.Execute()
        f.Free()
        pass

    def check_update(self):
        """Checks for new version."""
        # TODO - Update
        pass

    def menu_null(self):
        """For menu item which does nothing."""
        pass

    def run(self, *args):
        """Used for when user selects plugin entry from Edit > Plugins"""
        self.about()
        pass

    def term(self):
        """Termination call."""
        if self.monitor_thread:
            self.monitor_thread.kill()
        self.unset_debug_hooks()
        self.save_configuration()

    def set_debug_hooks(self):
        """Installs debugger hooks for automatic patching."""
        self.unset_debug_hooks()
        dap_msg("Installing debug hooks...")
        self.debug_hook = DebugAutoPatchPlugin.DebugHook()
        self.debug_hook.hook()
        self.debug_hook.steps = 0
        dap_msg("Done!")

    def unset_debug_hooks(self):
        """Remove any installed debug hooks."""
        try:
            if self.debug_hook:
                dap_msg("Removing previously installed debugger hooks...")
                self.debug_hook.unhook()
                dap_msg("Done!")
        except:
            pass

    def apply_byte_patch(self, patched_byte_ojb):
        """Applies a byte patch to current debugger memory."""
        # check if debugger is even running
        if not idaapi.is_debugger_on():
            dap_warn("Cannot apply patch - debugger is not currently on!")
            return 0

        try:
            # patched byte in debugger memory
            if not self.old_ida:
                result = idc.patch_dbg_byte(patched_byte_ojb.addr, patched_byte_ojb.patched)
            else:
                result = idc.PatchDbgByte(patched_byte_ojb.addr, patched_byte_ojb.patched)
            if result > 0:
                idaapi.invalidate_dbgmem_contents(patched_byte_ojb.addr, 1) # addr, size
            return result
        except Exception as e:
            dap_err("Error encountered while applying byte patch to memory!", str(e))
        except:
            dap_err("Unknown error encountered while applying byte patch to memory!")
        return 0

    def visit_patched_bytes(self):
        """Iterates through patched bytes and stores them in a buffer."""
        try:
            visitor = self.PatchVisitor()
            result = idaapi.visit_patched_bytes(0, idaapi.BADADDR, visitor)
            if result != 0:
                dap_err("visit_patched_bytes() returned unexpected result", "error code ({})".format(result))
                return []
            return visitor.patched_bytes
        except Exception as e:
            dap_err("Exception encountered while visiting patched bytes", str(e))
        except:
            dap_err("Unknown")

    def load_configuration(self):
        """Loads configuration from disk."""
        self.cfg = {}
        save_cfg = False
        # load configuration from file
        try:
            f = open(DAP_CONFIG_FILE_PATH, "rt")
            self.cfg = json.load(f)
            f.close()
        except IOError:
            dap_msg("Failed to load config file -- using defaults.")
            save_cfg = True
        except Exception as e:
            dap_err("Failed to load config file.", str(e))

        # Enables or disables patching at debug time
        if DapCfg.Enabled not in self.cfg:
            self.cfg[DapCfg.Enabled] = True
        # Primary patched application address - set to BADADDR = use application start
        if DapCfg.PrimaryPatchAddr not in self.cfg:
            self.cfg[DapCfg.PrimaryPatchAddr] = idaapi.BADADDR
        if save_cfg:
            self.save_configuration()

    def save_configuration(self):
        """Saves configuration to disk."""
        if self.cfg:
            try:
                json.dump(self.cfg, open(DAP_CONFIG_FILE_PATH, "wt"))
            except Exception as e:
                dap_err("Failed to save configuration file", str(e))
            else:
                dap_msg("Saved configuration to: {}".format(DAP_CONFIG_FILE_PATH))


def PLUGIN_ENTRY():
    global DAP_INSTANCE
    logging.basicConfig(format='[%(levelname)s] %(message)s\t(%(module)s:%(funcName)s)')
    logging.root.setLevel(logging.DEBUG)
    DAP_INSTANCE = DebugAutoPatchPlugin()
    return DAP_INSTANCE


repository: HACK-THE-WORLD/IDAPluginList
file: CTO/cto_func_lister_plugin.py
description: IDA插件集合，含项目名称及项目地址，每日定时Clone项目。
comment: CTO Function Lister
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_func_lister_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_func_lister")
ida_idaapi.require("icon")
ida_idaapi.require("syncdata")
ida_idaapi.require("qtutils")

class cto_func_lister_plugin_t(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_KEEP
    comment = "CTO Function Lister"
    toolbar_displayed_name = cto_base.cto_base.orig_title
    toolbar_name = toolbar_displayed_name + 'Toolbar'
    wanted_name = comment
    wanted_hotkey = "Alt-Shift-F"
    help = "Press '" + wanted_hotkey + "' to display the " + comment + " widget."
    
    action_name = "cto_func_lister: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_func_lister: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_func_lister()
            
        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_func_lister_plugin_t.exec_from_toolbar.action_name,
            cto_func_lister_plugin_t.comment,
            cto_func_lister_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_func_lister_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_base.cto_base.is_dark_mode_with_main():
            ida_kernwin.update_action_icon(cto_func_lister_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_func_lister(self):
        global g_cto_func_lister

        # 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

        # for the first message when this plugin is launched by a user by pressing shortcut key or going to menu.
        if 'g_cto_func_lister' 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))
            
        if self.g:
            # save several important data with pickle and closing hooks...
            self.g.Close(0)
        
        # reload the main modules
        ida_idaapi.require("cto_func_lister")

        # get sync data on a global variable
        sd = syncdata.sync_data()
        sync_data = sd.get()
        # execute the main function
        self.g = cto_func_lister.exec_cto_function_lister(cto_data=sync_data, debug=debug)
        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
        # 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_func_lister_plugin_t.exec_from_toolbar.action_name, act_icon_id)
        ida_kernwin.update_action_icon(self.menu_path + self.wanted_name, act_icon_id)
        
        # install ui hook to enable toolbar later
        self.ph = qtutils.enable_toolbar_t(self.toolbar_name)
        
        # show the messages after launching
        if 'g_cto_func_lister' not in globals():
            ida_kernwin.msg("Launched %s.%s" % (self.wanted_name, os.linesep))
        else:
            ida_kernwin.msg("Reloaded %s.%s" % (self.wanted_name, os.linesep))
            
        # put the instance in the global variables
        g_cto_func_lister = self.g
        if not hasattr(sys.modules["__main__"], "g_cto_func_lister"):
            setattr(sys.modules["__main__"], "g_cto_func_lister", self.g)
        else:
            sys.modules["__main__"].g_cto_func_lister = self.g
            
    def run(self, arg):
        self.exec_cto_func_lister()
        
    def term(self):
        if self.g:
            self.g.Close(0)
        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_func_lister_plugin_t.exec_from_toolbar.action_name)
        ida_kernwin.delete_toolbar(self.wanted_name)
        
        if hasattr(sys.modules["__main__"], "g_cto_func_lister"):
            delattr(sys.modules["__main__"], "g_cto_func_lister")

def PLUGIN_ENTRY():
    return cto_func_lister_plugin_t()

def main():
    global g_cto_func_lister
    g_cto_func_lister = cto_func_lister.exec_cto_function_lister()

if __name__ == '__main__':
    main()


repository: clarkb7/annotate_lineinfo
file: annotate_lineinfo_plugin.py
description: Annotate IDA with source and line number information from a PDB
created_at: 2019-03-24 00:42:03+00:00
pushed_at: 2021-04-15 05:43:49+00:00
url:
https://github.com/clarkb7/annotate_lineinfo/blob/8e6d2872ff099267c7e8894d741d7e58ec3fa38f/annotate_lineinfo_plugin.py

content:
"""
IDA plugin for annotate_lineinfo

Copyright (c) 2019 Branden Clark [github.com/clarkb7]
MIT License, see LICENSE for details.
"""
import idaapi
import annotate_lineinfo.annotate_lineinfo as ali

PLUGIN_COMMENT = "Annotate IDA with source and line number information from a PDB"
PLUGIN_HELP = "github.com/clarkb7/annotate_lineinfo"
PLUGIN_NAME = "annotate_lineinfo"
PLUGIN_WANTED_HOTKEY = 'Alt-Shift-A'

def ALI_MSG(msg,EOL="\n"):
    idaapi.msg("[{}] {}{}".format(PLUGIN_NAME, msg, EOL))

ACTION_ADD_ANNOTATION = 0
ACTION_DEL_ANNOTATION = 1

ali_plugin = None
try:
    idaapi.action_handler_t
except AttributeError:
    ALI_MSG("IDA action API unavailable")
    ALI_IDA_ACTION_API = False
else:
    ALI_IDA_ACTION_API = True
    class ALI_DISASM_SelectionHandler(idaapi.action_handler_t):
        """Dynamic action handler. Annotate selection with line info"""
        def __init__(self, action):
            idaapi.action_handler_t.__init__(self)
            self._action = action

        def activate(self, ctx):
            try:
                # from is a reserved keyword in python...
                cur_sel_from = getattr(ctx.cur_sel, "from")
                start,end = (x.at.toea() for x in [cur_sel_from, ctx.cur_sel.to])
            except AttributeError:
                _,start,end = idaapi.read_selection()
            length = end-start
            if self._action == ACTION_ADD_ANNOTATION:
                ali.ida_add_lineinfo_comment_to_range(ali_plugin.dia, start, length)
            elif self._action == ACTION_DEL_ANNOTATION:
                ali.ida_del_lineinfo_comment_from_range(start, length)

    class ALI_DISASM_FunctionHandler(idaapi.action_handler_t):
        """Dynamic action handler. Annotate function with line info"""
        def __init__(self, action):
            idaapi.action_handler_t.__init__(self)
            self._action = action
        def activate(self, ctx):
            ida_func = ctx.cur_func
            if self._action == ACTION_ADD_ANNOTATION:
                ali.ida_add_lineinfo_comment_to_func(ali_plugin.dia, ida_func)
            elif self._action == ACTION_DEL_ANNOTATION:
                ali.ida_del_lineinfo_comment_from_func(ida_func)

    class ALI_FUNCS_Handler(idaapi.action_handler_t):
        """Action handler. Annotate function with line info"""
        def __init__(self, action):
            idaapi.action_handler_t.__init__(self)
            self._action = action
        def activate(self, ctx):
            for pfn_id in ctx.chooser_selection:
                ida_func = idaapi.getn_func(pfn_id-1)
                if self._action == ACTION_ADD_ANNOTATION:
                    ali.ida_add_lineinfo_comment_to_func(ali_plugin.dia, ida_func)
                elif self._action == ACTION_DEL_ANNOTATION:
                    ali.ida_del_lineinfo_comment_from_func(ida_func)
            return 1
        def update(self, ctx):
            return idaapi.AST_ENABLE_FOR_FORM

    class ALI_MENU_AnnotateHandler(idaapi.action_handler_t):
        """Menu action handler. Annotate entire file with line info"""
        def __init__(self, action):
            idaapi.action_handler_t.__init__(self)
            self._action = action
        def activate(self, ctx):
            if self._action == ACTION_ADD_ANNOTATION:
                ali.ida_annotate_lineinfo_dia(ali_plugin.dia)
            elif self._action == ACTION_DEL_ANNOTATION:
                ali.ida_del_annotations()
            return 1
        def update(self, ctx):
            return idaapi.AST_ENABLE_ALWAYS

    class ALI_MENU_ChoosePDBHandler(idaapi.action_handler_t):
        """Menu action handler. Choose PDB file to load info from"""
        def activate(self, ctx):
            args = [False, "*.pdb", "Enter path to PDB file"]
            try:
                pdbpath = idaapi.ask_file(*args)
            except AttributeError:
                pdbpath = idc.AskFile(*args)
            if pdbpath is None:
                return 0
            if not ali_plugin.init_dia(inbin_path=pdbpath):
                return 0
            if not ali_plugin.attach_actions():
                return 0
            return 1
        def update(self, ctx):
            return idaapi.AST_ENABLE_ALWAYS

    class ALI_MENU_RetryPDBHandler(idaapi.action_handler_t):
        """Menu action handler. Retry auto-find PDB file"""
        def activate(self, ctx):
            if not ali_plugin.init_dia():
                return 0
            if not ali_plugin.attach_actions():
                return 0
            return 1
        def update(self, ctx):
            return idaapi.AST_ENABLE_ALWAYS

    class ALI_Hooks(idaapi.UI_Hooks):
        def finish_populating_tform_popup(self, form, popup):
            tft = idaapi.get_tform_type(form)
            if tft == idaapi.BWN_DISASM: # Disassembly view
                descs = []
                # Choose either selection or function annotation depending on cursor
                selection = idaapi.read_selection()
                if selection[0] == True:
                    descs.append(idaapi.action_desc_t(None,
                        'Annotate selection with line info',
                        ALI_DISASM_SelectionHandler(ACTION_ADD_ANNOTATION)))
                    descs.append(idaapi.action_desc_t(None,
                        'Remove annotations from selection',
                        ALI_DISASM_SelectionHandler(ACTION_DEL_ANNOTATION)))
                else:
                    func = idaapi.get_func(ScreenEA())
                    if func is not None:
                        descs.append(idaapi.action_desc_t(None,
                            'Annotate function with line info',
                            ALI_DISASM_FunctionHandler(ACTION_ADD_ANNOTATION)))
                        descs.append(idaapi.action_desc_t(None,
                            'Remove annotations from function',
                            ALI_DISASM_FunctionHandler(ACTION_DEL_ANNOTATION)))
                # Add corresponding action to popup menu
                for d in descs:
                    idaapi.attach_dynamic_action_to_popup(form, popup, d, None)
            elif tft == idaapi.BWN_FUNCS: # Functions view
                # Add action to popup menu
                idaapi.attach_action_to_popup(form, popup,
                    type(ali_plugin).action_wfuncs_add_name, None,
                    idaapi.SETMENU_INS)
                idaapi.attach_action_to_popup(form, popup,
                    type(ali_plugin).action_wfuncs_del_name, None,
                    idaapi.SETMENU_INS)

class ALI_plugin_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_PROC | idaapi.PLUGIN_HIDE
    comment = PLUGIN_COMMENT
    help = PLUGIN_HELP
    wanted_name = PLUGIN_NAME
    wanted_hotkey = ''

    menu_path = "Edit/Annotate lineinfo/"
    action_wfuncs_add_name = 'ali:wfuncs_add'
    action_wfuncs_add_label = "Annotate function(s) with line info"
    action_wfuncs_del_name = 'ali:wfuncs_del'
    action_wfuncs_del_label = "Remove annotations from function(s)"
    action_menu_annotate_name = 'ali:menu_annotate'
    action_menu_annotate_label = "Annotate entire input file"
    action_menu_delannotate_name = 'ali:menu_delannotate'
    action_menu_delannotate_label = "Remove all annotations"
    action_menu_loadpdb_name = 'ali:menu_loadpdb'
    action_menu_loadpdb_label = "Choose PDB file..."
    action_menu_retrypdb_name = 'ali:menu_retrypdb'
    action_menu_retrypdb_label = "Retry auto-find PDB"

    def init(self):
        self.dia = None
        self.hooks = None

        idaapi.autoWait()

        if not self.init_dia():
            ALI_MSG("Please specify PDB file path using '{}{}'".format(
                    type(self).menu_path,type(self).action_menu_loadpdb_label))

        if ALI_IDA_ACTION_API:
            # Register actions
            actions = [
                idaapi.action_desc_t(
                    type(self).action_wfuncs_add_name, type(self).action_wfuncs_add_label,
                    ALI_FUNCS_Handler(ACTION_ADD_ANNOTATION)),
                idaapi.action_desc_t(
                    type(self).action_wfuncs_del_name, type(self).action_wfuncs_del_label,
                    ALI_FUNCS_Handler(ACTION_DEL_ANNOTATION)),
                idaapi.action_desc_t(
                    type(self).action_menu_annotate_name, type(self).action_menu_annotate_label,
                    ALI_MENU_AnnotateHandler(ACTION_ADD_ANNOTATION), PLUGIN_WANTED_HOTKEY),
                idaapi.action_desc_t(
                    type(self).action_menu_delannotate_name, type(self).action_menu_delannotate_label,
                    ALI_MENU_AnnotateHandler(ACTION_DEL_ANNOTATION)),
                idaapi.action_desc_t(
                    type(self).action_menu_loadpdb_name, type(self).action_menu_loadpdb_label,
                    ALI_MENU_ChoosePDBHandler()),
                idaapi.action_desc_t(
                    type(self).action_menu_retrypdb_name, type(self).action_menu_retrypdb_label,
                    ALI_MENU_RetryPDBHandler()),
            ]
            for action in actions:
                if not idaapi.register_action(action):
                    ALI_MSG("Failed to register action: {}".format(action.name))
                    return idaapi.PLUGIN_SKIP

            # Attach actions
            if not self.attach_actions():
                return idaapi.PLUGIN_SKIP

        ALI_MSG("loaded!")
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        pass

    def term(self):
        ALI_MSG("unloading!")
        if ALI_IDA_ACTION_API:
            if self.hooks is not None:
                self.hooks.unhook()
            idaapi.unregister_action(type(self).action_wfuncs_add_name)
            idaapi.unregister_action(type(self).action_wfuncs_del_name)
            idaapi.unregister_action(type(self).action_menu_annotate_name)
            idaapi.unregister_action(type(self).action_menu_loadpdb_name)
            idaapi.unregister_action(type(self).action_menu_retrypdb_name)

    def init_dia(self, inbin_path=None, sympaths=None):
        if inbin_path is None:
            inbin_path = idaapi.get_input_file_path()
            if inbin_path is None:
                ALI_MSG("No file loaded")
                return False

        if sympaths is None:
            sympaths = []
            ida_sympath = ali.ida_get_sympath()
            if ida_sympath is not None:
                sympaths.append(ida_sympath)

        try:
            self.dia = ali.DIASession(inbin_path, sympaths=sympaths)
        except ValueError as e:
            ALI_MSG("Error loading PDB: {}".format(e))
            return False
        ALI_MSG("Loaded PDB info!")
        return True

    def attach_actions(self):
        if ALI_IDA_ACTION_API:
            # Menu actions
            menu_actions = [
                type(self).action_menu_loadpdb_name,
                type(self).action_menu_retrypdb_name,
            ]
            if self.ready():
                menu_actions += [
                    type(self).action_menu_annotate_name,
                    type(self).action_menu_delannotate_name,
                ]
            for name in menu_actions:
                if not idaapi.attach_action_to_menu(
                    type(self).menu_path,
                    name, idaapi.SETMENU_APP):
                    ALI_MSG("Failed to attach action: {}".format(name))
                    return False

            # UI Hooks
            if self.ready() and self.hooks is None:
                self.hooks = ALI_Hooks()
                if not self.hooks.hook():
                    ALI_MSG("Failed to install UI hooks")
                    return False
        return True

    def ready(self):
        return self.dia is not None

def PLUGIN_ENTRY():
    global ali_plugin
    ali_plugin = ALI_plugin_t()
    return ali_plugin


repository: ianmkim/snnpp
file: amx2.py
description: Spiking Neural Network implementation in pure C++ with minimal dependencies
url: https://github.com/ianmkim/snnpp/blob/3203272f7f58744c7852d82740825f492552317d/amx2.py
wanted_name: Aarch64 AMX
comment: Aarch64 Apple AMX extension
created_at: 2022-10-31 18:49:13+00:00
pushed_at: 2022-11-15 02:26:04+00:00
content:
import idaapi
import ida_hexrays


AMX_NONE = 0
AMX_OP0 = 1
AMX_OP1 = 2
AMX_OP2 = 3
AMX_OP3 = 4
AMX_OP4 = 5
AMX_OP5 = 6
AMX_OP6 = 7
AMX_OP7 = 8
AMX_OP8 = 9
AMX_OP9 = 10
AMX_OP10 = 11
AMX_OP11 = 12
AMX_OP12 = 13
AMX_OP13 = 14
AMX_OP14 = 15
AMX_OP15 = 16
AMX_OP16 = 17
AMX_OP17 = 18
AMX_OP18 = 19
AMX_OP19 = 20
AMX_OP20 = 21
AMX_OP21 = 22
AMX_OP22 = 23

OP_NAMES = {
	AMX_OP0: "AMXLDX",
	AMX_OP1: "AMXLDY",
	AMX_OP2: "AMXSTX",
	AMX_OP3: "AMXSTY",
	AMX_OP4: "AMXLDZ",
	AMX_OP5: "AMXSTZ",
	AMX_OP6: "AMXLDZI",
	AMX_OP7: "AMXSTZI",
	AMX_OP8: "AMXEXTRX", # amxextrx?
	AMX_OP9: "AMXEXTRY", # amxextry?
	AMX_OP10: "AMXFMA64",
	AMX_OP11: "AMXFMS64",
	AMX_OP12: "AMXFMA32",
	AMX_OP13: "AMXFMS32",
	AMX_OP14: "AMXMAC16",
	AMX_OP15: "AMXFMA16",
	AMX_OP16: "AMXFMS16",
	AMX_OP17: "AMX17", # amxset / amxclr
	AMX_OP18: "AMXVECINT",
	AMX_OP19: "AMXVECFP",
	AMX_OP20: "AMXMATINT",
	AMX_OP21: "AMXMATFP",
	AMX_OP22: "AMXGENLUT",
}

OP_INTRINSIC_NAMES = {
	AMX_OP0:  "__amx_ldx",
	AMX_OP1:  "__amx_ldy",
	AMX_OP2:  "__amx_stx",
	AMX_OP3:  "__amx_sty",
	AMX_OP4:  "__amx_ldz",
	AMX_OP5:  "__amx_stz",
	AMX_OP6:  "__amx_ldzi",
	AMX_OP7:  "__amx_stzi",
	AMX_OP8:  "__amx_extrx",
	AMX_OP9:  "__amx_extry",
	AMX_OP10: "__amx_fma64",
	AMX_OP11: "__amx_fms64",
	AMX_OP12: "__amx_fma32",
	AMX_OP13: "__amx_fms32",
	AMX_OP14: "__amx_mac16",
	AMX_OP15: "__amx_fma16",
	AMX_OP16: "__amx_fms16",
	AMX_OP17: "__amx_op17", # amxset / amxclr
	AMX_OP18: "__amx_vecint",
	AMX_OP19: "__amx_vecfp",
	AMX_OP20: "__amx_matint",
	AMX_OP21: "__amx_matfp",
	AMX_OP22: "__amx_genlut",
}

def decode_AMX(d, insn):
	if (d & 0xfffffC00) == 0x00201000:
		Xr = d & 31
		m = (d >> 5) & 31
		if m <= AMX_OP22 - AMX_OP0:
			#insn.itype = idaapi.ARM_nop
			insn.itype = idaapi.ARM_hlt
			insn.segpref = 14
			if m == 17:
				insn.Op1.type = idaapi.o_imm
				insn.Op1.value = Xr
				insn.Op1.dtype = idaapi.dt_byte
			else:
				insn.Op1.type = idaapi.o_reg
				insn.Op1.reg = Xr + 129
				insn.Op1.dtype = idaapi.dt_qword
			insn.insnpref = AMX_OP0 + m
			insn.size = 4
		return True
	return False

class Aarch64AMXHook(idaapi.IDP_Hooks):
	CUSTOM_INSTRUCTIONS = {idaapi.ARM_hlt}
	INDENT = 16
	def ev_ana_insn(self, outctx):
		return outctx.size if decode_AMX(idaapi.get_dword(outctx.ea), outctx) else 0

	def ev_emu_insn(self, insn):
		if insn.itype != idaapi.ARM_brk:
			return False
		return True

	def ev_out_mnem(self, outctx):
		if outctx.insn.itype in self.CUSTOM_INSTRUCTIONS:
			mnem = OP_NAMES.get(ord(outctx.insn.insnpref), None)
			if mnem is not None:
				outctx.out_custom_mnem(mnem, self.INDENT)
				return 1
		return 0

class MicroInstruction(ida_hexrays.minsn_t):

	def __init__(self, opcode, ea):
		ida_hexrays.minsn_t.__init__(self, ea)
		self.opcode = opcode
		self.l.zero()
		self.r.zero()
		self.d.zero()

class CallBuilder():

	def __init__(self, cdg, name, return_type=idaapi.tinfo_t(idaapi.BT_VOID)):
		self.emitted = False
		self.cdg = cdg
		self.callinfo = ida_hexrays.mcallinfo_t()
		self.callinfo.callee = idaapi.BADADDR
		self.callinfo.solid_args = 0
		self.callinfo.call_spd = 0
		self.callinfo.stkargs_top = 0
		self.callinfo.cc = idaapi.CM_CC_FASTCALL
		self.callinfo.return_type = return_type
		self.callinfo.flags = idaapi.FCI_SPLOK | idaapi.FCI_FINAL | idaapi.FCI_PROP
		self.callinfo.role = idaapi.ROLE_UNK

		glbhigh_off = cdg.mba.get_stack_region().off + cdg.mba.get_stack_region().size
		# what memory is visible to the call : GLBLOW - GLBHIGH
		self.callinfo.visible_memory.add(ida_hexrays.ivl_t(0x00, 0x100000))
		self.callinfo.visible_memory.add(ida_hexrays.ivl_t(glbhigh_off, 0xFFFFFFFFFFFFFFFF - glbhigh_off))
		# spoiled locations : GLBLOW - GLBHIGH
		self.callinfo.spoiled.mem.add(ida_hexrays.ivl_t(0x00, 0x100000))
		self.callinfo.spoiled.mem.add(ida_hexrays.ivl_t(glbhigh_off, 0xFFFFFFFFFFFFFFFF - glbhigh_off))

		self.callins = MicroInstruction(ida_hexrays.m_call, self.cdg.insn.ea)
		self.callins.l.make_helper(name)
		self.callins.d.t = ida_hexrays.mop_f
		self.callins.d.size = 0
		self.callins.d.f = self.callinfo

		if (return_type.is_void()):
			self.ins = self.callins
		else:
			self.callins.d.size = return_type.get_size()
			self.ins = MicroInstruction(ida_hexrays.m_mov, self.cdg.insn.ea)
			self.ins.l.t = ida_hexrays.mop_d
			self.ins.l.d = self.callins
			self.ins.l.size = self.callins.d.size
			self.ins.d.t = ida_hexrays.mop_r
			self.ins.d.r = 0x00
			self.ins.d.size = self.callins.d.size

	def add_register_argument(self, t, operand):
		ca = ida_hexrays.mcallarg_t()
		ca.t = idaapi.mop_r
		ca.r = operand
		ca.type = t
		ca.size = t.get_size()
		self.callinfo.args.push_back(ca)
		self.callinfo.solid_args += 1

	def set_return_register(self, reg):
		self.ins.d.r = reg

	def emit(self):
		if self.emitted == False:
			self.cdg.mb.insert_into_block(self.ins, self.cdg.mb.tail)
			self.emitted = True

class AMXFilter(ida_hexrays.microcode_filter_t):
	def __init__(self):
		ida_hexrays.microcode_filter_t.__init__(self)
		ida_hexrays.install_microcode_filter(self, True)

	def match(self, cdg):
		return cdg.insn.itype == idaapi.ARM_hlt and cdg.insn.insnpref != AMX_NONE

	def apply(self, cdg):
		op = ord(cdg.insn.insnpref)
		intrinsic_name = OP_INTRINSIC_NAMES.get(op, '__amx%d' % op)
		if cdg.insn.Op1.type == idaapi.o_reg:
			builder = CallBuilder(cdg, intrinsic_name)
			builder.add_register_argument(idaapi.tinfo_t(idaapi.BT_INT64 | idaapi.BTMT_UNSIGNED), cdg.load_operand(0))
			builder.emit()
		elif cdg.insn.Op1.type == idaapi.o_imm:
			if op == AMX_OP17 and cdg.insn.Op1.value == 0:
				builder = CallBuilder(cdg, '__amx_begin')
				builder.emit()
			elif op == AMX_OP17 and cdg.insn.Op1.value == 1:
				builder = CallBuilder(cdg, '__amx_end')
				builder.emit()
			else:
				builder = CallBuilder(cdg, '%s_%d' % (intrinsic_name, cdg.insn.Op1.value))
				builder.emit()

		return idaapi.MERR_OK


class Aarch64AMXPlugin(idaapi.plugin_t):
	flags = idaapi.PLUGIN_PROC | idaapi.PLUGIN_HIDE
	comment = "Aarch64 Apple AMX extension"
	wanted_hotkey = ""
	help = "Runs transparently"
	wanted_name = "Aarch64 AMX"
	hook = None
	enabled = 1

	def init(self):
		if idaapi.ph_get_id() != idaapi.PLFM_ARM or idaapi.BADADDR <= 0xFFFFFFFF:
			return idaapi.PLUGIN_SKIP
		if not ida_hexrays.init_hexrays_plugin():
			print("[-] {0} : no decompiler available, skipping".format(self.wanted_name))
			return idaapi.PLUGIN_SKIP
		print("%s init"%self.comment)
		self.hook = Aarch64AMXHook()
		self.hook.hook()
		self.filter = AMXFilter()
		return idaapi.PLUGIN_KEEP

	def run():
		pass

	def term(self):
		if self.hook is not None:
			self.hook.unhook()
		print("%s unloaded"%self.comment)

def PLUGIN_ENTRY():
	return Aarch64AMXPlugin()


repository: rouault/mapproxy_hips
file: mapproxy_hips/pluginmodule.py
description: HIPS plugin for mapproxy
created_at: 2022-05-16 14:42:47+00:00
pushed_at: 2022-07-13 09:22:20+00:00
url:
https://github.com/rouault/mapproxy_hips/blob/8986eca85f27969a5afe6fbc4d75c001168194e5/mapproxy_hips/pluginmodule.py

content:
# SPDX-License-Identifier: MIT
# Copyright (C) 2021-2022 Spatialys
# Funded by Centre National d'Etudes Spatiales (CNES): https://cnes.fr

from mapproxy.config.loader import register_source_configuration, register_service_configuration
from mapproxy.config.spec import add_subcategory_to_layer_md
from mapproxy.util.ext.dictspec.spec import anything
from mapproxy.script.util import register_command
from mapproxy.service.demo import register_extra_demo_server_handler, register_extra_demo_substitution_handler

from mapproxy_hips.config.hips_source_configuration import HIPSSourceConfiguration, hips_source_yaml_spec
from mapproxy_hips.config.hips_service_configuration import hips_service_creator, hips_service_yaml_spec
from mapproxy_hips.script.hipsallsky import hipsallsky_command
from mapproxy_hips.script.hipsseed import hipsseed_command
from mapproxy_hips.service.demo_extra import extra_demo_server_handler, extra_demo_substitution_handler

import logging
log = logging.getLogger('mapproxy.hips')

already_executed = False

def plugin_entrypoint():
    """ Entry point of the plugin, called by mapproxy """

    global already_executed
    if already_executed:
        return
    already_executed = True

    log.info('execute plugin_entrypoint')

    register_source_configuration('hips', HIPSSourceConfiguration,
                                  'hips', hips_source_yaml_spec())

    register_service_configuration('hips', hips_service_creator,
                                   'hips', hips_service_yaml_spec())
    # Add a 'hips' subcategory to layer spec to be able to define hips service
    # specific layer metadata
    add_subcategory_to_layer_md('hips', anything())

    # Register command line utilities as sub-commands of mapproxy-util
    register_command('hips-allsky', {
        'func': hipsallsky_command,
        'help': 'Generate HIPS allsky preview file.'
    })
    register_command('hips-seed', {
        'func': hipsseed_command,
        'help': 'Pre-generate HIPS tiles.'
    })

    register_extra_demo_server_handler(extra_demo_server_handler)
    register_extra_demo_substitution_handler(extra_demo_substitution_handler)


repository: BinaryAnalysisPlatform/bap-ida-python
file: plugins/bap/plugins/pseudocode_bap_taint.py
description: integration with IDA
wanted_name: BAP Taint Pseudocode
comment: BAP Taint Plugin for Pseudocode View
created_at: 2016-06-10 18:00:33+00:00
pushed_at: 2021-05-27 15:13:33+00:00
url:
https://github.com/BinaryAnalysisPlatform/bap-ida-python/blob/d8d4679de2f50bb75f556419565821d95404034e/plugins/bap/plugins/pseudocode_bap_taint.py

content:
"""
Hex-Rays Plugin to propagate taint information to Pseudocode View.

Requires BAP_Taint plugin, and installs callbacks into it.
"""

import idc
import idaapi

from bap.utils import hexrays
from bap.utils.bap_taint import BapTaint

colors = {
    'black':   0x000000,
    'red':     0xCCCCFF,
    'green':   0x99FF99,
    'yellow':  0xC2FFFF,
    'blue':    0xFFB2B2,
    'magenta': 0xFFB2FF,
    'cyan':    0xFFFFB2,
    'white':   0xFFFFFF,
    'gray':    0xEAEAEA,
}


def next_color(current_color, ea):
    coloring_order = [
        colors[c] for c in [
            'gray',
            'white',
            'red',
            'yellow',
        ]
    ]
    BGR_MASK = 0xffffff
    ea_color = idaapi.get_item_color(ea)
    if ea_color & BGR_MASK not in coloring_order:
        return current_color
    assert(current_color & BGR_MASK in coloring_order)
    ea_idx = coloring_order.index(ea_color & BGR_MASK)
    current_idx = coloring_order.index(current_color & BGR_MASK)
    if ea_idx >= current_idx:
        return ea_color
    else:
        return current_color


class PseudocodeBapTaint(hexrays.PseudocodeVisitor):
    """Propagate taint information from Text/Graph view to Pseudocode view."""

    flags = idaapi.PLUGIN_HIDE
    comment = "BAP Taint Plugin for Pseudocode View"
    help = "BAP Taint Plugin for Pseudocode View"
    wanted_name = "BAP Taint Pseudocode"

    def visit_line(self, line):
        line.widget.bgcolor = colors['gray']
        for addr in line.extract_addresses():
            line.widget.bgcolor = next_color(line.widget.bgcolor, addr)


def PLUGIN_ENTRY():
    """Install Pseudocode_BAP_Taint upon entry."""
    return PseudocodeBapTaint()


repository: tmr232/ida-plugins
file: function_flow/function_flow.py
description: IDA Pro Plugins
wanted_name: Function Flow
comment: Show Flow in Functions
created_at: 2015-05-27 12:35:29+00:00
pushed_at: 2015-05-31 13:05:09+00:00
url:
https://github.com/tmr232/ida-plugins/blob/c333ebecc9b3d641cdfdc9ab0d66f73e0b551fc3/function_flow/function_flow.py

content:
import networkx as nx

import idaapi

from sark import get_codeblock, get_flowchart, get_block_start, get_nx_graph

try:
    from sark.ui import ActionHandler

    use_new_ui = True
except:
    use_new_ui = False

COLOR_REACHABLE = 0x66EE11
COLOR_UNREACHABLE = 0x6611EE
COLOR_REACHING = 0x11EE66
COLOR_NOT_REACHING = 0x1166EE
COLOR_SOURCE = 0xEE6611
COLOR_NONE = 0xFFFFFFFF
COLOR_EXIT = 0x000048


def iter_exit_nodes(ea):
    for block in get_flowchart(ea):
        # Check if there are successors
        for successor in block.next:
            break
        else:
            yield block


def clear_func(ea):
    for block in get_flowchart(ea):
        block.color = COLOR_NONE


def mark_not_reaching_nodes(ea, source_color=COLOR_SOURCE, other_color=COLOR_NOT_REACHING):
    graph = get_nx_graph(ea)
    graph = graph.reverse()
    block_ea = get_block_start(ea)
    reaching = nx.descendants(graph, block_ea)
    for node_ea in graph.nodes_iter():
        if node_ea not in reaching:
            get_codeblock(node_ea).color = other_color

    get_codeblock(ea).color = source_color


def mark_reaching_nodes(ea, source_color=COLOR_SOURCE, other_color=COLOR_REACHING):
    graph = get_nx_graph(ea)
    graph = graph.reverse()
    block_ea = get_block_start(ea)
    for descendant in nx.descendants(graph, block_ea):
        get_codeblock(descendant).color = other_color

    get_codeblock(ea).color = source_color


def mark_unreachable_nodes(ea, source_color=COLOR_SOURCE, other_color=COLOR_UNREACHABLE):
    graph = get_nx_graph(ea)
    block_ea = get_block_start(ea)
    descendants = nx.descendants(graph, block_ea)
    for block in get_flowchart(ea):
        if block.startEA not in descendants:
            block.color = other_color

    get_codeblock(ea).color = source_color


def mark_reachable_nodes(ea, source_color=COLOR_SOURCE, other_color=COLOR_REACHABLE):
    graph = get_nx_graph(ea)
    block_ea = get_block_start(ea)
    for descendant in nx.descendants(graph, block_ea):
        get_codeblock(descendant).color = other_color

    get_codeblock(ea).color = source_color


def mark_exit_nodes(ea, node_color=COLOR_EXIT):
    for block in iter_exit_nodes(ea):
        block.color = node_color


if use_new_ui:
    #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    # Start of IDA >= 6.7 Code

    class MarkReachableNodesHandler(ActionHandler):
        TEXT = "Reachable"

        def _activate(self, ctx):
            clear_func(ctx.cur_ea)
            mark_reachable_nodes(ctx.cur_ea)


    class MarkUnReachableNodesHandler(ActionHandler):
        TEXT = "Unreachable"

        def _activate(self, ctx):
            clear_func(ctx.cur_ea)
            mark_unreachable_nodes(ctx.cur_ea)


    class MarkReachingNodesHandler(ActionHandler):
        TEXT = "Reaching"

        def _activate(self, ctx):
            clear_func(ctx.cur_ea)
            mark_reaching_nodes(ctx.cur_ea)


    class MarkNotReachingNodesHandler(ActionHandler):
        TEXT = "Not Reaching"

        def _activate(self, ctx):
            clear_func(ctx.cur_ea)
            mark_not_reaching_nodes(ctx.cur_ea)


    class MarkClearHandler(ActionHandler):
        TEXT = "Clear"

        def _activate(self, ctx):
            clear_func(ctx.cur_ea)


    class MarkExits(ActionHandler):
        TEXT = "Exits"

        def _activate(self, ctx):
            clear_func(ctx.cur_ea)
            mark_exit_nodes(ctx.cur_ea)

            idaapi.msg("\n" * 2)

            for block in iter_exit_nodes(ctx.cur_ea):
                idaapi.msg("Exit at 0x{:08X}\n".format(block.startEA))


    class Hooks(idaapi.UI_Hooks):
        def populating_tform_popup(self, form, popup):
            # You can attach here.
            pass

        def finish_populating_tform_popup(self, form, popup):
            # Or here, after the popup is done being populated by its owner.

            if idaapi.get_tform_type(form) == idaapi.BWN_DISASM:
                idaapi.attach_action_to_popup(form, popup, MarkReachableNodesHandler.get_name(), "Mark/")
                idaapi.attach_action_to_popup(form, popup, MarkUnReachableNodesHandler.get_name(), "Mark/")
                idaapi.attach_action_to_popup(form, popup, MarkReachingNodesHandler.get_name(), "Mark/")
                idaapi.attach_action_to_popup(form, popup, MarkNotReachingNodesHandler.get_name(), "Mark/")
                idaapi.attach_action_to_popup(form, popup, MarkExits.get_name(), "Mark/")
                idaapi.attach_action_to_popup(form, popup, MarkClearHandler.get_name(), "Mark/")


    class FunctionFlow(idaapi.plugin_t):
        flags = idaapi.PLUGIN_PROC
        comment = "Show Flow in Functions"
        help = "Show code flow inside functions"
        wanted_name = "Function Flow"
        wanted_hotkey = ""

        def init(self):
            MarkReachableNodesHandler.register()
            MarkUnReachableNodesHandler.register()
            MarkReachingNodesHandler.register()
            MarkNotReachingNodesHandler.register()
            MarkExits.register()
            MarkClearHandler.register()

            self.hooks = Hooks()
            self.hooks.hook()
            return idaapi.PLUGIN_KEEP

        def term(self):
            MarkReachableNodesHandler.unregister()
            MarkUnReachableNodesHandler.unregister()
            MarkReachingNodesHandler.unregister()
            MarkNotReachingNodesHandler.unregister()
            MarkExits.unregister()
            MarkClearHandler.unregister()

        def run(self, arg):
            pass

        # End of IDA >= 6.7 Code
        #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

else:  # Old (< 6.7) ui code
    #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    # Start of IDA < 6.7 Code

    def mark_reachable():
        ea = idaapi.get_screen_ea()
        clear_func(ea)
        mark_reachable_nodes(ea)


    def mark_unreachable():
        ea = idaapi.get_screen_ea()
        clear_func(ea)
        mark_unreachable_nodes(ea)


    def mark_reaching():
        ea = idaapi.get_screen_ea()
        clear_func(ea)
        mark_reaching_nodes(ea)


    def mark_not_reaching():
        ea = idaapi.get_screen_ea()
        clear_func(ea)
        mark_not_reaching_nodes(ea)


    def mark_exists():
        ea = idaapi.get_screen_ea()
        clear_func(ea)
        mark_exit_nodes(ea)

        idaapi.msg("\n" * 2)

        for block in iter_exit_nodes(ea):
            idaapi.msg("Exit at 0x{:08X}\n".format(block.startEA))


    def mark_clear():
        ea = idaapi.get_screen_ea()
        clear_func(ea)


    class FunctionFlow(idaapi.plugin_t):
        flags = idaapi.PLUGIN_PROC
        comment = "Show Flow in Functions"
        help = "Show code flow inside functions"
        wanted_name = "Function Flow"
        wanted_hotkey = ""

        def init(self):
            idaapi.add_menu_item("View/Mark/", "Reachable", None, 0, mark_reachable, tuple())
            idaapi.add_menu_item("View/Mark/", "Un-Reachable", None, 0, mark_unreachable, tuple())
            idaapi.add_menu_item("View/Mark/", "Reaching", None, 0, mark_reaching, tuple())
            idaapi.add_menu_item("View/Mark/", "Not Reaching", None, 0, mark_not_reaching, tuple())
            idaapi.add_menu_item("View/Mark/", "Exists", None, 0, mark_exists, tuple())
            idaapi.add_menu_item("View/Mark/", "Clear", None, 0, mark_clear, tuple())

            return idaapi.PLUGIN_KEEP

        def term(self):
            pass

        def run(self, arg):
            pass

        # End of IDA < 6.7 Code
        #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~


def PLUGIN_ENTRY():
    return FunctionFlow()


repository: obaby/CommentView-4-IDAPRO-7.0
file: showcmts.py
description: CommentView 4 IDAPRO 7.0
wanted_name: Comments Viewer by obaby
comment: IDA Comments Viewer: generate all comments of the idb
created_at: 2018-01-24 08:47:35+00:00
pushed_at: 2020-02-27 02:23:01+00:00
url:
https://github.com/obaby/CommentView-4-IDAPRO-7.0/blob/dc8e76547988d5066cfdbe8b916a8ef5f6542ee7/showcmts.py

content:
'''
 '----------------------------------------------------------------------------------------------------------'
 'IDA Comments Viewer for IDA pro 7.0'
 'Version:1.0 alpha'
 'plugin by obaby'
 'http://www.h4ck.org.cn http://findu.co'
 '----------------------------------------------------------------------------------------------------------'
'''


import idaapi
from ida_ida import *
import idautils
import idc
from idaapi import Choose2
import ida_kernwin


class chooser_handler_t(idaapi.action_handler_t):
    def __init__(self, thing):
        idaapi.action_handler_t.__init__(self)
        self.thing = thing

    def activate(self, ctx):
        #sel = []
        #for i in xrange(len(ctx.chooser_selection)):
        #    sel.append(str(ctx.chooser_selection.at(i)))
        #print "command %s selected @ %s" % (self.thing, ", ".join(sel))
        pass

    def update(self, ctx):
        return idaapi.AST_ENABLE_FOR_FORM if idaapi.is_chooser_tform(ctx.form_type) else idaapi.AST_DISABLE_FOR_FORM


class MyChoose2(Choose2):

    def __init__(self, title, nb = 5, flags=0, width=None, height=None, embedded=False, modal=False):
        Choose2.__init__(
            self,
            title,
            [ ["Address", 16], ["T", 2], ["Instruction/Data", 60], ["Comment", 100]],
            flags = flags,
            width = width,
            height = height,
            embedded = embedded)
        self.n = 0
        self.items = self.get_all_comments() #[ self.make_item() for x in xrange(0, nb+1) ]
        self.icon = 5
        self.selcount = 0
        self.modal = modal
        self.popup_names = [] #["Inzert", "Del leet", "Ehdeet", "Ree frech"]

        #print("created %s" % str(self))

    def OnClose(self):
        print "closed", str(self)

    def OnEditLine(self, n):
        self.items[n][1] = self.items[n][1] + "*"
        #print("editing %d" % n)

    def OnInsertLine(self):
        self.items.append(self.make_item())
        #print("insert line")

    def OnSelectLine(self, n):
        self.selcount += 1

        ea = int(self.items[n][0], 16)
        idc.jumpto(ea)
        #Warning("[%02d] selectline '%s'" % (self.selcount, n))

    def OnGetLine(self, n):
        #print("getline %d" % n)
        return self.items[n]

    def OnGetSize(self):
        n = len(self.items)
        #print("getsize -> %d" % n)
        return n

    def OnDeleteLine(self, n):
        #print("del %d " % n)
        del self.items[n]
        return n

    def OnRefresh(self, n):
        #print("refresh %d" % n)
        return n

    def OnGetIcon(self, n):
        r = self.items[n]
        t = self.icon + r[1].count("*")
        #print "geticon", n, t
        return t

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

    def make_item(self):
        r = [str(self.n), "func_%04d" % self.n]
        self.n += 1
        return r

    def check_isin_filter(self, cmt):
        cmt_str = str(cmt)
        if cmt_str.startswith('void') or cmt_str.startswith('char') \
                or cmt_str.startswith('int') or cmt_str.startswith('switch') \
                or cmt_str.startswith('jump') or cmt_str.startswith('size_t') \
                or cmt_str.startswith('dw') or cmt_str.startswith('nSize') \
                or cmt_str.startswith('hFile') or cmt_str.startswith('lp'):
            return True
        else:
            return False

    def get_all_comments(self):
        cmts = []
        for seg in idautils.Segments():
            print 'Anylising:', idc.SegName(seg), hex(idc.SegStart(seg)), hex(idc.SegEnd(seg)) + '\n'
            ea = idc.SegStart(seg)
            start = idc.SegStart(seg)
            end = idc.SegEnd(seg)
            while ea < end:

                if ea != idc.BADADDR:
                    cmt = idc.GetCommentEx(ea, True)
                    if cmt:
                        if self.check_isin_filter(cmt):
                            print " Address: ",format(ea, '#16X'),'In fliter,IGN:', cmt
                        else:
                            current_cmt = [format(ea, '#16X'), 'R', idc.GetDisasm(ea), cmt]
                            cmts.append(current_cmt)
                            self.n += 1
                            print " Address: ", format(ea, '#16X'), 'R', 'Comment:', cmt

                    cmt2 = idc.GetCommentEx(ea, False)
                    if cmt2:
                        if self.check_isin_filter(cmt2):
                            print " Address: ",format(ea, '#16X'),'In fliter,IGN:', cmt2
                        else:
                            current_cmt = [format(ea, '#16X'), 'N', idc.GetDisasm(ea), cmt2]
                            cmts.append(current_cmt)
                            self.n += 1
                            print " Address: ",format(ea, '#16X'), 'N', 'Comment:', cmt2
                ea = idc.next_head(ea, end)
        return cmts


    def OnGetLineAttr(self, n):
        pass
        #print("getlineattr %d" % n)
        #if n == 1:
        #    return [0xFF0000, 0]


 # -----------------------------------------------------------------------
def test_choose2(modal=False):
    global c
    c = MyChoose2("Comments List", nb=10, modal=modal)
    r = c.show()
    #c.get_all_comments() # get all comments
    form = idaapi.get_current_tform()
    for thing in ["A", "B"]:
        idaapi.attach_action_to_popup(form, None, "choose2:act%s" % thing)


# -----------------------------------------------------------------------
def test_choose2_embedded():
    global c
    c = MyChoose2("Comments List", nb=12, embedded=True, width=123, height=222)
    r = c.Embedded()
    if r == 1:
        try:
            if test_embedded:
                o, sel = _idaapi.choose2_get_embedded(c)
                print("o=%s, type(o)=%s" % (str(o), type(o)))
                test_embedded(o)
        finally:
            c.Close()


class show_cmts_plugin_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = "IDA Comments Viewer: generate all comments of the idb"
    help = "Bugs and report: http://www.h4ck.org.cn"
    wanted_name = "Comments Viewer by obaby"
    wanted_hotkey = "Ctr-Alt-F8"

    def init(self):
        print '----------------------------------------------------------------------------------------------------------'
        print 'IDA Comments Viewer for IDA pro 7.0'
        print 'Version:1.0 alpha'
        print 'plugin by obaby'
        print 'http://www.h4ck.org.cn http://findu.co'
        print '----------------------------------------------------------------------------------------------------------'
        return idaapi.PLUGIN_OK

    def run(self, arg):
        print "Start to analyzing all comments in idb..."
        ida_kernwin.show_wait_box('Analyzing comments in progress, this will take a while.')
        test_choose2(False)
        ida_kernwin.hide_wait_box('Analyzing comments in progress, this will take a while.')
        print "Finished,have a good time"



    def term(self):
        ida_kernwin.hide_wait_box('Analyzing comments in progress, this will take a while.')





def PLUGIN_ENTRY():
    return show_cmts_plugin_t()


repository: ax330d/hrdev
file: hrdev.py
description: Hex-Rays Decompiler Enhanced View
url: https://github.com/ax330d/hrdev/blob/f2667243add5297acaabc4c65fcb21b6eba3c7e5/hrdev.py
wanted_name: HRDEV
comment: Hex-Rays Decompiler Enhanced View
created_at: 2015-02-22 19:59:05+00:00
pushed_at: 2018-09-21 14:53:08+00:00
content:
#!c:\\python27\python.exe
# -*- coding: utf-8 -*-

'''Just a wrapper module.'''

import idaapi
from hrdev_plugin import Plugin

try:
    from PySide import QtCore, QtGui
except:
    from PyQt5 import QtCore, QtWidgets


__author__ = 'Arthur Gerkis'
__version__ = '0.0.4 (beta)'


class HRDEVPluginEntry(idaapi.plugin_t):
    '''HRDEV Plugin Entry.'''
    flags = idaapi.PLUGIN_UNL
    comment = "Hex-Rays Decompiler Enhanced View"

    help = ""
    wanted_name = "HRDEV"
    wanted_hotkey = "Alt+F5"

    def init(self):
        print "HRDEV plugin is loaded. Use Alt+F5 hotkey to decompile."
        return idaapi.PLUGIN_OK

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

    def term(self):
        return


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

    def activate(self, ctx):
        Plugin().run()
        return 1

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


idaapi.register_action(idaapi.action_desc_t(
    'hrdev:decompile',
    'Generate pseudocode (HRDEV)',
    HRDEVMenuHandler(),
    'Alt+F5',
    'Generate pseudocode (HRDEV)'))


idaapi.attach_action_to_menu(
    'View/Open subviews/Generate pseudocode',
    'hrdev:decompile',
    idaapi.SETMENU_APP)


def PLUGIN_ENTRY():
    return HRDEVPluginEntry()


repository: nihilus/krypton
file: Krypton.py
url: https://github.com/nihilus/krypton/blob/979eea2558247221e904f54aedae8a871e27c793/Krypton.py
created_at: 2014-08-31 20:27:58+00:00
pushed_at: 2014-08-31 20:29:07+00:00
content:
doc = """
#################################################################################
#										#
#  Krypton is an IDA Plugin that assists one in executing a function from IDB	#
#  (IDA database) using IDA's powerful Appcall feature				#
#										#
#  krypton takes xrefs from a given function (say a possible decoder) to  	#
#  find all function calls to it and then parses and finds the parameters used 	#
#  (including prototype, no of arguments, and the arguments themselves) from 	#
#  instructions and uses them to execute the function using appcall, this is 	#
#  most useful in analysing a malware binary with encryption			#
#										#
#  To Install, Copy this file into IDA Plugin folder,				#
#  Std path: "C:\Program Files\IDA\plugins\"					#
#  NOTE: Appcall emulation feature requires IDA >= 5.6				#
#										#
#  Author: Karthik Selvaraj at Symantec Corporation				#
#  Email : neoxfx at gmail dot com						#
#										#
#  Krypton 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/>.	#
#										#
#################################################################################
"""

version     = '1.0'
pluginname  = "Krypton"
pluginVname = pluginname + " " + str(version)
debug       = False
CreateIDC   = True
verbose     = 0

import sys,string,os

def kprint(*s):
    print "[%s]" % pluginname,
    for k in s:
        print k,
    print


try:
    import idautils
    import idc
    from idaapi import *
except:
    kprint ("Error: Use within IDA (>= 5.6)!!")
    print doc
    sys.exit(0)


gsplitChar = "|"
class HelperRoutines:
    """
    This class contains helper routines that the rest of the classes use.
    """
    debug                 = debug
    createIDC             = CreateIDC
    version               = version
    verbose               = verbose
    interConnectInstr     = 'jmp'         # basic block interconnect instruction, in case of spaghetti code
    argPassMech           = "push"
    argPassType           = "__cdecl"
    FILTER                = ''.join([(len(repr(chr(x)))==3) and chr(x) or '.' for x in range(256)])
    logPointer            = None
    logPointerIdc         = None

    def getDataValue(self, addr, size):
        if (size == 1):
            return Byte(addr)
        elif (size == 2):
            return Word(addr)
        elif (size == 4):
            return Dword(addr)
        else:
            self.log(2, "Error! while parsing data: 0x%x --> size: %d" % (addr, size))
            raise

    def log(self, v, cmt):
        if (self.debug) and (v <= self.verbose):
            if self.logPointer == None:
                self.openLogFile("")
            self.logPointer.write(cmt+"\n")

    def logIdc(self, cmt):
        if (self.createIDC):
            #cmt = cmt.translate(self.FILTER)
            if self.logPointerIdc == None:
                self.openIdcLogFile("")
            self.logPointerIdc.write(cmt+"\n")

    def flushlog(self):
        if (self.logPointer != None):
            self.logPointer.flush()
        if (self.logPointerIdc != None):
            self.logPointerIdc.flush()

    def logClose(self):
        if (self.logPointer != None):
            self.logPointer.close()
        if (self.logPointerIdc != None):
            self.logPointerIdc.close()

    def openLogFile(self, fnamePrefix):
        if not self.debug:
            return
        if (self.logPointer != None):
            self.logPointer.close()
        self.logPointer = open(GetIdbPath()+"_%s_%s.log"%(pluginVname, fnamePrefix), "w")
        s = "//%s - LOG file %s" % (pluginVname+fnamePrefix, GetIdbPath())
        self.logPointer.write(s + "\n//" + "=" * len(s) + "\n")

    def openIdcLogFile(self, fnamePrefix):
        if not self.createIDC:
            return
        if (self.logPointerIdc != None):
            self.logPointerIdc.close()
        self.logPointerIdc = open(GetIdbPath()+"_%s_%s.idc"%(pluginVname, fnamePrefix), "w")
        s = "//%s - IDC file %s" % (pluginVname+fnamePrefix, GetIdbPath())
        self.logPointerIdc.write(s + "\n//" + "=" * len(s) + "\n")

    def getXrefs(self, ea):
        a = []
        try:
            a = list(CodeRefsTo(ea, 1))
        except:
            a = []
        return a



class Function(HelperRoutines):

    name               = ""
    ea                 = None
    nargs              = 0
    xrefs              = []
    proto              = ""
    argsSizeArray      = []
    runUntilEA         = {'start':BADADDR, 'end':BADADDR}
    ESCAPE_CODE_BEGIN  = 0xF0
    ArgArrCodes        = {"CONST":0xFE, "BUFFER":0xFF}

    def __init__(self, ea, proto, argSizeArr, argPassMech="push"):

        if ((not isLoaded(ea)) or (proto == None) or (type(proto).__name__ != "str") or (argSizeArr == None)):
            self.log(2, "Error! Bad function parameters! :0x%x" % ea)
            raise
        try:
            self.name   = GetFunctionName(ea)
        except:
            self.name   = "0x%x" % ea

        self.ea     = ea

        if (type(argSizeArr).__name__ == "list"):
            self.nargs  = len(argSizeArr)
            self.argsSizeArray  = argSizeArr
        elif (type(argSizeArr).__name__ == "str"):
            b = eval(argSizeArr)
            self.argsSizeArray = b
            self.nargs = len(self.argsSizeArray)
        else:
            kprint("type:%s" % type(argSizeArr).__name__)
            kprint("arr:%s, %d" % (argSizeArr, len(argSizeArr)))
            kprint("arg array issue!")
            raise

        self.xrefs  = self.getXrefs(ea)

        if len(self.xrefs) == 0:
            self.log(2, "Error! No Xref to function ea :0x%x" % ea)
            raise
        self.proto  = proto
        self.argPassMech    = argPassMech


    def getInstrDataPair(self, i):

        if (i > self.nargs):
            self.log(1, "Error! Index out of Array!! %d" % i)
            raise
        (k,v) = self.argsSizeArray[i]
        if(type(k).__name__ == "str"):
            try:
                k = self.ArgArrCodes[k.upper()]
            except:
                self.log(1, "Error! Unknown code in Argument Array at index: %d" % i)
                kprint("Error! Unknown code \"%s\" in Argument Array \"%s\"" % (k, repr(self.argsSizeArray)))
                raise
        return (k,v)

    def getData(self):
        return self.ea, len(self.xrefs), self.proto, self.argsSizeArray, self.argPassMech

    def getEA(self):
        return self.ea

    def getName(self):
        return self.name

    def updateName(self):
        try:
            oproto = self.proto
            oname  = self.name
            self.name   = GetFunctionName(self.ea)
            if (self.name != oname):
                self.proto = self.proto.replace(oname, self.name)
        except:
            self.name   = oname
            self.proto  = oproto

    def setInitializerExecution(self, initStart, initEnd):
        if ((initStart == BADADDR) or (initEnd == BADADDR)):
            return False
        self.runUntilEA['start'] = initStart
        self.runUntilEA['end'] = initEnd
        kprint("Initializer set for function \"%s\"->EA_RANGE(0x%x, 0x%x)" % (self.name, initStart, initEnd))
        return True

    def __repr__(self):
        global gsplitChar
        return " 0x%x %c %s %c %s %c %s %c (%04d)" % (self.ea, gsplitChar, self.proto, gsplitChar, repr(self.argsSizeArray).split("\""), gsplitChar, self.argPassMech, gsplitChar, len(self.xrefs))



IDBbackedup = False
class DebuggerCtrl(HelperRoutines):
    """
    This class implements debugger control functions, this class makes Appcall calling possible.
    """
    oldDebuggerOptions = None
    modulePath         = ""
    moduleArgs         = ""
    sourceDir          = ""

    def __init__(self):
        self.modulePath = GetInputFilePath()
        self.moduleArgs = ""
        self.sourceDir  = os.path.dirname(GetIdbPath())

    def setSafeGuard(self):
        global IDBbackedup
        self.oldDebuggerOptions = SetDebuggerOptions(DOPT_ENTRY_BPT|DOPT_LIB_MSGS|DOPT_THREAD_BPT|DOPT_START_BPT)
        if not IDBbackedup:
            nam = GetIdbPath()+".bak"
            kprint("Taking IDB Backup to \"%s\"" % nam)
            SaveBase(nam, 0)
            IDBbackedup = True

    def start(self):
        self.setSafeGuard()
        if not os.path.isfile(self.modulePath):
            self.modulePath = idc.AskFile(0,"*.*","Please select the right Path of the module/executable loaded in IDA")
            idc.SetInputFilePath(self.modulePath)
        LoadDebugger("win32", 0) # load win32 as default
        if (StartDebugger(self.modulePath, self.moduleArgs, self.sourceDir) == -1):
            self.log(1, "Error! Failed to launch debugger!")
        GetDebuggerEvent(WFNE_ANY, -1) # handle first event


    def stop(self):
        global pvv
        global decResults
        global decViews

        StopDebugger()

        # workaround for mimic'king persistent view, as StopDebugger will kill all views created
        # hook processExit event and re-create views
        SetDebuggerOptions(self.oldDebuggerOptions)

        evt = 0
        while (GetDebuggerEvent(WFNE_ANY, -1) != PROCESS_EXIT):
            evt += 1

        try:
            pvv = Protoviewer()

            for (k,v) in decViews.iteritems():
                kprint("attempt closing for decView 0x%x" % k)
                v.Close()
                del v

            decViews = {}
            for (ea,x) in decResults.iteritems():
                (fpt, res) = x
                f = protoStr2FunctionInstance(fpt)
                if (f != None):
                    decViews[ea] = decViewer(f)
        except Exception,e:
            kprint("Exception!! @DebuggerCtrl:Stop(), ", e)



class Decryptor(HelperRoutines):
    """
    This Class that executes a given function using Appcall with parameters from all its call references
    """

    tempBuffers           = []
    trackBackLimit        = 10            # instruction trace back limit before the decoder function call
    func                  = None          # instance of class function to hold all attributes
    params                = {}            # parameter hash tree for each reference
    debugCtrl             = None
    decRes                = {}

    def __init__(self, decFunc):
        nprefix = "_%s" % decFunc.getName()
        self.openLogFile(nprefix)
        self.openIdcLogFile(nprefix)
        self.log(1, "\n\nSTART DECRYPTOR\n" + "-"*50)
        self.log(1, "Decryptor function parameters => %s" %  repr(decFunc))
        self.func = decFunc
        self.debugCtrl = DebuggerCtrl()
        decRes = {}

    def _buffer(self, size=256):
        return  Appcall.byref("\x00" * size)

    def cleanup(self):
        self.debugCtrl.stop()
        self.logClose()

    def setResults(self):
        global decResults
        if self.func.getEA() in decResults:
            del decResults[self.func.getEA()]
        if len(self.decRes) > 0:
            decResults[self.func.getEA()] = (repr(self.func), self.decRes,)
            return True
        return False

    def run(self):

        for ref in self.func.xrefs:
            self.log(4, "[-] checking near " + hex(ref))
            self.parseArgs(ref)

        self.debugCtrl.start()
        self.log(1, "[*] debugger launched")

        if (self.func.runUntilEA['end'] != BADADDR) and (self.func.runUntilEA['start'] != BADADDR):
            idc.SetRegValue(self.func.runUntilEA['start'], 'EIP')
            idc.RunTo(self.func.runUntilEA['end'])

        self.decRes = None
        self.decRes = {}

        for ref in self.func.xrefs:
            try:
                arg = self.params[ref]
            except:
                self.log(3, "[-!] Couldn't find arg, skipping call at 0x%x" % ref)
                continue

            try:
                if (self.execute(ref, arg) == -1):
                    break
            except Exception,e:
                kprint("Exception!! @decryptor:run, ", e)
                pass

        self.log(1, "[*] debugger stopped")
        self.log(1, "\n\nEND DECRYPTOR\n" + "-"*50)
        if not self.setResults():
            kprint("No result from function run \"%s\"" % self.func.getName())

        return True

    def execute(self, ref, argv):

        dFunc = Appcall.typedobj(self.func.proto)
        dFunc.ea = self.func.ea
        buffers = []

        try:
            self.log(1, ('[-] Found call @ 0x%x to Decoder - ' % ref) + self.formattedFuncPrint(argv))

            cmd = "dFunc("
            ind = 0
            ibd = 0
            for i in range(self.func.nargs):
                insSize, dataSize = self.func.getInstrDataPair(i)
                if (insSize == self.func.ArgArrCodes["CONST"]):  # const value by ref
                    t = Appcall.byref(dataSize)
                    cmd += "t, "
                elif (insSize == self.func.ArgArrCodes["BUFFER"]):
                    buffers.append(self._buffer(dataSize))
                    cmd += "buffers[%d], " % ibd
                    ibd += 1
                else:
                    cmd += "argv[%d], " % ind
                    ind += 1

            cmd = cmd[:-2] + ")"
            self.log(3, "\tcmd: \"%s\"" % cmd)

            try:
                # This eval is used to convert a python string to a python code statement
                # Appcall is constructed in the string cmd with right parameters before being called
                s = 0
                s = eval(cmd)
            except Exception,e:
                dcode = GetDebuggerEvent(WFNE_ANY, -1)
                if ((dcode == 1) or (dcode == 0x4) or (dcode == 0x200) or (dcode == 0x1000)):
                    pass
                else:
                    dstr = "[!] debugger error code(0x%x), %s" % (dcode, repr(e))
                    self.log(3, dstr)
                    if str(e.__str__()).startswith("Appcall:") or (dcode == -2) or (dcode == -1) or (dcode == 0x40) or (dcode == 0x400):
                        kprint("[!!] Cannot continue! %s" % dstr)
                        kprint("[!!] Stopping script!!")
                        return -1

            self.log(3, "[+] Buffer post decryption @ 0x%x" % ref)

            obbuf = []
            for buff in buffers:
                o = self.parseResult(buff);
                self.log(3, "\tbuffer = \"%s\"" % o)
                if ((type(o).__name__ == 'str') and ((len(o) > 3) or (not repr(o).startswith("\'\\x")))):
                    obbuf.append(o)

            if isLoaded(s):
                try:
                    sdo = GetString(s)
                    if ((sdo != None) and (type(sdo).__name__ == 'str')):
                        self.log(3, "\tRet (eax) buffer = \"%s\"" % sdo)
                        if (not repr(sdo).startswith("\'\\x")):
                            obbuf.append(sdo)
                except:
                    pass

            if isLoaded(Dword(s)):
                try:
                    sdo = GetString(Dword(s))
                    if ((sdo != None) and (type(sdo).__name__ == 'str')):
                        self.log(3, "\tRet Poi(eax) buffer = \"%s\"" % sdo)
                        if (not repr(sdo).startswith("\'\\x")):
                            obbuf.append(sdo)
                except:
                    pass

            rs = max(obbuf, key=len)
            rs = rs.replace("\n", "\\n")
            self.decRes[ref] = rs
            self.logIdc("MakeRptCmt(0x%x, \"%s\");" % (ref, rs.replace("\"", "\\\"")))
            return 0
        except Exception,e:
            kprint("Exception!! @decryptor:execute, ", e)
            pass

    def parseResult(self, buff):
        i = 0
        unicodeStr = ""
        asciiStr = buff.cstr().strip()
        nullFound = False
        # Parse out the unicode string
        while i + 1 < len(buff.value) and not nullFound:
            # Unicode vaules come in as little endian
            aChar = struct.unpack("<H", buff.value[i:i+2])[0]

            if aChar != 0:
                unicodeStr += unichr(aChar)
            else:
                nullFound = True
            i += 2
        # Works on the assumption the buffer will be null-ed out after
        # character data, has worked in practice
        if len(unicodeStr) > len(asciiStr):
            return unicodeStr.encode('ascii', 'ignore')
        else:
            return asciiStr

    def parseArgs(self, ref):

        i = self.trackBackLimit
        argsToParse = self.func.nargs
        argv = []
        ref_p = ref
        while (i and argsToParse):
            insSize, dataSize = self.func.getInstrDataPair(self.func.nargs-argsToParse)
            if (insSize > self.func.ESCAPE_CODE_BEGIN):
                argsToParse -= 1
            else:
                if (dataSize > insSize):
                    self.log(3, "[!!] invalid argument size pair (%d, %d) for arg no: %d" % (insSize, dataSize, self.func.nargs-argsToParse))
                    raise
                ref_p = idc.PrevHead(ref_p, 0)
                if ref_p == BADADDR:
                    break
                disstr = idc.GetDisasm(ref_p)
                if disstr.startswith(self.func.argPassMech) and (idc.ItemSize(ref_p) == insSize):
                    arVal = self.getDataValue(ref_p + insSize - dataSize, dataSize)
                    argv.append(arVal)
                    self.log(4, "Arg found: (0x%x)" % arVal)
                    argsToParse -= 1
                    if not argsToParse:
                        break
                if len(self.getXrefs(ref_p)) == 1:
                    nref = self.getXrefs(ref_p)[0]
                    refins = idc.GetDisasm(nref)
                    if refins.startswith(self.interConnectInstr):
                        ref_p = nref
                i -= 1
        if len(argv) > 0:
            self.log(4, "[-] Adding args for call at 0x%x->%s" % (ref,repr(argv)))
            self.params[ref] = argv

    def formattedFuncPrint(self, argv):

        out = self.func.name + "_" + hex(self.func.ea) + "("
        for arg in argv:
            try:
                out += hex(arg) + ', '
            except:
                out += repr(arg) + ', '
        out = out[:-2]
        out += ")"
        return out


class decViewer(idaapi.simplecustviewer_t):
    decfn  = None

    def __init__(self, fn, reRun = False):
        try:
            if (fn == None):
                raise "No valid decrypt function passed!"
            self.decfn = fn
            if self.Create(reRun):
                self.Show()
            else:
                kprint("unable to create decryptor view!")
                raise
        except Exception,e:
            kprint("Exception!! @decViewer, ", e)

    def Create(self, reRun):
        global pluginVname
        global decResults

        self.menu_write2IDB = None
        self.menu_refresh   = None
        self.menu_write2IDC = None

        if not idaapi.simplecustviewer_t.Create (self, pluginVname + " - " + self.decfn.getName() + " result"):
            kprint("Unable to create view \"%s\"" % pluginVname + " - " + self.decfn.getName() + " result")
            return False

        if reRun or (self.decfn.getEA() not in decResults):
            try:
                if (self.decfn != None):
                    dec = Decryptor(self.decfn)
                    dec.run()
                    dec.cleanup()

            except Exception,e:
                kprint("Exception!! @decViewer:Create in calling decryptor, ", e)

        self.refresh()
        return True

    def refresh(self):
        global decResults
        global gsplitChar
        try:
            self.ClearLines()
            if self.decfn.getEA() in decResults:
                (ea,x) = decResults[self.decfn.getEA()]
                for (k,v) in x.iteritems():
                    cline = idaapi.COLSTR("  0x%08X  %c " % (k, gsplitChar), idaapi.SCOLOR_NUMBER)
                    cline += idaapi.COLSTR("%s" % v, idaapi.SCOLOR_STRING)
                    self.AddLine(cline)
            self.Refresh()
        except Exception,e:
            kprint("Exception!! @decViewer:refresh, ", e)
        return True

    def write2IDB(self):
        global decResults
        try:
            if self.decfn.getEA() in decResults:
                (ea,x) = decResults[self.decfn.getEA()]
                for (k,v) in x.iteritems():
                    MakeRptCmt(k, v)
                kprint("Decrypted strings Written to IDB at respective references!")
        except Exception,e:
            kprint("Exception!! @decViewer:write2IDB, ", e)
        return True


    def OnPopup(self):
        try:
            self.ClearPopupMenu()
            if not self.Count():
                self.AddPopupMenu("-")
                self.menu_refresh   = self.AddPopupMenu("(r) Refresh")
            else:
                self.AddPopupMenu("-")
                self.menu_refresh   = self.AddPopupMenu("(r) Refresh")
                self.menu_write2IDB = self.AddPopupMenu("(w) write these to IDB")

        except Exception,e:
            kprint("Exception!! @decViewer:OnPopup, ", e)
        return True

    def OnPopupMenu(self, menu_id):
        if menu_id == self.menu_refresh:
            self.refresh()

        elif menu_id == self.menu_write2IDB:
            self.write2IDB()

        else:
            return False

        return True

    def GotoLn(self, n):
        global gsplitChar
        if n != 0:
            (ln,b,c) = self.GetLine(n)
            if ln == None:
                return True
            ln = tag_remove(ln)
            l = ln.split(gsplitChar)
            nu = int(l[0],16)
            Jump(nu)
            return True

    def OnDblClick(self, shift):
        n = self.GetLineNo()
        self.GotoLn(n)

    def OnKeydown(self, vkey, shift):
        # escape
        if vkey == 27:
            pass
        # enter
        elif vkey == 13:
            n = self.GetLineNo()
            self.GotoLn(n)

        elif vkey == ord("R"):
            self.refresh()

        elif vkey == ord("W"):
            self.write2IDB()

        else:
            return False

        return True



class ProtoAnalyser(HelperRoutines):
    """
    This class Analyses Function Prototypes and Argument value pairs that can be later used in Appcall
    It analyses Top referenced functions, topNum, default is 10

    """

    allXrefFuncs       = []
    topNum             = 0
    functionlist       = []
    MAX_REFS_TO_CHECK  = 10
    TRACK_BACK_LIMIT   = 10
    MAX_ARGS           = 10
    protoArray         = {}
    argPassMechInsSize = 1
    GCD_THRESOLD       = 70
    no_of_funcs        = 0


    def __init__(self, topN=25):
        self.openLogFile("ProtoAnalyser")
        self.no_of_funcs   = idaapi.get_func_qty()
        if topN > self.no_of_funcs:
            self.topNum = self.no_of_funcs
        else:
            self.topNum = topN
        self.allXrefFuncs  = self.getTopXrefFunctions()

    def cleanup(self):
        self.functionlist  = []
        self.protoArray    = {}
        pass


    def getTopXrefFunctions(self):

        funcs = []
        for fnum in xrange(self.no_of_funcs):
            f_ea = idaapi.getn_func(fnum).startEA
            if f_ea != BADADDR:
                funcs.append((GetFunctionName(f_ea), f_ea, len(self.getXrefs(f_ea)),))
        funcs = sorted(funcs, key=lambda x:x[2], reverse=True)
        return funcs

    def analysePrototype(self, ea, force=False):

        try:
            if not isLoaded(ea):
                self.log(2, "Error! Bad address: 0x%x" % ea)
                raise

            try:
                fname = GetFunctionName(ea)
            except:
                fname = "0x%x" % ea

            xrefs = self.getXrefs(ea)
            if len(xrefs) == 0:
                self.log(3, "Discarding function: %s as there is no xref" % fname)
                return None

            tmpSorterDict  = {}
            mr = self.MAX_REFS_TO_CHECK
            if (len(xrefs) < self.MAX_REFS_TO_CHECK):
                mr = len(xrefs)

            xrefs = xrefs[0:mr]
            self.log(3, "\n%s\nAnalysing prototype for function 0x%x" % ('-'*50,ea))

            if (not force) and (GetFunctionFlags(ea) & (FUNC_HIDDEN|FUNC_LIB)):
                self.log(3, "Discarding function: %s as it is a LIB/HIDDEN function" % fname)
                return None

            idbProto = print_type(ea, False)
            gt = GuessType(ea)
            if (((idbProto != None) and (idbProto.find("__usercall") != -1)) or ((gt != None) and (gt.find("__usercall") != -1))):
                self.argPassMech = "mov"
                self.argPassType = "__usercall"
            else:
                self.argPassType = "__cdecl"
                self.argPassMech = "push"

            for ref in xrefs:
                argvArray      = []
                i = self.TRACK_BACK_LIMIT
                ref_p = ref

                proto = "int %s %s(" % (self.argPassType, fname)

                while (i):
                    g = 0
                    ref_p = idc.PrevHead(ref_p, 0)
                    if ((ref_p == BADADDR) or idc.GetDisasm(ref_p).startswith("call")):
                        g = 2
                        break
                    disstr = idc.GetDisasm(ref_p)
                    if disstr.startswith(self.argPassMech):
                        itmSize = idc.ItemSize(ref_p)
                        if (itmSize == 5):
                            dataVal = self.getDataValue(ref_p + self.argPassMechInsSize, itmSize - self.argPassMechInsSize)
                            if isLoaded(dataVal) and (not isCode(GetFlags(dataVal))):
                                proto += "void*,"
                            else:
                                proto += "int,"
                        elif (itmSize == 2):
                            proto += "int,"
                        else:
                            proto += "int*,"
                            g = 1

                        if g == 0:
                            argvArray.append((itmSize, (itmSize - self.argPassMechInsSize),))
                        else: # push reg
                            argvArray.append(("BUFFER", 0x400,)) # flag to use buffer

                    if len(self.getXrefs(ref_p)) == 1:
                        nref = self.getXrefs(ref_p)[0]
                        refins = idc.GetDisasm(nref)
                        if refins.startswith(self.interConnectInstr):
                            ref_p = nref
                    i -= 1

                if (len(argvArray) == 0):
                    proto += "void,"

                if ((i == 0) or (g == 2)):
                    proto = proto[:-1] + ");"

                self.protoArray[proto] = argvArray
                self.log(4, "Proto, args pattern found at 0x%x is: %s->%s" % (ref, proto, repr(argvArray)))

                try:
                    tmpSorterDict[proto] = tmpSorterDict[proto] + 1
                except:
                    tmpSorterDict[proto] = 1

            itms = sorted(tmpSorterDict.iteritems(), key=lambda (k,v):(v,k), reverse=True)
            (prt, n,) = itms[0]

            if (len(prt.split(",")) > self.MAX_ARGS):
                self.log(2, "too many arguments in proto from protoanalyser, ignoring function")
                return None

            if (idbProto != None):
                self.log(3, "proto from IDB: %s" % idbProto)
                nIdbPrtArgs = len(idbProto.split(","))
                if (nIdbPrtArgs > self.MAX_ARGS):
                    self.log(3, "too many arguments in proto from IDB, ignoring it and using parser found proto")
                    idbProto = prt
                self.log(2, "proto found by parser: %s, occurance: %d" % (prt, n))
                if (len(argvArray) > nIdbPrtArgs):
                    argvArray = argvArray[:nIdbPrtArgs]
                    self.protoArray[prt] = argvArray
                fn = Function(ea, idbProto, self.protoArray[prt], self.argPassMech)
            else:
                fn = Function(ea, prt, self.protoArray[prt], self.argPassMech)

            self.log(2, repr(fn))
            n = (n*100)/mr

            if (not force) and (self.argPassType != "__usercall") and (n < self.GCD_THRESOLD):
                self.log(3, "[GCD of protos from refs is %d%%, discarding above function]" % n)
                return None
            self.log(3, "[GCD of protos from refs is %d%%, considering above function as a candidate, ArgPassMech:%s]" % (n,self.argPassMech))
            return fn

        except Exception,e:
            kprint("Exception!! @analysePrototype, ", e)


    def IsFunctionPresent(self, f):
        for x in self.functionlist:
            if x.getEA() == f.getEA():
                return True
        return False

    def updateProtoFunctionNames(self):
        for x in self.functionlist:
            x.updateName()

    def addfunc(self, f):
        if (f != None):
            if not self.IsFunctionPresent(f):
                self.functionlist.append(f)
                kprint("function \"%s\" added in ProtoAnalyser" % f.getName())
            else:
                kprint("function \"%s\" Already present!" % f.getName())

    def run(self):
        self.cleanup()
        self.log(1, "START prototype analysis...\n" + "-"*50)
        i = self.topNum
        for x in self.allXrefFuncs:
            f = self.analysePrototype(x[1])
            if f != None:
                self.addfunc(f)
                i -= 1
            if (i == 0):
                break
        self.log(1, "END prototype analysis...\n" + "-"*50)
        self.flushlog()


def protoStr2FunctionInstance(s):
    global gsplitChar
    if (s != None) and (type(s).__name__ == "str"):
        np = s.split(gsplitChar)
        ea = int(np[0].strip(),16)
        if isLoaded(ea):
            try:
                return Function(ea, np[1].strip(), np[2].strip(), np[3].strip())
            except:
                kprint("unable to edit item!")
    return None

class Protoviewer (idaapi.simplecustviewer_t):
    global gsplitChar
    splitChar = gsplitChar

    def __init__(self):
        global pa
        if pa == None:
            pa = ProtoAnalyser()
            pa.run()
        self.Create()
        self.Show()

    def Create (self):
        global pa
        global pluginVname

        self.menu_jumpto        = None
        self.menu_deleteitem    = None
        self.menu_edititem      = None
        self.menu_additem       = None
        self.menu_refresh       = None
        self.menu_runAsDecryptor= None
        self.menu_quitPlugin    = None
        self.menu_setInitializerExecution = None

        if not idaapi.simplecustviewer_t.Create (self, pluginVname + " - Top Xref'ed functions"):
            return False
        if pa:
            self.refresh()
        else:
            self.ClearLines()
        return True

    def Cleanup(self):
        global pa
        self.Close()
        if pa != None:
            pa.cleanup()
            pa = None
        return True

    def refresh(self):
        self.ClearLines()
        self.checkItemUpdates()
        header = " Virtual Addr " + self.splitChar + " " * 30 + "  Function Prototype " +  " " * 30 + self.splitChar
        header += " " * 11 + "   Argument Array  " + " " * 11 + self.splitChar + " Type " + self.splitChar + " xrefs"
        self.AddLine(header)
        for line in self.get_colored_lines():
            self.AddLine(line)
        self.Refresh()

    def get_colored_line(self, n):
        ea, nrefs, proto, argSizeArr, argPassMech = self.get_item(n).getData()
        cline = idaapi.COLSTR("  0x%08X  %c " % (ea,self.splitChar), idaapi.SCOLOR_NUMBER)
        proto_v = idaapi.COLSTR("%-80s" % proto, idaapi.SCOLOR_CODNAME)

        argArr_v = idaapi.COLSTR("%c %-40s" % (self.splitChar, repr(argSizeArr)), idaapi.SCOLOR_STRING)
        mech_v = idaapi.COLSTR("%c %-5s" % (self.splitChar,argPassMech), idaapi.SCOLOR_KEYWORD)

        nrefs_v = idaapi.COLSTR("%c (%03d)" % (self.splitChar,nrefs), idaapi.SCOLOR_AUTOCMT)
        cline += proto_v + argArr_v + mech_v + nrefs_v
        return cline

    def get_colored_lines(self):
        lines = []
        for i in xrange (self.get_number_of_items()):
            l = self.get_colored_line(i)
            lines.append(l)
        return lines

    def get_number_of_items(self):
        global pa
        return len(pa.functionlist)

    def checkItemUpdates(self):
        global pa
        pa.updateProtoFunctionNames()

    def get_item(self, n):
        global pa
        return pa.functionlist[n]

    def insert_item(self, n, v):
        global pa
        pa.functionlist.insert(n, v)

    def append_item(self, v):
        global pa
        pa.functionlist.append(v)

    def remove_item(self, n):
        global pa
        pa.functionlist.pop(n)

    def get_item_EA(self, n):
        f = self.get_item(n)
        return f.getEA()

    def get_item_name(self, n):
        f = self.get_item(n)
        return f.getName()

    def update_item(self, n, v):
        global pa
        pa.functionlist[n] = v

    def add_item(self):
        kprint("[TIP] Did you know, you can place your cursor over desired function in dissasembly view and invoke Ctrl+Y to add")
        newProto = AskStr ("0x000000 | int __cdecl func(int, int, int); | [(5,4),(\"const\",0xb0),(\"buffer\",0x100),] | \"push\" | 0", "Enter new prototype")
        if newProto != None:
            try:
                f = protoStr2FunctionInstance(newProto)
                if (f != None):
                    self.append_item(f)
            except:
                kprint("unable to add item!")
        self.refresh()
        return True

    def edit_item(self):
        n = self.GetLineNo()
        if n != 0:
            (ln,b,c) = self.GetLine(n)
            if ln == None:
                return True
            ln = tag_remove(ln)
            l = ln.split(self.splitChar)
            nln = ""
            for x in l:
                nln += x.strip() + " |"
            nln = nln[:-1]
            nps = AskStr (nln, "Please fix prototype")
            if nps != None:
                f = protoStr2FunctionInstance(nps)
                if (f != None):
                    try:
                        self.update_item(n-1, f)
                    except:
                        kprint("unable to edit item!")
            self.refresh()
        return True

    def delete_item(self, ask = True):
        if self.get_number_of_items():
            result = 1
            if ask:
                result = AskYN(1, "Delete item?")
            if result == 1:
                ln = self.GetLineNo()
                if ln != 0:
                    self.remove_item(ln-1)
                    self.refresh()
        return True


    def OnDblClick(self, shift):
        n = self.GetLineNo()
        if n != 0:
            Jump(self.get_item_EA(n-1))
        return True


    def OnHint(self, n):
        global pa
        ret = None
        if n != 0:
            ea = self.get_item_EA(n-1)
            xrefs = pa.getXrefs(ea)
            if len(xrefs) > 0:
                nx = xrefs[0]
                hintStr = ""
                for i in xrange(0, 10):
                    hintStr = idaapi.COLSTR(idc.GetDisasm(nx), idaapi.SCOLOR_MACRO) + "\n" + hintStr
                    nx = idc.PrevHead(nx, 0)
                    if (nx == BADADDR):
                        break
                t = "Code @ First xref 0x%x\n" % xrefs[0]
                hintStr = idaapi.COLSTR(t, idaapi.SCOLOR_NUMBER) + hintStr
                ret = (12, hintStr)
        return ret


    def setInitializerExecution(self):
        n = self.GetLineNo()
        if n!= 0:
            f = self.get_item(n-1)
        else:
            return False

        initCodeSt = idc.BeginEA()
        initCodeSt = idc.AskAddr(initCodeSt, "Enter the start address of the Code Initializer that you want to run before Decryption Routine")
        if initCodeSt == BADADDR:
            kprint ("[!] Cancelled")
            return False

        endCodeSt = initCodeSt + idc.ItemSize(initCodeSt)
        endCodeSt = idc.AskAddr(endCodeSt, "Enter the end address of the Code Initializer that you want to run before Decryption Routine")
        if endCodeSt == BADADDR:
            kprint ("[!] Cancelled")
            return False

        return f.setInitializerExecution(initCodeSt, endCodeSt)


    def runDecryptor(self):
        global decViews
        global decResults
        reRun = False
        n = self.GetLineNo()
        if n!= 0:
            f = self.get_item(n-1)
        else:
            return False
        rf = f.getEA()
        if rf not in decResults:
            ans = AskYN(1, "Run function \"%s\" as decryptor?" % self.get_item_name(n-1))
            if ans != 1:
                return False
        else:
            ans = AskYN(0, "Re-Run function \"%s\" as decryptor?" % self.get_item_name(n-1))
            if ans != 1:
                reRun = False
            else:
                reRun = True
        try:
            v = decViews[rf]
            v.Close()
        except:
            #kprint("view not found for 0x%x creting new one" % rf)
            pass

        decViews[rf] = decViewer(f, reRun)
        return True


    def show_xrefs(self):
        n = self.GetLineNo()
        if n != 0:
            ea = self.get_item_EA(n-1)
            open_xrefs_window(ea)

    def quit_plugin(self):
        global kryptonInst
        kprint("Quiting Plugin, restart with Ctrl+F8");
        self.Cleanup()
        kryptonInst.term()

    def OnKeydown(self, vkey, shift):
        # escape
        if vkey == 27:
            pass
        # enter
        elif vkey == 13:
            n = self.GetLineNo()
            if n != 0:
                Jump (self.get_item_EA(n-1))

        elif vkey == ord('D'):
            self.delete_item()

        elif vkey == ord("E"):
            self.edit_item()

        elif vkey == ord("A"):
            self.add_item()

        elif vkey == ord("R"):
            self.refresh()

        elif vkey == ord("K"):
            self.runDecryptor()

        elif vkey == ord("X"):
            self.show_xrefs()

        elif vkey == ord("I"):
            self.setInitializerExecution()

        elif vkey == ord("Q"):
            self.quit_plugin()

        else:
            return False

        return True


    def OnPopup(self):
        self.ClearPopupMenu()
        if not self.Count():
            self.AddPopupMenu("-")
            self.menu_refresh = self.AddPopupMenu("(r) Refresh")
        else:
            self.AddPopupMenu("-")
            self.menu_refresh = self.AddPopupMenu("(r) Refresh")
            self.menu_refresh = self.AddPopupMenu("(x) Xrefs")
            self.AddPopupMenu("-")
            self.menu_jumpto = self.AddPopupMenu("(Enter) Jump to item address")
            self.menu_edititem = self.AddPopupMenu("(e) Edit Prototype")
            self.AddPopupMenu("-")
            self.menu_additem = self.AddPopupMenu("(a) Add Function Manually")
            self.menu_deleteitem = self.AddPopupMenu("(d) Delete item")
            self.AddPopupMenu("-")
            self.menu_runAsDecryptor  = self.AddPopupMenu("(k) Run As Decryptor")
            self.AddPopupMenu("-")
            self.menu_setInitializerExecution  = self.AddPopupMenu("(i) Set Initializer Execution")
            self.AddPopupMenu("-")
            self.menu_quitPlugin = self.AddPopupMenu("(q) Quit Plugin")

        return True


    def OnPopupMenu(self, menu_id):
        if menu_id == self.menu_refresh:
            self.refresh()

        elif menu_id == self.menu_jumpto:
            n = self.GetLineNo()
            if n!= 0:
                Jump(self.get_item_EA(n-1))

        elif menu_id == self.menu_runAsDecryptor:
            self.runDecryptor()

        elif menu_id == self.menu_deleteitem:
            self.delete_item()

        elif menu_id == self.menu_additem:
            self.add_item()

        elif menu_id == self.menu_edititem:
            self.edit_item()

        elif menu_id == self.menu_setInitializerExecution:
             self.setInitializerExecution()

        elif menu_id == self.menu_quitPlugin:
            self.quit_plugin()

        else:
            return False

        return True

    def OnCursorPosChanged(self):
        self.refresh()
        return True


def addFunctionProtoCB():
    global pa
    global pvv
    try:
        ea = get_func(here()).startEA
        if isLoaded(ea):
            if pvv != None and pa != None:
                f = pa.analysePrototype(ea, True)
                if f != None:
                    pa.addfunc(f)
                    pvv.refresh()
            else:
                kprint("function not added as no instance of protoviewer found!")
        else:
            kprint ("0x%x bad Address!" % here())
    except Exception,e:
        kprint("Exception!! @addFunctionProtoCB, ", e)

def reCreateClosedProtoView():
    global pvv
    if pvv:
        pvv.Close()
    pvv = Protoviewer()


def getProcessorName():
    inf = idaapi.get_inf_structure()
    return inf.procName.strip('\x00')


# Globals
pvv         = None  # protoViewer
decViews    = {}    # decryptViewer instance array for each function decrypted
pa          = None  # protoAnalyser
decResults  = {}    # decryptor results array
kryptonInst = None
PLUGIN_MODE = True

class kplugin_t (idaapi.plugin_t):
    flags = 0
    comment = ""
    help = ""
    wanted_name = pluginVname
    wanted_hotkey = "Ctrl-F8"
    m1 = m3 = None

    def init(self):
        if getProcessorName().startswith('metapc'):
            return PLUGIN_KEEP
        return PLUGIN_SKIP

    def run(self, arg):
        global pvv
        Wait()
        if not pvv:
            pvv = Protoviewer()
            self.m1 = idaapi.add_menu_item("Edit/", pluginVname + " - Add func proto for Decryptor", "Ctrl+Y", 0, addFunctionProtoCB, None)
            if self.m1 == None:
                kprint ("add menu item failed!")
            self.m3 = idaapi.add_menu_item("View/Open subviews/", pluginVname, "", 0, reCreateClosedProtoView, None)
            if self.m3 == None:
                kprint ("add menu item failed!")
        else:
            if not pvv.IsFocused():
                pvv.Close()
                pvv = Protoviewer()

    def term(self):
        global pvv
        if pvv != None:
            pvv.Cleanup()
            pvv = None
        if self.m1 != None:
            idaapi.del_menu_item(self.m1)
        if self.m3 != None:
            idaapi.del_menu_item(self.m3)
        pass

if (PLUGIN_MODE):
    def PLUGIN_ENTRY ():
        global kryptonInst
        kryptonInst = kplugin_t()
        return kryptonInst
else:
    Wait()
    if not pvv:
        pvv = Protoviewer()
    else:
        if not pvv.IsFocused():
            pvv.Close()
            pvv = Protoviewer()


