# plugins

repository: fireundubh/IDA7-Alleycat
file: alleycat.py
description: Alleycat plugin by devttys0, ported to IDA 7
wanted_name: AlleyCat
created_at: 2018-04-29 00:17:44+00:00
pushed_at: 2018-04-29 00:25:01+00:00
url:
https://github.com/fireundubh/IDA7-Alleycat/blob/33c210e6e868cf387953087543f8cb404869d5ee/alleycat.py

content:
import time

import idaapi
import idautils
import idc

# This limits the depth of any individual path, as well as the maximum
# number of paths that will be searched; this is needed for practical
# reasons, as IDBs with tens of thousands of functions take a long time
# to exhaust all possible paths without some practical limitation.
#
# This is global so it's easy to change from the IDAPython prompt.
ALLEYCAT_LIMIT = 10000


class AlleyCatException(Exception):
    pass


class AlleyCat(object):
    """
    Class which resolves code paths. This is where most of the work is done.
    """

    def __init__(self, start, end, quiet=False):
        """
        Class constructor.

        @start - The start address.
        @end   - The end address.

        Returns None.
        """
        global ALLEYCAT_LIMIT
        self.limit = ALLEYCAT_LIMIT
        self.paths = []
        self.quiet = quiet

        # We work backwards via xrefs, so we start at the end and end at the start
        if not self.quiet:
            print 'Generating call paths from %s to %s...' % (self._name(end), self._name(start))
        self._build_paths(start, end)

    @staticmethod
    def _name(ea):
        name = idc.Name(ea)
        if not name:
            name = idc.GetFuncOffset(ea)
            if not name:
                name = '0x%X' % ea
        return name

    def _add_path(self, path):
        if path not in self.paths:
            self.paths.append(path)

    def _build_paths(self, start, end=idc.BADADDR):
        partial_paths = [[start]]

        # Loop while there are still unresolve paths and while all path sizes have not exceeded ALLEYCAT_LIMIT
        while partial_paths and len(self.paths) < self.limit and len(partial_paths) < self.limit:
            # Initialize a unique set of callers for this iteration
            callers = set()

            # Callee is the last entry of the first path in partial paths.
            # The first path list will change as paths are completed and popped from the list.
            callee = partial_paths[0][-1]

            # Find all unique functions that reference the callee, assuming this path has not
            # exceeded ALLEYCAT_LIMIT.
            if len(partial_paths[0]) < self.limit:
                for xref in idautils.XrefsTo(callee):
                    caller = AlleyCatFunctionPaths.get_code_block(xref.frm)
                    if caller and caller.startEA not in callers:
                        callers.add(caller.startEA)

            # If there are callers to the callee, remove the callee's current path
            # and insert new ones with the new callers appended.
            if callers:
                base_path = partial_paths.pop(0)
                for caller in callers:

                    # Don't want to loop back on ourselves in the same path
                    if caller in base_path:
                        continue

                    # If we've reached the desired end node, don't go any further down this path
                    if caller == end:
                        self._add_path((base_path + [caller])[::-1])
                    else:
                        partial_paths.append(base_path + [caller])
            # Else, our end node is not in this path, so don't include it in the finished path list.
            elif end not in partial_paths[0]:
                partial_paths.pop(0)
            # If there were no callers then this path has been exhaused and should be
            # popped from the partial path list into the finished path list.
            elif end in partial_paths[0]:
                # Paths start with the end function and end with the start function; reverse it.
                self._add_path(partial_paths.pop(0)[::-1])


class AlleyCatFunctionPaths(AlleyCat):
    def __init__(self, start_ea, end_ea, quiet=False):

        # We work backwards via xrefs, so we start at the end and end at the start
        try:
            start = idaapi.get_func(end_ea).startEA
        except:
            raise AlleyCatException('Address 0x%X is not part of a function!' % end_ea)

        try:
            end = idaapi.get_func(start_ea).startEA
        except:
            end = idc.BADADDR

        super(AlleyCatFunctionPaths, self).__init__(start, end, quiet)

    @staticmethod
    def get_code_block(ea):
        return idaapi.get_func(ea)


class AlleyCatCodePaths(AlleyCat):

    def __init__(self, start_ea, end_ea, quiet=False):
        end_func = idaapi.get_func(end_ea)
        start_func = idaapi.get_func(start_ea)

        if not start_func:
            raise AlleyCatException('Address 0x%X is not part of a function!' % start_ea)
        if not end_func:
            raise AlleyCatException('Address 0x%X is not part of a function!' % end_ea)
        if start_func.startEA != end_func.startEA:
            raise AlleyCatException('The start and end addresses are not part of the same function!')

        self.func = start_func
        self.blocks = [block for block in idaapi.FlowChart(self.func)]

        # We work backwards via xrefs, so we start at the end and end at the start
        end_block = self._get_code_block(start_ea)
        start_block = self._get_code_block(end_ea)

        if not end_block:
            raise AlleyCatException('Failed to find the code block associated with address 0x%X' % start_ea)
        if not start_block:
            raise AlleyCatException('Failed to find the code block associated with address 0x%X' % end_ea)

        super(AlleyCatCodePaths, self).__init__(start_block.startEA, end_block.startEA, quiet)

    def _get_code_block(self, ea):
        for block in self.blocks:
            if block.startEA <= ea < block.endEA:
                return block
        return None


# Everything below here is just IDA UI/Plugin stuff ###


# noinspection PyAttributeOutsideInit
class AlleyCatGraphHistory(object):
    """
    Manages include/exclude graph history.
    """

    INCLUDE_ACTION = 0
    EXCLUDE_ACTION = 1

    def __init__(self):
        self.reset()

    def reset(self):
        self.history = []
        self.includes = []
        self.excludes = []
        self.history_index = 0
        self.include_index = 0
        self.exclude_index = 0

    def update_history(self, action):
        if self.excludes and len(self.history) - 1 != self.history_index:
            self.history = self.history[0:self.history_index + 1]
        self.history.append(action)
        self.history_index = len(self.history) - 1

    def add_include(self, obj):
        if self.includes and len(self.includes) - 1 != self.include_index:
            self.includes = self.includes[0:self.include_index + 1]
        self.includes.append(obj)
        self.include_index = len(self.includes) - 1
        self.update_history(self.INCLUDE_ACTION)

    def add_exclude(self, obj):
        if len(self.excludes) - 1 != self.exclude_index:
            self.excludes = self.excludes[0:self.exclude_index + 1]
        self.excludes.append(obj)
        self.exclude_index = len(self.excludes) - 1
        self.update_history(self.EXCLUDE_ACTION)

    def get_includes(self):
        return set(self.includes[0:self.include_index + 1])

    def get_excludes(self):
        return set(self.excludes[0:self.exclude_index + 1])

    def undo(self):
        if self.history:
            if self.history[self.history_index] == self.INCLUDE_ACTION:
                if self.include_index >= 0:
                    self.include_index -= 1
            elif self.history[self.history_index] == self.EXCLUDE_ACTION:
                if self.exclude_index >= 0:
                    self.exclude_index -= 1

            self.history_index -= 1
            if self.history_index < 0:
                self.history_index = 0

    def redo(self):
        self.history_index += 1
        if self.history_index >= len(self.history):
            self.history_index = len(self.history) - 1

        if self.history[self.history_index] == self.INCLUDE_ACTION:
            if self.include_index < len(self.includes) - 1:
                self.include_index += 1
        elif self.history[self.history_index] == self.EXCLUDE_ACTION:
            if self.exclude_index < len(self.excludes) - 1:
                self.exclude_index += 1


class AlleyCatGraph(idaapi.GraphViewer):
    """
    Displays the graph and manages graph actions.
    """

    def __init__(self, results, title='AlleyCat Graph'):
        idaapi.GraphViewer.__init__(self, title)
        self.results = results

        self.nodes_ea2id = {}
        self.nodes_id2ea = {}
        self.edges = {}
        self.end_nodes = []
        self.edge_nodes = []
        self.start_nodes = []

        self.history = AlleyCatGraphHistory()
        self.include_on_click = False
        self.exclude_on_click = False

    # noinspection PyAttributeOutsideInit
    def Show(self):
        """
        Display the graph.

        Returns True on success, False on failure.
        """
        if not idaapi.GraphViewer.Show(self):
            return False
        else:
            self.cmd_undo = self.AddCommand('Undo', '')
            self.cmd_redo = self.AddCommand('Redo', '')
            self.cmd_reset = self.AddCommand('Reset graph', '')
            self.cmd_exclude = self.AddCommand('Exclude node', '')
            self.cmd_include = self.AddCommand('Include node', '')
            self.cmd_unhighlight = self.AddCommand('Temporarily un-highlight all paths', '')
            return True

    def OnRefresh(self):
        # Clear the graph before refreshing
        self.clear()
        self.nodes_ea2id = {}
        self.nodes_id2ea = {}
        self.edges = {}
        self.end_nodes = []
        self.edge_nodes = []
        self.start_nodes = []

        includes = self.history.get_includes()
        excludes = self.history.get_excludes()

        for path in self.results:
            parent_node = None

            # Check to see if this path contains all nodes marked for explicit inclusion
            if (set(path) & includes) != includes:
                continue

            # Check to see if this path contains any nodes marked for explicit exclusion
            if (set(path) & excludes) != set():
                continue

            for ea in path:
                # If this node already exists, use its existing node ID
                if ea in self.nodes_ea2id:
                    this_node = self.nodes_ea2id[ea]
                # Else, add this node to the graph
                else:
                    this_node = self.AddNode(self.get_name_by_ea(ea))
                    self.nodes_ea2id[ea] = this_node
                    self.nodes_id2ea[this_node] = ea

                # If there is a parent node, add an edge between the parent node and this one
                if parent_node is not None:
                    self.AddEdge(parent_node, this_node)
                    if this_node not in self.edges[parent_node]:
                        self.edges[parent_node].append(this_node)

                # Update the parent node for the next loop
                parent_node = this_node
                if parent_node not in self.edges:
                    self.edges[parent_node] = []

                # Highlight this node in the disassembly window
                # self.highlight(ea)

            try:
                # Track the first, last, and next to last nodes in each path for
                # proper colorization in self.OnGetText.
                self.start_nodes.append(self.nodes_ea2id[path[0]])
                self.end_nodes.append(self.nodes_ea2id[path[-1]])
                self.edge_nodes.append(self.nodes_ea2id[path[-2]])
            except:
                pass

        return True

    def OnGetText(self, node_id):
        color = idc.DEFCOLOR

        if node_id in self.edge_nodes:
            color = 0x00ffff
        elif node_id in self.start_nodes:
            color = 0x00ff00
        elif node_id in self.end_nodes:
            color = 0x0000ff

        return self[node_id], color

    def OnHint(self, node_id):
        hint = ''

        try:
            for edge_node in self.edges[node_id]:
                hint += '%s\n' % self[edge_node]
        except:
            pass

        return hint

    def OnCommand(self, cmd_id):
        if self.cmd_undo == cmd_id:
            if self.include_on_click or self.exclude_on_click:
                self.include_on_click = False
                self.exclude_on_click = False
            else:
                self.history.undo()
            self.Refresh()
        elif self.cmd_redo == cmd_id:
            self.history.redo()
            self.Refresh()
        elif self.cmd_include == cmd_id:
            self.include_on_click = True
        elif self.cmd_exclude == cmd_id:
            self.exclude_on_click = True
        elif self.cmd_reset == cmd_id:
            self.include_on_click = False
            self.exclude_on_click = False
            self.history.reset()
            self.Refresh()
        # elif self.cmd_unhighlight == cmd_id:
        #     self.unhighlight_all()

    def OnClick(self, node_id):
        if self.include_on_click:
            self.history.add_include(self.nodes_id2ea[node_id])
            self.include_on_click = False
        elif self.exclude_on_click:
            self.history.add_exclude(self.nodes_id2ea[node_id])
            self.exclude_on_click = False
        self.Refresh()

    def OnDblClick(self, node_id):
        xref_locations = []
        node_ea = self.get_ea_by_name(self[node_id])

        if node_id in self.edges:
            for edge_node_id in self.edges[node_id]:

                edge_node_name = self[edge_node_id]
                edge_node_ea = self.get_ea_by_name(edge_node_name)

                if edge_node_ea != idc.BADADDR:
                    for xref in idautils.XrefsTo(edge_node_ea):
                        # Is the specified node_id the source of this xref?
                        if self.match_xref_source(xref, node_ea):
                            xref_locations.append((xref.frm, edge_node_ea))

        if xref_locations:
            xref_locations.sort()

            print ''
            print 'Path Xrefs from %s:' % self[node_id]
            print '-' * 100
            for (xref_ea, dst_ea) in xref_locations:
                print '%-50s  =>  %s' % (self.get_name_by_ea(xref_ea), self.get_name_by_ea(dst_ea))
            print '-' * 100
            print ''

            idc.Jump(xref_locations[0][0])
        else:
            idc.Jump(node_ea)

    def OnClose(self):
        # TODO: Add a 'do not ask again' feature?
        # if idc.AskYN(1, "Path nodes have been highlighted in the disassembly window. Undo highlighting?") == 1:
        #     self.unhighlight_all()
        pass

    @staticmethod
    def match_xref_source(xref, source):
        # TODO: This must be modified if support for graphing function blocks is added.
        return (xref.type != idc.fl_F) and (idc.GetFunctionAttr(xref.frm, idc.FUNCATTR_START) == source)

    @staticmethod
    def get_ea_by_name(name):
        """
        Get the address of a location by name.

        @name - Location name

        Returns the address of the named location, or idc.BADADDR on failure.
        """
        # This allows support of the function offset style names (e.g., main+0C)
        # TODO: Is there something in the IDA API that does this already??
        # TODO: AppCall maybe? http://www.hexblog.com/?p=112 -fireundubh
        ea = idc.BADADDR
        if '+' in name:
            (func_name, offset) = name.split('+')
            base_ea = idc.LocByName(func_name)
            if base_ea != idc.BADADDR:
                try:
                    ea = base_ea + int(offset, 16)
                except:
                    ea = idc.BADADDR
        else:
            ea = idc.LocByName(name)
            if ea == idc.BADADDR:
                try:
                    ea = int(name, 0)
                except:
                    ea = idc.BADADDR
        return ea

    def clear(self):
        # Clears the graph and unhighlights the disassembly
        self.Clear()
        # self.unhighlight_all()

    @staticmethod
    def get_name_by_ea(ea):
        """
        Get the name of the specified address.

        @ea - Address.

        Returns a name for the address, one of idc.Name, idc.GetFuncOffset or 0xXXXXXXXX.
        """
        name = idc.Name(ea)
        if not name:
            name = idc.GetFuncOffset(ea)
            if not name:
                name = '0x%X' % ea
        return name

    @staticmethod
    def colorize_node(ea, color):
        # Colorizes an entire code block
        func = idaapi.get_func(ea)
        if not func:
            return

        for block in idaapi.FlowChart(func):
            if block.startEA <= ea < block.endEA:
                ea = block.startEA
                while ea < block.endEA:
                    idaapi.set_item_color(ea, color)
                    ea = idc.NextHead(ea)
                break

    def highlight(self, ea):
        # Highlights an entire code block
        # self.colorize_node(ea, color=0xFFFBCC)
        pass

    def unhighlight(self, ea):
        # Unhighlights an entire code block
        self.colorize_node(ea, idc.DEFCOLOR)

    def unhighlight_all(self):
        # Unhighlights all code blocks
        for path in self.results:
            for ea in path:
                self.unhighlight(ea)


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

    def activate(self, ctx):
        target = idapathfinder_t.current_function()

        if not target:
            return 1

        sources = idapathfinder_t.get_user_selected_functions(many=True)
        if not sources:
            return 1

        idapathfinder_t.find_and_plot_paths(sources, [target])
        return 0

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


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

    def activate(self, ctx):
        source = idapathfinder_t.current_function()

        if not source:
            return 1

        targets = idapathfinder_t.get_user_selected_functions(many=True)
        if not targets:
            return 1

        idapathfinder_t.find_and_plot_paths([source], targets)
        return 0

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


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

    def activate(self, ctx):
        target = idc.ScreenEA()
        source = idapathfinder_t.current_function()

        if not source:
            return 1

        idapathfinder_t.find_and_plot_paths([source], [target], AlleyCatCodePaths)
        return 0

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


# noinspection PyPep8Naming
class idapathfinder_t(idaapi.plugin_t):
    flags = 0
    comment = ''
    help = ''
    wanted_name = 'AlleyCat'
    wanted_hotkey = ''

    def __init__(self):
        self.graph = None
        self.menu_action_find_from = idaapi.action_desc_t('alleycat:find_from', 'Find paths to the current function from...', AlleycatActionHandlerFindPathsFrom())
        self.menu_action_find_to = idaapi.action_desc_t('alleycat:find_to', 'Find paths from the current function to...', AlleycatActionHandlerFindPathsTo())
        self.menu_action_find_current = idaapi.action_desc_t('alleycat:find_current', 'Find paths in the current function to the current code block', AlleycatActionHandlerFindPathsToCurrentBlock())
        idaapi.register_action(self.menu_action_find_from)
        idaapi.register_action(self.menu_action_find_to)
        idaapi.register_action(self.menu_action_find_current)

    # noinspection PyMethodMayBeStatic
    def init(self):
        idaapi.attach_action_to_menu('View/Graphs/Find paths to the current function from...', 'alleycat:find_from', idaapi.SETMENU_APP)
        idaapi.attach_action_to_menu('View/Graphs/Find paths from the current function to...', 'alleycat:find_to', idaapi.SETMENU_APP)
        idaapi.attach_action_to_menu('View/Graphs/Find paths in the current function to the current code block', 'alleycat:find_current', idaapi.SETMENU_APP)
        return idaapi.PLUGIN_KEEP

    # noinspection PyMethodMayBeStatic
    def term(self):
        idaapi.detach_action_from_menu('View/Graphs/Find paths to the current function from...', 'alleycat:find_from')
        idaapi.detach_action_from_menu('View/Graphs/Find paths from the current function to...', 'alleycat:find_to')
        idaapi.detach_action_from_menu('View/Graphs/Find paths in the current function to the current code block', 'alleycat:find_current')
        return None

    def run(self, arg):
        pass

    @staticmethod
    def current_function():
        result = idaapi.get_func(idc.ScreenEA())
        if result:
            return result.startEA
        else:
            print 'No linear address found at cursor'

    @staticmethod
    def find_and_plot_paths(sources, targets, klass=AlleyCatFunctionPaths):
        results = []

        for target in targets:
            for source in sources:
                s = time.time()
                r = klass(source, target).paths
                e = time.time()
                print 'Found %d paths in %f seconds.' % (len(r), (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 not results:
            return

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

        idapathfinder_t.graph = AlleyCatGraph(results, 'Path Graph')
        idapathfinder_t.graph.Show()

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

        while True:
            func = 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 func or func == idc.BADADDR or func == current_function:
                break
            elif func not in functions:
                functions.append(func)

            if not many:
                break

        return functions


def PLUGIN_ENTRY():
    return idapathfinder_t()


repository: Aspirin493/xorstr-decrypt-idaplugin
file: xordecrypt.py
wanted_name: XOR String Decoder
comment: XOR String Decoder
created_at: 2023-04-15 14:40:56+00:00
pushed_at: 2023-04-15 14:42:29+00:00
description:
This is a plugin for IDA Pro that will help you deobfuscate xorstr strings in a windows x64 application.

url:
https://github.com/Aspirin493/xorstr-decrypt-idaplugin/blob/d40afbdff85f239e86847b1f68b43196ca3c444e/xordecrypt.py

content:
import idaapi
import idc

# Constants for XOR string analysis
MAX_XOR_STR_LEN = 64
MAX_XOR_STR_COUNT = 256

class XorStringDecoder(idaapi.plugin_t):
    flags = 0
    comment = "XOR String Decoder"
    help = "XOR String Decoder"
    wanted_name = "XOR String Decoder"
    wanted_hotkey = ""

    def init(self):
        return idaapi.PLUGIN_OK

    def run(self, arg):
        print("[XOR String Decoder] Starting...")

        # List of found XOR strings
        xor_strings = []

        # List of byte strings that were found in the program
        byte_strings = []

        # Search for all byte strings in the program
        for i in range(idaapi.get_segm_qty()):
            seg = idaapi.getnseg(i)
            if not seg:
                continue

            ea = seg.start_ea
            end_ea = seg.end_ea
            while ea < end_ea:
                byte_string = idaapi.get_strlit_contents(ea, MAX_XOR_STR_LEN, idaapi.ASCSTR_C)
                if byte_string:
                    byte_strings.append(byte_string)
                    ea += len(byte_string) + 1
                else:
                    ea += 1

        # Analysis of found byte strings and search for xor strings
        for byte_string in byte_strings:
            if len(byte_string) < 4 or len(xor_strings) >= MAX_XOR_STR_COUNT:
                continue

            xor_string = ""
            key = 0x0

            # Looking for a key to deobfuscate a string
            for i in range(2, MAX_XOR_STR_LEN):
                k = ord(byte_string[0]) ^ ord(byte_string[i])
                valid_key = True
                for j in range(1, len(byte_string)):
                    if (j == i):
                        continue
                    if (ord(byte_string[j]) ^ k) != ord(byte_string[j-1]):
                        valid_key = False
                        break
                if valid_key:
                    key = k
                    break

            # If the key is found, deobfuscate the string
            if key:
                for b in byte_string:
                    xor_string += chr(ord(b) ^ key)
                xor_strings.append(xor_string)

        # Showing the found xor strings to the user
        if xor_strings:
            print("[XOR String Decoder] Found %d xor strings:" % len(xor_strings))
            for xor_string in xor_strings:
                print("[XOR String Decoder]   %s" % xor_string)
        else:
            print("[XOR String Decoder] No xor strings found.")

    def term(self):
        pass

def PLUGIN_ENTRY():
    return XorStringDecoder()


repository: mcdulltii/obfDetect
file: obfDetect.py
description: IDA plugin to pinpoint obfuscated code
url: https://github.com/mcdulltii/obfDetect/blob/c1c36d22268b4874edb852cec517060a580b1546/obfDetect.py
wanted_name: Obfuscation Detection
comment: Calculates binary obfuscation heuristics
created_at: 2021-08-17 11:35:18+00:00
pushed_at: 2022-04-29 15:57:07+00:00
content:
from idc import *
import idaapi
from idaapi import *
import idautils
import ida_kernwin

from obfDetect import gui

PLUGIN_VERSION = "1.7"
IDAVERSIONS = "IDA PRO 7.4+"
AUTHORS = "mcdulltii"
DATE = "2022"

def banner():
    banner_options = (PLUGIN_VERSION, AUTHORS, DATE, IDAVERSIONS)
    banner_titles = "Obfuscation Detection v%s - (c) %s - %s - %s" % banner_options
    # print plugin banner
    print("\n---[" + banner_titles + "]---\n")

# Obfuscation Detection Handler
class ObfDetectHandler(action_handler_t):
    def __init__(self):
        action_handler_t.__init__(self)

    # Run script when invoked.
    def activate(self, ctx):
        gui.show_gui()

    def update(self, ctx):
        return AST_ENABLE_ALWAYS

class obfDetect_plugin_t(plugin_t):
    flags = PLUGIN_HIDE
    comment = "Calculates binary obfuscation heuristics"
    help = "Obfuscation Detection"
    wanted_name = "Obfuscation Detection"
    wanted_hotkey = ""

    def editor_menuaction(self):
        action_desc = action_desc_t(
            'my:detectoraction',  # The action name. This acts like an ID and must be unique
            'Obfuscation Detection',  # The action text.
            ObfDetectHandler(),  # The action handler.
            'Ctrl+Shift+H',  # Optional: the action shortcut
            '',  # Optional: the action tooltip (available in menus/toolbar)
            122  # icon number
        )

        # Register the action
        register_action(action_desc)

        attach_action_to_menu(
            'File/Editor...',  # The relative path of where to add the action
            'my:detectoraction',  # The action ID (see above)
            SETMENU_APP)  # We want to append the action after the 'Manual instruction...

        form = ida_kernwin.get_current_widget()
        attach_action_to_popup(form, None, "my:detectoraction", None)

    def init(self):
        """
        This is called by IDA when it is loading the plugin.
        """
        # attempt plugin initialization
        try:
            self._install_plugin()
        # failed to initialize or integrate the plugin, log and skip loading
        except Exception as e:
            form = ida_kernwin.get_current_widget()
            pass
        return PLUGIN_KEEP

    def _install_plugin(self):
        """
        Initialize & integrate the plugin into IDA.
        """
        self.editor_menuaction()
        self._init()

    def run(self, arg = 0):
        # We need the calls again if we wanna load it via File/Plugins/editor
        msg("Obfuscation Detection loaded.\nUse Alt+E hot key to quick load.\n")
        handler = ObfDetectHandler()
        handler.activate(self)

    def term(self):
        pass

def PLUGIN_ENTRY():
    banner()
    return obfDetect_plugin_t()


repository: dockerizeme/dockerizeme
file: hard-gists/b5e7a80691ed5e44e7fea1964bae18dc/snippet.py
description: Automatically generate environment configurations as Dockerfiles
wanted_name: DynHints
comment: Display dynamically-generated hints.
created_at: 2018-08-12 21:21:04+00:00
pushed_at: 2022-12-07 22:45:41+00:00
url:
https://github.com/dockerizeme/dockerizeme/blob/408f3fa3d36542d8fc1236ba1cac804de6f14b0c/hard-gists/b5e7a80691ed5e44e7fea1964bae18dc/snippet.py

content:
'''
example of:
  - using ctypes with the IDA SDK
  - providing custom UI hints with dynamic data from Python

in this silly example, we display UI hints with the current timestamp.
a more useful plugin might inspect the hovered line, and display some documentation.

Author: Willi Ballenthin <william.ballenthin@fireeye.com>
Licence: Apache 2.0
'''
import sys
import ctypes

import idc
import idaapi
import idautils


def get_ida_ctypes():
    '''
    get the ida sdk dll.
    
    Args: None

    Returns:
      ctypes.CDLL: the IDA SDK DLL

    via: http://www.hexblog.com/?p=695
    '''
    idaname = 'ida64' if __EA64__ else 'ida'
    if sys.platform == 'win32':
        return ctypes.windll[idaname + '.wll']
    elif sys.platform == 'linux2':
        return ctypes.cdll['lib' + idaname + '.so']
    elif sys.platform == 'darwin':
        return ctypes.cdll['lib' + idaname + '.dylib']
    else:
        raise RuntimeError('unknown platform: ' + sys.platform)


c_long_p = ctypes.POINTER(ctypes.c_long)
c_int_p = ctypes.POINTER(ctypes.c_int)
c_char_pp = ctypes.POINTER(ctypes.c_char_p)
    

# ctypes declaration for the `hook_type_t` type.
#
# the idaapi calling convention is stdcall, so we use WINFUNCTYPE;
#  it is *not* cdecl (ctypes.CFUNCTYPE).
#
# relevant idasdk documentation: 
#
#    typedef int idaapi hook_cb_t(
#      void *user_data,
#      int notification_code,
#      va_list va);
HookCb = ctypes.WINFUNCTYPE(
    # return type
    ctypes.c_int,      # int idaapi

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


class HOOK_TYPES:
    '''
    types of events that be hooked to with hook_to_notification_point().
    this corresponds to `hook_type_t`.
    '''
    # Hook to the processor module.
    HT_IDP = 0

    # Hook to the user interface.
    HT_UI = 1

    # Hook to the debugger.
    HT_DBG = 2

    # Hook to the database events.
    HT_IDB = 3

    # Internal debugger events.
    HT_DEV = 4

    # Custom/IDA views notifications.
    HT_VIEW = 5

    # Output window notifications.
    HT_OUTPUT = 6

    HT_LAST = 7


def do_hook_to_notification_point(dll, hook_type, cb, user_data):
    '''
    register a callback for a class of events in IDA.
    
    Args:
      dll (ctypes.CDLL): the IDA SDK DLL
      hook_type (int): hook type from HOOK_TYPES enum
      cb (HookCb): the callback to register
      user_data (ctypes.c_void_p): context provided to callback when invoked
    
    Returns: None

    relevant idasdk documentation: 

        idaman bool ida_export 	
        hook_to_notification_point (
            hook_type_t hook_type,
            hook_cb_t *cb,
            void *user_data)
   '''
    
    hook_to_notification_point = dll.hook_to_notification_point
    hook_to_notification_point.argtypes = [
        ctypes.c_int,      # hook_type_t hook_type
        HookCb,            # hook_cb_t  *cb
        ctypes.c_void_p,   # void       *user_data
    ]

    hook_to_notification_point(hook_type, cb, user_data)


def do_unhook_from_notification_point(dll, hook_type, cb, user_data):
    '''
    unregister a callback.
    
    Args:
      dll (ctypes.CDLL): the IDA SDK DLL
      hook_type (int): hook type from HOOK_TYPES enum
      cb (HookCb): the callback to register
      user_data (ctypes.c_void_p): context provided to callback when invoked
    
    Returns: None

    relevant idasdk documentation: 

        idaman bool ida_export 	
        unhook_from_notification_point (
            hook_type_t hook_type,
            hook_cb_t *cb,
            void *user_data = NULL)
   '''
    
    unhook_from_notification_point = dll.unhook_from_notification_point
    unhook_from_notification_point.argtypes = [
        ctypes.c_int,      # hook_type_t hook_type
        HookCb,            # hook_cb_t  *cb
        ctypes.c_void_p,   # void       *user_data
    ]

    unhook_from_notification_point(hook_type, cb, user_data)


class UI_NOTIFICATIONS:

    # get single line hint for the given address.
    # this is *not* used for the disasm listing hover tooltips.
    # this is used for nav bar hover tooltips.
    # 
    # relevant idasdk documentation: 
    #
    #     cb: ui wants to display a simple hint for an address.
    #     Use this event to generate a custom hint
    #     See also more generic ::ui_get_item_hint
    #     \param ea       (::ea_t)
    #     \param buf      (char *)
    #     \param bufsize  (size_t)
    #     \return true if generated a hint
    UI_GET_EA_HINT = 79

    # i'm not sure where this is called.
    #
    # relevant idasdk documentation: 
    #
    #     cb: ui wants to display multiline hint for an item.
    #     See also more generic ::ui_get_custom_viewer_hint
    #     \param ea (ea_t, or item id like a structure or enum member)
    #     \param max_lines             (int) maximal number of lines
    #     \param[out] important_lines  (int *) number of important lines. if zero, output is ignored
    #     \param[out] hint             (::qstring *) the output string
    #     \return true if generated a hint
    UI_GET_ITEM_HINT = 80

    # get multiple line hint for the given address
    # this is invoked for at least the following views:
    #   - disasm
    #   - enums
    #   - structures
    # 
    # relevant idasdk documentation: 
    #
    #     cb: ui wants to display a hint for a viewer (idaview or custom).
    #     \param viewer (TCustomControl*) viewer
    #     \param place                 (::place_t *) current position in the viewer
    #     \param[out] important_lines  (int *) number of important lines.
    #                                         if zero, the result is ignored
    #     \param[out] hint             (::qstring *) the output string
    #     \return true if generated a hint
    UI_GET_CUSTOM_VIEWER_HINT = 91


class DynHintsPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_KEEP
    comment = "Display dynamically-generated hints."

    help = "Display dynamically-generated hints."
    wanted_name = "DynHints"
    wanted_hotkey = "Ctrl-]"

    def init(self):
        import datetime
        self.dll = get_ida_ctypes()

        # can't use a bound method as a callback (since `self` doesn't get provided)
        #  so we'll create a closure that has access to `self`.
        # 
        # via: http://stackoverflow.com/a/7261524/87207
        def cb(user_data, notification_code, va_list):
            '''
            example hook_cb_t function that handles custom viewer hints.

            Args:
              user_data (ctypes.c_void_p): context supplied to callback registration
              notification_code (int): one of the UI_NOTIFICATIONS enum values
              va_list (ctypes.c_void_p): varargs that must be manually parsed

            Returns:
              int: see notifiication code documentation for interpretation

            Notes:
              - This is a closure that expects to have 
            '''

            # this function is called *a lot*, so don't do any heavy lifting
            #  until you know its the event you want.

            # ctypes doesn't support varargs in callback functions.
            # so, we need to parse the remaining arguments ourselves.
            #
            # on windows, varargs are sequential stack locations.
            # so, lets access the members like an array of ints/pointers.
            #
            # WARNING: the following section that manually parses varargs is
            #  probably architecture and platfrom dependent!
            va_list = ctypes.cast(va_list, c_long_p)

            if notification_code == UI_NOTIFICATIONS.UI_GET_EA_HINT:
                # ea is just a number: 
                #
                #     typedef uint32 	ea_t
                #
                # via: https://www.hex-rays.com/products/ida/support/sdkdoc/pro_8h.html#a7b0aeaed04e477c02cf8ea3452002d1a
                ea = va_list[0]
                buf = ctypes.cast(va_list[1], ctypes.c_char_p)
                bufsize = va_list[2]

                print('ui_get_ea_hint:')
                print('>.. notification code: %s' % (notification_code))
                print('>.. ea: %s' % (hex(ea)))
                print('>.. buf: %s' % (buf))
                print('>.. bufsize: %s' % (hex(bufsize)))

                the_hint = datetime.datetime.now().isoformat(' ')

                self.dll.qstrncpy(buf, ctypes.c_char_p(the_hint), bufsize)
                print('<.. buf: %s' % (buf))

                return 1

            elif notification_code == UI_NOTIFICATIONS.UI_GET_CUSTOM_VIEWER_HINT:
                viewer = ctypes.cast(va_list[0], c_long_p)
                place = ctypes.cast(va_list[1], c_long_p)
                important_lines = ctypes.cast(va_list[2], c_long_p)
                hint = ctypes.cast(va_list[3], c_char_pp)

                if not place:
                    print('ui_get_custom_viewer_hint: invalid place')
                    return 0

                print('ui_get_custom_viewer_hint:')
                print('>.. notification code: %s' % (notification_code))
                print('>.. important lines: %s %s' % (important_lines, important_lines.contents))
                print('>.. hint: %s %s' % (hint, hint.contents))

                # so, we'd like to fetch the EA of the current view.
                # ideally, we'd do:
                #
                #     ea = place.toea()
                #
                # but `place` is a raw c++ object pointer, and ctypes isn't that smart.
                # next best would be to do something like:
                #
                #     place = self.dll.get_custom_viewer_place(viewer);
                #
                # however, this doesn't work because `get_custom_viewer_place` is not an exported routine.
                # it seems to be part of the IDA SDK static lib to which plugins link.
                #
                # next best would be to use `idaapi.get_custom_viewer_place`:
                #
                #     place = idaapi.get_custom_viewer_place(viewer);
                #
                # but, this doesn't work because we're mixing a ctypes pointer with a swig function.
                # so, we'll fall back to querying the current viewer, and fetching the place from there.

                # let's only display for disassembly listings
                #
                # i only know how to test the view/form type using the `get_tform_type` function.
                # therefore, we'll first query the current tform, and subsequently the current custom_viewer.
                tform = idaapi.get_current_tform()
                if idaapi.get_tform_type(tform) != idaapi.BWN_DISASM:
                    return 0

                viewer = idaapi.get_current_viewer()

                # `place` is a tuple (though techincally, a list), with elements:
                #  - place_t proxy
                #  - x position in characters
                #  - y position in characters from top of screen/form (-1 in graph view)
                place, x, y = idaapi.get_custom_viewer_place(viewer, True)

                the_hint = '0x%08X: %s' % (place.toea(), datetime.datetime.now().isoformat(' '))

                important_lines[0] = ctypes.c_long(1)
                # we don't have access to the qstring c++ class methods,
                #  so we'll use a dummy routine to correctly set our qstring contents.
                # `replace_tabs` assigns from a char * to a qstring *.
                #
                # relevant idasdk documentation: 
                #
                #    idaman THREAD_SAFE bool ida_export
                #    replace_tabs (
                #        qstring *out,
                #        const char *str,
                #        int tabsize)
                self.dll.replace_tabs(hint, ctypes.c_char_p(the_hint), 4)

                print('<.. important lines: %s %s' % (important_lines, important_lines.contents))
                print('<.. hint: %s %s' % (hint, hint.contents))

                return 1

            return 0


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

        # need to keep a ref around, or the param gets garbage collected
        self.ctx = ctypes.c_long(69)

        return idaapi.PLUGIN_OK

    def run(self, arg):
        print('hints: run')
        do_hook_to_notification_point(self.dll, HOOK_TYPES.HT_UI, self.cb, ctypes.byref(self.ctx))

    def term(self):
        print('hints: term')
        do_unhook_from_notification_point(self.dll, HOOK_TYPES.HT_UI, self.cb, ctypes.byref(self.ctx))


def PLUGIN_ENTRY():
    return DynHintsPlugin()


repository: interruptlabs/heimdallr-ida
file: heimdallr.py
description: A IDA plugin to enable linking to locations in an IDB with a disas:// URI
created_at: 2023-01-16 11:26:59+00:00
pushed_at: 2023-10-18 12:26:55+00:00
url:
https://github.com/interruptlabs/heimdallr-ida/blob/d0b86e960b37cb9b2ba998575c1a4e68a0462e63/heimdallr.py

content:
import heimdallr_utils.plugin
import platform, os, sys, json

from pathlib import Path
import ida_idaapi


default_config = """
{
  "ida_location": "",
  "idb_path": [
    ""
  ],
"heimdallr_client" : "heimdallr_client"
}
"""

def PLUGIN_ENTRY():
    """Entry point searched for by IDA"""
    ida_idaapi.require("heimdallr_utils.plugin")
    current_heimdallr_instance = heimdallr_utils.plugin.heimdallrRPC()
    return current_heimdallr_instance

def reload():
    """Reloads plugin logic - for testing.
    Doesn't appear to reload any part of the state maintained in the plugin, but can reload helper functions.
    """
    ida_idaapi.require("heimdallr_utils.plugin")


def macos_get_app_path(executable_path : Path):
    """Recursively looks for the root application folder for IDA under MacOS
    """
    for parent in executable_path.parents:
        if ".app" in parent.name:
            return parent
    raise RuntimeError("Could not resolve MacOS Application Path")

def install():
    """Carries out initial installation of the IDA plugin. Creates a symlink to the user specific plugin directory
    and creates the Heimdallr configuration directory.
    """
    completed = False
    if platform.system() == "Windows":
        idauser_path = Path(os.path.expandvars("%APPDATA%/Hex-Rays/IDA Pro/"))
        heimdallr_path = Path(os.path.expandvars("%APPDATA%/heimdallr/"))
    else:
        idauser_path = Path(os.path.expandvars("$HOME/.idapro/"))
        heimdallr_path = Path(os.path.expandvars("$HOME/.config/heimdallr/"))
    
    plugin_path = idauser_path / "plugins"
    plugin_path.mkdir(parents=True, exist_ok=True)
    installee = Path(__file__)
    install_path = plugin_path / installee.name
    if install_path.exists() and installee.is_symlink():
        installee.unlink()
    if not install_path.exists():
        os.symlink(installee, install_path)
        print(f"Symlinked plugin to {install_path}")
        completed = True
    elif installee.is_relative_to(plugin_path):
        completed = True
    else:
        print(f"Error: Existing plugin exists at {install_path}")
    
    config_path = heimdallr_path / "settings.json"
    heimdallr_path.mkdir(parents=True, exist_ok=True)
    if not config_path.exists():
        config = json.loads(default_config)
        ida_path = Path(sys.executable)
        if platform.system() == "Darwin":
            ida_path = macos_get_app_path(ida_path)

        config['ida_location'] = str(ida_path)
        with open(config_path, "w") as fd:
            json.dump(config, fd)
        print(f"Wrote settings file to {config_path}")
    else:
        print(f"Skipped settings file genereation - already exists @ {config_path}")
    if completed:
        print("Installation completed! Restart IDA to activate plugin.")

def uninstall():
    """Uninstalls the IDA Plugin by removing the symlink in the IDA Plugin directory. Maintains configuration file but
    gives instructions to uninstall manually
    """
    completed = True
    if platform.system() == "Windows":
        idauser_path = Path(os.path.expandvars("%APPDATA%/Hex-Rays/IDA Pro/"))
        heimdallr_path = Path(os.path.expandvars("%APPDATA%/heimdallr/"))
    else:
        idauser_path = Path(os.path.expandvars("$HOME/.idapro/"))
        heimdallr_path = Path(os.path.expandvars("$HOME/.config/heimdallr/"))
    
    plugin_path = idauser_path / "plugins"
    installee = Path(__file__)
    install_path = plugin_path / installee.name
    if install_path.exists() and installee.is_symlink():
        os.remove(install_path)
        print(f"Removed symlink at {install_path}")
        completed = True
    elif installee.is_relative_to(plugin_path):
        print(f"Did not remove plugin files as they are directly installed at {install_path}")
    
    if completed:
        print(f"""Uninstall can be completed by:
        - Removing the configuration directory at {heimdallr_path}
        - Removing the package with pip `pip3 remove heimdallr-ida`
        """)


repository: fox-it/mkYARA
file: mkyara/ida/mkyara_plugin.py
description: Generating YARA rules based on binary code
wanted_name: mkYARA
comment: mkYARA Plugin
created_at: 2019-03-28 13:07:36+00:00
pushed_at: 2021-10-07 08:45:57+00:00
url:
https://github.com/fox-it/mkYARA/blob/8147f91809db815687e55ddcdc88d18538940704/mkyara/ida/mkyara_plugin.py

content:
from __future__ import print_function
import idaapi
import idc
import idautils

from mkyara import (
    YaraGenerator,
)
from capstone import (
    CS_ARCH_X86,
    CS_MODE_16,
    CS_MODE_32,
    CS_MODE_64
)
from PyQt5 import QtGui, QtWidgets
from PyQt5.QtCore import Qt

INSTRUCTION_SET_MAPPING = {
    'metapc': CS_ARCH_X86,
}


def get_input_file_hash():
    return idautils.GetInputFileMD5()


def get_selection():
    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)
    return start, end


def get_inf_structure_bitness(info):
    bits = 16
    if info.is_64bit():
        bits = 64
    elif info.is_32bit():
        bits = 32
    return bits


def get_arch_info():
    info = idaapi.get_inf_structure()
    proc = info.procName.lower()
    bits = get_inf_structure_bitness(info)
    instruction_set = None
    instruction_mode = None

    if proc == 'metapc':
        instruction_set = CS_ARCH_X86
        if bits == 16:
            instruction_mode = CS_MODE_16
        elif bits == 32:
            instruction_mode = CS_MODE_32
        elif bits == 64:
            instruction_mode = CS_MODE_64
    return instruction_set, instruction_mode


class YaraRuleDialog(QtWidgets.QDialog):
    def __init__(self, parent, start_addr, end_addr, yara_rule):
        super(YaraRuleDialog, self).__init__(parent)
        self.start_addr = start_addr
        self.end_addr = end_addr
        self.yara_rule = yara_rule
        self.populate_form()

    def populate_form(self):
        self.setWindowTitle('mkYARA :: Generated Yara Rule')
        self.resize(800, 600)
        self.layout = QtWidgets.QVBoxLayout(self)
        self.top_layout = QtWidgets.QHBoxLayout()
        self.bottom_layout = QtWidgets.QHBoxLayout()
        self.bottom_layout.setAlignment(Qt.AlignRight | Qt.AlignBottom)
        # layout.addStretch()

        self.layout.addWidget(QtWidgets.QLabel("Generated Yara rule from 0x{:x} to 0x{:x}".format(self.start_addr, self.end_addr)))
        self.text_edit = QtWidgets.QTextEdit()
        font = QtGui.QFont()
        font.setFamily("Consolas")
        font.setStyleHint(QtGui.QFont.Monospace)
        font.setFixedPitch(True)
        font.setPointSize(10)
        self.text_edit.setFont(font)
        metrics = QtGui.QFontMetrics(font)
        self.text_edit.setTabStopWidth(4 * metrics.width(' '))

        self.text_edit.insertPlainText(self.yara_rule)
        self.layout.addWidget(self.text_edit)

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

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

    def ok_btn_clicked(self):
        self.close()


class mkYARAPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_FIX
    comment = "mkYARA Plugin"
    help = "mkYARA Plugin"
    wanted_name = "mkYARA"
    wanted_hotkey = ""
    dialog = None

    def init(self):
        loose_yara_action = idaapi.action_desc_t(
            'mkYARA:generate_loose_yara',   # The action name. This acts like an ID and must be unique
            'Generate Loose Yara Rule ',  # The action text.
            generic_handler(lambda: self.generate_yara_rule("loose")),   # The action handler.
            None,      # Optional: the action shortcut
            'Generate loose yara rule',  # Optional: the action tooltip (available in menus/toolbar)
            199  # Optional: the action icon (shows when in menus/toolbars)
        )

        normal_yara_action = idaapi.action_desc_t(
            'mkYARA:generate_normal_yara',   # The action name. This acts like an ID and must be unique
            'Generate Normal Yara Rule ',  # The action text.
            generic_handler(lambda: self.generate_yara_rule("normal")),   # The action handler.
            'Ctrl+Y',      # Optional: the action shortcut
            'Generate normal yara rule',  # Optional: the action tooltip (available in menus/toolbar)
            199  # Optional: the action icon (shows when in menus/toolbars)
        )

        strict_yara_action = idaapi.action_desc_t(
            'mkYARA:generate_strict_yara',   # The action name. This acts like an ID and must be unique
            'Generate Strict Yara Rule ',  # The action text.
            generic_handler(lambda: self.generate_yara_rule("strict")),   # The action handler.
            None,      # Optional: the action shortcut
            'Generate strict yara rule',  # Optional: the action tooltip (available in menus/toolbar)
            199  # Optional: the action icon (shows when in menus/toolbars)
        )

        data_yara_action = idaapi.action_desc_t(
            'mkYARA:generate_data_yara',   # The action name. This acts like an ID and must be unique
            'Generate Data Yara Rule ',  # The action text.
            generic_handler(lambda: self.generate_yara_rule("normal", is_data=True)),   # The action handler.
            None,      # Optional: the action shortcut
            'Generate data yara rule',  # Optional: the action tooltip (available in menus/toolbar)
            199  # Optional: the action icon (shows when in menus/toolbars)
        )

        idaapi.register_action(loose_yara_action)
        idaapi.register_action(normal_yara_action)
        idaapi.register_action(strict_yara_action)
        idaapi.register_action(data_yara_action)
        self.ui_hooks = mkYARAUIHooks()
        self.ui_hooks.hook()
        print('mkYARA :: Plugin Started')
        return idaapi.PLUGIN_KEEP

    def generate_yara_rule(self, mode, is_data=False):
        start, end = get_selection()
        size = end - start
        data = idaapi.get_bytes(start, size)
        ins_set, ins_mode = get_arch_info()
        yr_gen = YaraGenerator(mode, ins_set, ins_mode)
        yr_gen.add_chunk(data, offset=start, is_data=is_data)
        rule_obj = yr_gen.generate_rule()
        file_hash = get_input_file_hash()
        rule_obj.metas["hash"] = "\"{}\"".format(file_hash)
        rule = rule_obj.get_rule_string()
        self.dialog = YaraRuleDialog(None, start, end, rule)
        self.dialog.show()

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

    def run(self, arg):
        pass


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

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

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


class mkYARAUIHooks(idaapi.UI_Hooks):
    def populating_widget_popup(self, form, popup):
        pass

    def finish_populating_widget_popup(self, form, popup):
        idaapi.attach_action_to_popup(form, popup, "mkYARA:generate_loose_yara", "mkYARA/")
        idaapi.attach_action_to_popup(form, popup, "mkYARA:generate_normal_yara", "mkYARA/")
        idaapi.attach_action_to_popup(form, popup, "mkYARA:generate_strict_yara", "mkYARA/")
        idaapi.attach_action_to_popup(form, popup, "mkYARA:generate_data_yara", "mkYARA/")


plugin = mkYARAPlugin()
def PLUGIN_ENTRY():
    global plugin
    return plugin


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

content:
#!/usr/bin/env python
# Jay Smith
# jay.smith@mandiant.com
# 
########################################################################
# Copyright 2014 Mandiant/FireEye
# Copyright 2019 FireEye
#
# Mandiant/Fireye licenses this file to you under the Apache License, Version
# 2.0 (the "License"); you may not use this file except in compliance with the
# License. You may obtain a copy of the License at:
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied. See the License for the specific language governing
# permissions and limitations under the License.
########################################################################
#
# Mostly a glorified wrapper around the apply_callee_tinfo() idasdk function.
# Useful for when IDA doesn't apply stack analysis to an indirect call,
# and you can identify the function prototype during reverse engineering.
#
########################################################################

import sys

import idc 
import idautils  
import idaapi

idaapi.require('flare')
idaapi.require('flare.apply_callee_type')
idaapi.require('flare.jayutils')

PLUGIN_HELP = "This is help"
PLUGIN_NAME = "ApplyCalleeType"
PREFERRED_SHORTCUT = "Alt-J"
PLUGIN_COMMENT = "Apply callee type to indirect call location"
ACTION_NAME = 'flare:apply_callee_type'
MENU_PATH = "Edit/Operand type/Manual"

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

ex_addmenu_item_ctx = None 

def installMenuIda7():
    class ApplyCalleeHandler(idaapi.action_handler_t):
        def activate(self, ctx):
            doApplyCallee()
            return 1

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

    ret = idaapi.register_action(idaapi.action_desc_t(
            ACTION_NAME,            # Name. Acts as an ID. Must be unique.
            PLUGIN_NAME,            # Label. That's what users see.
            ApplyCalleeHandler(),   # Handler. Called when activated, and for updating
            PREFERRED_SHORTCUT,     # Shortcut (optional)
            PLUGIN_COMMENT          # Tooltip (optional)
            ))
    if not ret:
        print('Failed to register action. Bailing out')
        return
    # Insert the action in the menu
    if idaapi.attach_action_to_menu(MENU_PATH, ACTION_NAME, idaapi.SETMENU_APP):
        print("Attached to menu.")
    else:
        print("Failed attaching to menu.")

    setattr(sys.modules['idaapi'], '_apply_callee_type_plugin_installFlag', True)

def installMenu():
    #hack -> stashing a flag under idaapi to prevent multiple menu items from appearing
    if hasattr(sys.modules['idaapi'], '_apply_callee_type_plugin_installFlag'):
        #print('Skipping menu install: already present')
        return
    if using_ida7api:
        return installMenuIda7()
    global ex_addmenu_item_ctx
    ex_addmenu_item_ctx = idaapi.add_menu_item(
        MENU_PATH, 
        PLUGIN_NAME, 
        PREFERRED_SHORTCUT, 
        0, 
        doApplyCallee, 
        tuple("hello world")
    )
    if ex_addmenu_item_ctx  is None:
        print('Failed to init apply_callee_type_plugin')

    setattr(sys.modules['idaapi'], '_apply_callee_type_plugin_installFlag', True)


class apply_callee_type_plugin_t(idaapi.plugin_t):
    flags = 0
    comment = PLUGIN_COMMENT
    help = PLUGIN_HELP
    wanted_name = PLUGIN_NAME
    wanted_hotkey = ""


    def init(self):
        idaapi.msg('apply_callee_type_plugin:init\n')

        installMenu()
        return idaapi.PLUGIN_OK

    def run(self, arg):
        #idaapi.msg('apply_callee_type_plugin:run\n')
        flare.apply_callee_type.main()

    def term(self):
        #idaapi.msg('apply_callee_type_plugin:term\n')
        #if self.ex_addmenu_item_ctx is not None:
        #    idaapi.del_menu_item(ex_addmenu_item_ctx)
        pass

def PLUGIN_ENTRY():
    try:
        return apply_callee_type_plugin_t()
    except Exception as err:
        import traceback
        msg("Error: %s\n%s" % (str(err), traceback.format_exc()))
        raise

def doApplyCallee(*args):
    #idaapi.msg('doApplyCallee:Calling now\n')
    flare.apply_callee_type.main()


repository: patois/NECromancer
file: NECromancer.py
description: IDA Pro V850 Processor Module Extension
url: https://github.com/patois/NECromancer/blob/2838214fed2746191d129c9a3e7eef900cd25590/NECromancer.py
wanted_name: NECromancer
created_at: 2017-08-20 12:25:42+00:00
pushed_at: 2018-05-08 16:06:18+00:00
content:
from ida_lines import COLOR_INSN, COLOR_MACRO 
from ida_idp import CUSTOM_INSN_ITYPE, IDP_Hooks, ph_get_regnames, ph_get_id, PLFM_NEC_V850X
from ida_bytes import get_bytes
from ida_idaapi import plugin_t, PLUGIN_PROC, PLUGIN_HIDE, PLUGIN_SKIP, PLUGIN_KEEP
from ida_ua import o_displ, o_reg, o_imm, dt_dword, OOF_ADDR
from struct import unpack

###############################################################################
#
#
#   _____ _____ _____                                   
#  |   | |   __|     |___ ___ _____ ___ ___ ___ ___ ___ 
#  | | | |   __|   --|  _| . |     | .'|   |  _| -_|  _|
#  |_|___|_____|_____|_| |___|_|_|_|__,|_|_|___|___|_|  
# ____________________________________________________________________________
#
# NECromancer - a NEC V850X instruction extender plugin for IDA Pro
# -----------------------------------------------------------------
#
# This plugin extends the V850E1 IDA processor module by adding support
# for certain V850E2 instructions on a per need basis. Rather than modifying
# the source code of the V850E1 IDA processor module, this script has been
# developed as an exercise in writing a processor module extension in
# IDAPython, particularly for version 7 of IDA and onwards.
#
###############################################################################

# ----------------------------------------------------------------------------
#
# necromancer: noun | nec*ro*man*cer | a supposed practice of magic involving
# communication with the diseased
#
# history and changelog:
# ----------------------
# 2017.01.31 - initial version
#              support for "divq", "divqu", shl (reg1, reg2, reg3),
#              shr (reg1, reg2, reg3) and "sar" instructions
# 2017.02.01 - support for LD.HU (disp23[reg1], reg3), "feret" and "eiret"
#              instructions
# 2017.02.02 - support for ST.H (reg3, disp23[reg1]) instruction,
#              bugfixes, cleanup
# 2017.02.03 - support for sign extending 23bit displacement values,
#              "sch1l", "sch1r", "caxi" and "fetrap" instructions
# 2017.08.20 - IDA 7 compatibility
# 2017.09.03 - Full IDA 7 compatibility (not requiring compatibility layer)
# 2017.12.03 - Bugfixes (with thanks to https://github.com/Quorth)
# 2018.05.08 - Fixed decoding of fetrap instruction
#
#
# based on V850E2S User's Manual: Architecture, available at:
# https://www.renesas.com/en-eu/doc/products/mpumcu/doc/v850/r01us0037ej0100_v850e2.pdf
#
# ------------------------------------------------------------------------------

__author__ = "Dennis Elser"

DEBUG_PLUGIN = True

NEWINSN_COLOR = COLOR_MACRO if DEBUG_PLUGIN else COLOR_INSN

# from V850 processor module
N850F_USEBRACKETS = 0x01
N850F_OUTSIGNED = 0x02


class NewInstructions:
    (NN_divq,
    NN_divqu,
    NN_sar,
    NN_shl,
    NN_shr,
    NN_feret,
    NN_eiret,
    NN_ld_hu,
    NN_st_h,
    NN_sch1l,
    NN_sch1r,
    NN_caxi,
    NN_fetrap) = range(CUSTOM_INSN_ITYPE, CUSTOM_INSN_ITYPE+13)
    
    lst = {NN_divq:"divq",
           NN_divqu:"divqu",
           NN_sar:"sar",
           NN_shl:"shl",
           NN_shr:"shr",
           NN_feret:"feret",
           NN_eiret:"eiret",
           NN_ld_hu:"ld.hu",
           NN_st_h:"st.h",
           NN_sch1l:"sch1l",
           NN_sch1r:"sch1r",
           NN_caxi:"caxi",
           NN_fetrap:"fetrap"}


#--------------------------------------------------------------------------
class v850_idp_hook_t(IDP_Hooks):
    def __init__(self):
        IDP_Hooks.__init__(self)

    def parse_r1(self, w):
        return w & 0x1F

    def parse_r2(self, w):
        return (w & 0xF800) >> 11

    def parse_r3(self, w):
        return self.parse_r2(w)

    def sign_extend(self, disp, nbits):
        val = disp
        if val & (1 << (nbits-1)):
            val |= ~((1 << nbits)-1)
        return val

    def decode_instruction(self, insn):
        buf = get_bytes(insn.ea, 2)
        hw1 = unpack("<H", buf)[0]

        op = (hw1 & 0x7E0) >> 5 # take bit5->bit10

        # Format I
        if op == 2 and (hw1 >> 11) == 0 and (hw1 & 0x1F) != 0:
            # TODO add vector4 parsing
            insn.itype = NewInstructions.NN_fetrap
            insn.size = 2
            return True

        # Format XIV
        elif op == 0x3D and ((hw1 & 0xFFE0) >> 5) == 0x3D:
            buf = get_bytes(insn.ea+2, 2)
            hw2 = unpack("<H", buf)[0]
            subop = hw2 & 0x1F
            
            if subop == 0x07: # ld.hu
                insn.itype = NewInstructions.NN_ld_hu

                insn.Op1.type = o_displ                
                insn.Op2.type = o_reg
                

                insn.Op1.specflag1 = N850F_USEBRACKETS | N850F_OUTSIGNED
                insn.Op1.reg = self.parse_r1(hw1)

                buf = get_bytes(insn.ea+4, 2)
                hw3 = unpack("<H", buf)[0]
                
                insn.Op1.addr = self.sign_extend(((hw3 << 6) | ((hw2 & 0x7E0) >> 5)) << 1, 23)
                insn.Op1.dtyp = dt_dword
                insn.Op2.reg = self.parse_r2(hw2)
                insn.Op2.dtyp = dt_dword
                
                insn.size = 6
                return True

            elif subop == 0xD:  # st.h
                insn.itype = NewInstructions.NN_st_h

                insn.Op1.type = o_reg
                insn.Op2.type = o_displ                
                
                insn.Op2.specflag1 = N850F_USEBRACKETS | N850F_OUTSIGNED
                insn.Op2.reg = self.parse_r1(hw1)

                buf = get_bytes(insn.ea+4, 2)
                hw3 = unpack("<H", buf)[0]
                
                insn.Op2.addr = self.sign_extend(((hw3 << 6) | ((hw2 & 0x7E0) >> 5)) << 1, 23)
                insn.Op2.dtyp = dt_dword
                insn.Op1.reg = self.parse_r2(hw2)
                insn.Op1.dtyp = dt_dword
                
                insn.size = 6
                return True

        # Format II
        elif op == 0x15: # sar imm5, reg2
            insn.itype = NewInstructions.NN_sar                    

            insn.Op1.type = o_imm
            insn.Op2.type = o_reg

            insn.Op1.value = hw1 & 0x1F
            insn.Op2.reg = self.parse_r2(hw1)

            insn.size = 2
            return True


        # Format IX, X, XI
        elif op == 0x3F:
            buf = get_bytes(insn.ea+2, 2)
            hw2 = unpack("<H", buf)[0]
            subop = hw2 & 0x7FF


            if hw1 & 0x7FF == 0x7E0:
                if hw1 == 0x7E0:
                    if hw2 == 0x14A: # feret
                        insn.itype = NewInstructions.NN_feret
                        insn.size = 4
                        return insn.size
                    elif hw2 == 0x0148: # eiret
                        insn.itype = NewInstructions.NN_eiret
                        insn.size = 4
                        return True

                elif subop == 0x366: # sch1l reg2, reg3
                    insn.itype = NewInstructions.NN_sch1l

                    insn.Op1.type = o_reg
                    insn.Op2.type = o_reg

                    insn.Op1.reg = self.parse_r2(hw1)
                    insn.Op2.reg = self.parse_r3(hw2)

                    insn.size = 4
                    return True
                
                elif subop == 0x362: # sch1r reg2, reg3    
                    insn.itype = NewInstructions.NN_sch1r

                    insn.Op1.type = o_reg
                    insn.Op2.type = o_reg

                    insn.Op1.reg = self.parse_r2(hw1)
                    insn.Op2.reg = self.parse_r3(hw2)

                    insn.size = 4
                    return True

            insn_handled = False
            
            if subop == hw2 == 0xA0: # sar reg1, reg2
                insn.itype = NewInstructions.NN_sar                    

                insn.Op1.type = o_reg
                insn.Op2.type = o_reg

                insn.Op1.reg = self.parse_r1(hw1)
                insn.Op2.reg = self.parse_r2(hw1)

                insn.size = 4
                return True

            elif subop == 0xEE: # caxi [reg1], reg2, reg3
                insn.itype = NewInstructions.NN_caxi                    

                insn.Op1.type = o_displ
                insn.Op1.addr = 0
                insn.Op2.type = o_reg
                insn.Op3.type = o_reg

                insn.Op1.reg = self.parse_r1(hw1)
                insn.Op2.reg = self.parse_r2(hw1)
                insn.Op2.reg = self.parse_r3(hw2)

                insn.size = 4
                return True
                
                
            elif subop == 0x2FC:  # divq reg1, reg2, reg3
                insn.itype = NewInstructions.NN_divq
                insn.size = 4
                insn_handled = True

            elif subop == 0x2FE: # divqu reg1, reg2, reg3
                insn.itype = NewInstructions.NN_divqu
                insn.size = 4
                insn_handled = True

            elif subop == 0xA2: # sar reg1, reg2, reg3
                insn.itype = NewInstructions.NN_sar                    
                insn.size = 4
                insn_handled = True
 
            elif subop == 0xC2: # shl reg1, reg2, reg3
                insn.itype = NewInstructions.NN_shl                    
                insn.size = 4
                insn_handled = True

            elif subop == 0x82: # shr reg1, reg2, reg3
                insn.itype = NewInstructions.NN_shr                   
                insn.size = 4
                insn_handled = True

            if insn_handled:
                insn.Op1.type = o_reg
                insn.Op2.type = o_reg
                insn.Op3.type = o_reg

                insn.Op1.reg = self.parse_r1(hw1)
                insn.Op2.reg = self.parse_r2(hw1)
                insn.Op3.reg = self.parse_r3(hw2)
                return True
            
        return False

    def ev_ana_insn(self, insn):
        if insn.ea & 1:
            return False

        return self.decode_instruction(insn)

    def ev_out_mnem(self, outctx):
        insntype = outctx.insn.itype
        global NEWINSN_COLOR

        if (insntype >= CUSTOM_INSN_ITYPE) and (insntype in NewInstructions.lst):
            mnem = NewInstructions.lst[insntype]
            outctx.out_tagon(NEWINSN_COLOR)
            outctx.out_line(mnem)
            outctx.out_tagoff(NEWINSN_COLOR)

            # TODO: how can MNEM_width be determined programatically?
            MNEM_WIDTH = 8
            width = max(1, MNEM_WIDTH - len(mnem))
            outctx.out_line(' ' * width)

            return True
        return False

    def ev_emu_insn(self, insn):
        if insn.itype in [NewInstructions.NN_eiret, NewInstructions.NN_feret]:
            return True
        return False

    def ev_out_operand(self, outctx, op):
        insn = outctx.insn
        if insn.itype in [NewInstructions.NN_ld_hu, NewInstructions.NN_st_h]:
            if op.type == o_displ:
                
                outctx.out_value(op, OOF_ADDR)
                brackets = insn.ops[op.n].specflag1 & N850F_USEBRACKETS
                if brackets:
                    outctx.out_symbol('[')
                outctx.out_register(ph_get_regnames()[op.reg])
                if brackets:
                    outctx.out_symbol(']')
                return True
        return False


#--------------------------------------------------------------------------
class NECromancer_t(plugin_t):
    flags = PLUGIN_PROC | PLUGIN_HIDE
    comment = ""
    wanted_hotkey = ""
    help = "Adds support for additional V850X instructions"
    wanted_name = "NECromancer"

    def __init__(self):
        self.prochook = None

    def init(self):
        if ph_get_id() != PLFM_NEC_V850X:
            return PLUGIN_SKIP

        self.prochook = v850_idp_hook_t()
        self.prochook.hook()
        print "%s intialized." % NECromancer_t.wanted_name
        return PLUGIN_KEEP

    def run(self, arg):
        pass

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

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


repository: danigargu/syms2elf
file: syms2elf.py
description: A plugin for Hex-Ray's IDA Pro and radare2 to export the symbols recognized to the ELF symbol table
url: https://github.com/danigargu/syms2elf/blob/329c2ce160ce15b28d9f7da99dd6f8403cab18a6/syms2elf.py
created_at: 2016-03-03 17:19:22+00:00
pushed_at: 2022-08-17 13:22:52+00:00
content:
#!/usr/bin/python
# -*- coding: utf-8 -*-
#
#  sym2elf
#  28/11/2014
#
#  Daniel García <danigargu [at] gmail.com>
#  Jesús Olmos   <sha0 [at] badchecksum.net>
#
#  ELF parser - based on the lib of ROPgadget tool by Jonathan Salwan
#  http://shell-storm.org/project/ROPgadget/
# 
#  This program is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software  Foundation, either  version 3 of  the License, or
#  (at your option) any later version.
#

PLUG_NAME    = "syms2elf"

import os
import sys

from ctypes import *
from struct import unpack

USE_R2  = False
USE_IDA = False

try:
    import idaapi
    USE_IDA = True
except:
    if "radare2" in os.environ.get('PATH',''):
        USE_R2  = True
    else:
        print("ERROR: The plugin must be run in IDA or radare2")
        sys.exit(0)

SHN_UNDEF = 0
STB_GLOBAL_FUNC = 0x12

class SHTypes:
    SHT_NULL      = 0
    SHT_PROGBITS  = 1
    SHT_SYMTAB    = 2
    SHT_STRTAB    = 3
    SHT_RELA      = 4
    SHT_HASH      = 5
    SHT_DYNAMIC   = 6
    SHT_NOTE      = 7
    SHT_NOBITS    = 8
    SHT_REL       = 9
    SHT_SHLIB     = 10
    SHT_DYNSYM    = 11
    SHT_NUM       = 12
    SHT_LOPROC    = 0x70000000
    SHT_HIPROC    = 0x7fffffff
    SHT_LOUSER    = 0x80000000
    SHT_HIUSER    = 0xffffffff

class ELFFlags:
    ELFCLASS32  = 0x01
    ELFCLASS64  = 0x02
    EI_CLASS    = 0x04
    EI_DATA     = 0x05
    ELFDATA2LSB = 0x01
    ELFDATA2MSB = 0x02
    EM_386      = 0x03
    EM_X86_64   = 0x3e
    EM_ARM      = 0x28
    EM_MIPS     = 0x08
    EM_SPARCv8p = 0x12
    EM_PowerPC  = 0x14
    EM_ARM64    = 0xb7

class SymFlags:
    STB_LOCAL   = 0
    STB_GLOBAL  = 1
    STB_WEAK    = 2
    STT_NOTYPE  = 0
    STT_OBJECT  = 1
    STT_FUNC    = 2
    STT_SECTION = 3
    STT_FILE    = 4
    STT_COMMON  = 5
    STT_TLS     = 6

class Elf32_Ehdr_LSB(LittleEndianStructure):
    _fields_ =  [
                    ("e_ident",         c_ubyte * 16),
                    ("e_type",          c_ushort),
                    ("e_machine",       c_ushort),
                    ("e_version",       c_uint),
                    ("e_entry",         c_uint),
                    ("e_phoff",         c_uint),
                    ("e_shoff",         c_uint),
                    ("e_flags",         c_uint),
                    ("e_ehsize",        c_ushort),
                    ("e_phentsize",     c_ushort),
                    ("e_phnum",         c_ushort),
                    ("e_shentsize",     c_ushort),
                    ("e_shnum",         c_ushort),
                    ("e_shstrndx",      c_ushort)
                ]
 
class Elf64_Ehdr_LSB(LittleEndianStructure):
    _fields_ =  [
                    ("e_ident",         c_ubyte * 16),
                    ("e_type",          c_ushort),
                    ("e_machine",       c_ushort),
                    ("e_version",       c_uint),
                    ("e_entry",         c_ulonglong),
                    ("e_phoff",         c_ulonglong),
                    ("e_shoff",         c_ulonglong),
                    ("e_flags",         c_uint),
                    ("e_ehsize",        c_ushort),
                    ("e_phentsize",     c_ushort),
                    ("e_phnum",         c_ushort),
                    ("e_shentsize",     c_ushort),
                    ("e_shnum",         c_ushort),
                    ("e_shstrndx",      c_ushort)
                ]

class Elf32_Phdr_LSB(LittleEndianStructure):
    _fields_ =  [
                    ("p_type",          c_uint),
                    ("p_offset",        c_uint),
                    ("p_vaddr",         c_uint),
                    ("p_paddr",         c_uint),
                    ("p_filesz",        c_uint),
                    ("p_memsz",         c_uint),
                    ("p_flags",         c_uint),
                    ("p_align",         c_uint)
                ]

class Elf64_Phdr_LSB(LittleEndianStructure):
    _fields_ =  [
                    ("p_type",          c_uint),
                    ("p_flags",         c_uint),
                    ("p_offset",        c_ulonglong),
                    ("p_vaddr",         c_ulonglong),
                    ("p_paddr",         c_ulonglong),
                    ("p_filesz",        c_ulonglong),
                    ("p_memsz",         c_ulonglong),
                    ("p_align",         c_ulonglong)
                ]

class Elf32_Shdr_LSB(LittleEndianStructure):
    _fields_ =  [
                    ("sh_name",         c_uint),
                    ("sh_type",         c_uint),
                    ("sh_flags",        c_uint),
                    ("sh_addr",         c_uint),
                    ("sh_offset",       c_uint),
                    ("sh_size",         c_uint),
                    ("sh_link",         c_uint),
                    ("sh_info",         c_uint),
                    ("sh_addralign",    c_uint),
                    ("sh_entsize",      c_uint)
                ]

class Elf64_Shdr_LSB(LittleEndianStructure):
    _fields_ =  [
                    ("sh_name",         c_uint),
                    ("sh_type",         c_uint),
                    ("sh_flags",        c_ulonglong),
                    ("sh_addr",         c_ulonglong),
                    ("sh_offset",       c_ulonglong),
                    ("sh_size",         c_ulonglong),
                    ("sh_link",         c_uint),
                    ("sh_info",         c_uint),
                    ("sh_addralign",    c_ulonglong),
                    ("sh_entsize",      c_ulonglong)
                ]

class Elf32_Ehdr_MSB(BigEndianStructure):
    _fields_ =  [
                    ("e_ident",         c_ubyte * 16),
                    ("e_type",          c_ushort),
                    ("e_machine",       c_ushort),
                    ("e_version",       c_uint),
                    ("e_entry",         c_uint),
                    ("e_phoff",         c_uint),
                    ("e_shoff",         c_uint),
                    ("e_flags",         c_uint),
                    ("e_ehsize",        c_ushort),
                    ("e_phentsize",     c_ushort),
                    ("e_phnum",         c_ushort),
                    ("e_shentsize",     c_ushort),
                    ("e_shnum",         c_ushort),
                    ("e_shstrndx",      c_ushort)
                ]
 
class Elf64_Ehdr_MSB(BigEndianStructure):
    _fields_ =  [
                    ("e_ident",         c_ubyte * 16),
                    ("e_type",          c_ushort),
                    ("e_machine",       c_ushort),
                    ("e_version",       c_uint),
                    ("e_entry",         c_ulonglong),
                    ("e_phoff",         c_ulonglong),
                    ("e_shoff",         c_ulonglong),
                    ("e_flags",         c_uint),
                    ("e_ehsize",        c_ushort),
                    ("e_phentsize",     c_ushort),
                    ("e_phnum",         c_ushort),
                    ("e_shentsize",     c_ushort),
                    ("e_shnum",         c_ushort),
                    ("e_shstrndx",      c_ushort)
                ]

class Elf32_Phdr_MSB(BigEndianStructure):
    _fields_ =  [
                    ("p_type",          c_uint),
                    ("p_offset",        c_uint),
                    ("p_vaddr",         c_uint),
                    ("p_paddr",         c_uint),
                    ("p_filesz",        c_uint),
                    ("p_memsz",         c_uint),
                    ("p_flags",         c_uint),
                    ("p_align",         c_uint)
                ]

class Elf64_Phdr_MSB(BigEndianStructure):
    _fields_ =  [
                    ("p_type",          c_uint),
                    ("p_flags",         c_uint),
                    ("p_offset",        c_ulonglong),
                    ("p_vaddr",         c_ulonglong),
                    ("p_paddr",         c_ulonglong),
                    ("p_filesz",        c_ulonglong),
                    ("p_memsz",         c_ulonglong),
                    ("p_align",         c_ulonglong)
                ]

class Elf32_Shdr_MSB(BigEndianStructure):
    _fields_ =  [
                    ("sh_name",         c_uint),
                    ("sh_type",         c_uint),
                    ("sh_flags",        c_uint),
                    ("sh_addr",         c_uint),
                    ("sh_offset",       c_uint),
                    ("sh_size",         c_uint),
                    ("sh_link",         c_uint),
                    ("sh_info",         c_uint),
                    ("sh_addralign",    c_uint),
                    ("sh_entsize",      c_uint)
                ]

class Elf64_Shdr_MSB(BigEndianStructure):
    _fields_ =  [
                    ("sh_name",         c_uint),
                    ("sh_type",         c_uint),
                    ("sh_flags",        c_ulonglong),
                    ("sh_addr",         c_ulonglong),
                    ("sh_offset",       c_ulonglong),
                    ("sh_size",         c_ulonglong),
                    ("sh_link",         c_uint),
                    ("sh_info",         c_uint),
                    ("sh_addralign",    c_ulonglong),
                    ("sh_entsize",      c_ulonglong)
                ]

class Elf32_Sym_LSB(LittleEndianStructure):
    _fields_ =  [
                    ("st_name",         c_uint),
                    ("st_value",        c_uint),
                    ("st_size",         c_uint),
                    ("st_info",         c_ubyte),
                    ("st_other",        c_ubyte),
                    ("st_shndx",        c_ushort)
                ]

class Elf64_Sym_LSB(LittleEndianStructure):
    _fields_ =  [
                    ("st_name",         c_uint),
                    ("st_info",         c_ubyte),
                    ("st_other",        c_ubyte),
                    ("st_shndx",        c_ushort),
                    ("st_value",        c_ulonglong),
                    ("st_size",         c_ulonglong)
                ]

class Elf32_Sym_MSB(BigEndianStructure):
    _fields_ =  [
                    ("st_name",         c_uint),
                    ("st_value",        c_uint),
                    ("st_size",         c_uint),
                    ("st_info",         c_ubyte),
                    ("st_other",        c_ubyte),
                    ("st_shndx",        c_ushort)
                ]

class Elf64_Sym_MSB(BigEndianStructure):
    _fields_ =  [
                    ("st_name",         c_uint),
                    ("st_info",         c_ubyte),
                    ("st_other",        c_ubyte),
                    ("st_shndx",        c_ushort),
                    ("st_value",        c_ulonglong),
                    ("st_size",         c_ulonglong)
                ]


""" This class parses the ELF """
class ELF:
    def __init__(self, binary):
        self.binary    = bytearray(binary)
        self.ElfHeader = None
        self.shdr_l    = []
        self.phdr_l    = []
        self.syms_l    = []
        self.e_ident   = self.binary[:15]
        self.ei_data   = unpack("<B", self.e_ident[ELFFlags.EI_DATA:ELFFlags.EI_DATA+1])[0] # LSB/MSB
        
        self.__setHeaderElf()
        self.__setShdr()
        self.__setPhdr()

    def is_stripped(self):
        if not self.get_symtab():
            return True
        if not self.get_strtab():
            return True
        return False

    def strip_symbols(self):        
        sh2delete = 2
        size2dec  = 0
        end_shdr  = self.ElfHeader.e_shoff + (self.sizeof_sh() * self.ElfHeader.e_shnum)

        symtab = self.get_symtab()
        strtab = self.get_strtab()

        if not symtab or not strtab:
            return False

        log("Stripping binary...")

        if symtab.sh_offset < end_shdr:
            size2dec += symtab.sh_size

        if strtab.sh_offset < end_shdr:
            size2dec += strtab.sh_size

        self.ElfHeader.e_shoff -= size2dec
        self.ElfHeader.e_shnum -= sh2delete

        e_shnum = self.ElfHeader.e_shnum
        e_shoff = self.ElfHeader.e_shoff
        sz_striped = (e_shoff + (e_shnum * self.sizeof_sh()))        

        if strtab.sh_offset > symtab.sh_offset:
            self.cut_at_offset(strtab.sh_offset, strtab.sh_size)  
            self.cut_at_offset(symtab.sh_offset, symtab.sh_size)
        else:
            self.cut_at_offset(symtab.sh_offset, symtab.sh_size)
            self.cut_at_offset(strtab.sh_offset, strtab.sh_size)

        self.binary = self.binary[0:sz_striped]
        self.write(0, self.ElfHeader)
        return True

    def get_symtab(self):
        shstrtab = bytes(self.get_shstrtab_data())
        for sh in self.shdr_l:
            sh_name = shstrtab[sh.sh_name:].split(b"\0")[0]
            if  sh.sh_type == SHTypes.SHT_SYMTAB and \
                (sh.sh_name == SHN_UNDEF or sh_name == ".symtab"):
                return sh
        return None

    def get_strtab(self):
        shstrtab = bytes(self.get_shstrtab_data())
        for sh in self.shdr_l:
            sh_name = shstrtab[sh.sh_name:].split(b"\0")[0]
            if  sh.sh_type == SHTypes.SHT_STRTAB and \
                (sh.sh_name == SHN_UNDEF or sh_name == ".strtab"):
                return sh
        return None

    def getArchMode(self):
        if self.ElfHeader.e_ident[ELFFlags.EI_CLASS] == ELFFlags.ELFCLASS32: 
            return 32
        elif self.ElfHeader.e_ident[ELFFlags.EI_CLASS] == ELFFlags.ELFCLASS64: 
            return 64
        else:
            log("[Error] ELF.getArchMode() - Bad Arch size")
            return None

    """ Parse ELF header """
    def __setHeaderElf(self):
        e_ident = self.binary[:15]

        ei_class = unpack("<B", e_ident[ELFFlags.EI_CLASS:ELFFlags.EI_CLASS+1])[0]
        ei_data  = unpack("<B", e_ident[ELFFlags.EI_DATA:ELFFlags.EI_DATA+1])[0]

        if ei_class != ELFFlags.ELFCLASS32 and ei_class != ELFFlags.ELFCLASS64:
            log("[Error] ELF.__setHeaderElf() - Bad Arch size")
            return None

        if ei_data != ELFFlags.ELFDATA2LSB and ei_data != ELFFlags.ELFDATA2MSB:
            log("[Error] ELF.__setHeaderElf() - Bad architecture endian")
            return None

        if ei_class == ELFFlags.ELFCLASS32: 
            if   ei_data == ELFFlags.ELFDATA2LSB: self.ElfHeader = Elf32_Ehdr_LSB.from_buffer_copy(self.binary)
            elif ei_data == ELFFlags.ELFDATA2MSB: self.ElfHeader = Elf32_Ehdr_MSB.from_buffer_copy(self.binary)
        elif ei_class == ELFFlags.ELFCLASS64: 
            if   ei_data == ELFFlags.ELFDATA2LSB: self.ElfHeader = Elf64_Ehdr_LSB.from_buffer_copy(self.binary)
            elif ei_data == ELFFlags.ELFDATA2MSB: self.ElfHeader = Elf64_Ehdr_MSB.from_buffer_copy(self.binary)

    """ Write the section header to self.binary """
    def write_shdr(self):
        off = self.ElfHeader.e_shoff
        for sh in self.shdr_l:
            self.write(off, sh)
            off += off + sizeof(sh) 

    """ Parse Section header """
    def __setShdr(self):
        shdr_num = self.ElfHeader.e_shnum
        base = self.binary[self.ElfHeader.e_shoff:]
        shdr_l = []

        e_ident = self.binary[:15]
        ei_data = unpack("<B", e_ident[ELFFlags.EI_DATA:ELFFlags.EI_DATA+1])[0]

        for i in range(shdr_num):
            if self.getArchMode() == 32:
                if   ei_data == ELFFlags.ELFDATA2LSB: shdr = Elf32_Shdr_LSB.from_buffer_copy(base)
                elif ei_data == ELFFlags.ELFDATA2MSB: shdr = Elf32_Shdr_MSB.from_buffer_copy(base)
            elif self.getArchMode() == 64:
                if   ei_data == ELFFlags.ELFDATA2LSB: shdr = Elf64_Shdr_LSB.from_buffer_copy(base)
                elif ei_data == ELFFlags.ELFDATA2MSB: shdr = Elf64_Shdr_MSB.from_buffer_copy(base)

            self.shdr_l.append(shdr)
            base = base[self.ElfHeader.e_shentsize:]

        string_table = self.binary[(self.shdr_l[self.ElfHeader.e_shstrndx].sh_offset):]
        for i in range(shdr_num):
            self.shdr_l[i].str_name = string_table[self.shdr_l[i].sh_name:].split(b'\0')[0]

    """ Parse Program header """
    def __setPhdr(self):
        pdhr_num = self.ElfHeader.e_phnum
        base = self.binary[self.ElfHeader.e_phoff:]
        phdr_l = []

        e_ident = self.binary[:15]
        ei_data = unpack("<B", e_ident[ELFFlags.EI_DATA:ELFFlags.EI_DATA+1])[0]

        for i in range(pdhr_num):
            if self.getArchMode() == 32:
                if   ei_data == ELFFlags.ELFDATA2LSB: phdr = Elf32_Phdr_LSB.from_buffer_copy(base)
                elif ei_data == ELFFlags.ELFDATA2MSB: phdr = Elf32_Phdr_MSB.from_buffer_copy(base)
            elif self.getArchMode() == 64:
                if   ei_data == ELFFlags.ELFDATA2LSB: phdr = Elf64_Phdr_LSB.from_buffer_copy(base)
                elif ei_data == ELFFlags.ELFDATA2MSB: phdr = Elf64_Phdr_MSB.from_buffer_copy(base)

            self.phdr_l.append(phdr)
            base = base[self.ElfHeader.e_phentsize:]

    def get_section_id(self, sh_name):
        for idx, sh in enumerate(self.shdr_l):
            if sh.str_name == sh_name.encode('ascii'):
                return idx
        return None

    def get_shstrtab_data(self):
        sh = self.shdr_l[self.ElfHeader.e_shstrndx]
        if sh.sh_type == SHTypes.SHT_STRTAB:
            return self.binary[sh.sh_offset:sh.sh_offset+sh.sh_size]
        return None

    def get_sym_at_offset(self, off):
        if self.getArchMode() == 32:
            if   ei_data == ELFFlags.ELFDATA2LSB: sym = Elf32_Sym_LSB.from_buffer_copy(self.binary[off:])
            elif ei_data == ELFFlags.ELFDATA2MSB: sym = Elf32_Sym_MSB.from_buffer_copy(self.binary[off:])
        elif self.getArchMode() == 64:
            if   ei_data == ELFFlags.ELFDATA2LSB: sym = Elf64_Sym_LSB.from_buffer_copy(self.binary[off:])
            elif ei_data == ELFFlags.ELFDATA2MSB: sym = Elf64_Sym_MSB.from_buffer_copy(self.binaryGetInputFilePath[off:])
        return sym

    def get_entrypoint(self):
        return self.e_entry

    def sizeof_sh(self):
        size = None
        if self.getArchMode() == 32:
            size = sizeof(Elf32_Shdr_LSB())
        elif self.getArchMode() == 64:
            size = sizeof(Elf64_Shdr_LSB())
        return size

    def sizeof_sym(self):
        size = None
        if self.getArchMode() == 32:
            size = sizeof(Elf32_Sym_LSB)
        elif self.getArchMode() == 64:
            size = sizeof(Elf64_Sym_LSB)
        return size

    def append_section_header(self, section):
        sh = None

        if self.getArchMode() == 32:
            if   self.ei_data == ELFFlags.ELFDATA2LSB: sh = Elf32_Shdr_LSB()
            elif self.ei_data == ELFFlags.ELFDATA2MSB: sh = Elf32_Shdr_MSB()
        elif self.getArchMode() == 64:
            if   self.ei_data == ELFFlags.ELFDATA2LSB: sh = Elf64_Shdr_LSB()
            elif self.ei_data == ELFFlags.ELFDATA2MSB: sh = Elf64_Shdr_MSB()

        sh.sh_name      = section["name"]
        sh.sh_type      = section["type"]
        sh.sh_flags     = section["flags"]
        sh.sh_addr      = section["addr"]
        sh.sh_offset    = section["offset"]
        sh.sh_size      = section["size"]
        sh.sh_link      = section["link"]
        sh.sh_info      = section["info"]
        sh.sh_addralign = section["addralign"]
        sh.sh_entsize   = section["entsize"]

        self.binary.extend(sh)

    def append_symbol(self, symbol):
        if self.getArchMode() == 32:
            if   self.ei_data == ELFFlags.ELFDATA2LSB: sym = Elf32_Sym_LSB()
            elif self.ei_data == ELFFlags.ELFDATA2MSB: sym = Elf32_Sym_MSB()
        elif self.getArchMode() == 64:
            if   self.ei_data == ELFFlags.ELFDATA2LSB: sym = Elf64_Sym_LSB()
            elif self.ei_data == ELFFlags.ELFDATA2MSB: sym = Elf64_Sym_MSB()

        sym.st_name   = symbol["name"]
        sym.st_value  = symbol["value"]
        sym.st_size   = symbol["size"]
        sym.st_info   = symbol["info"]
        sym.st_other  = symbol["other"]
        sym.st_shndx  = symbol["shndx"]

        self.binary.extend(sym)

    def get_binary(self):
        return self.binary

    def write(self, offset, data):
        self.binary[offset:offset+sizeof(data)] = data

    def expand_at_offset(self, offset, data):
        self.binary = self.binary[:offset] + data + self.binary[offset:]

    def cut_at_offset(self, offset, size):
        self.binary = self.binary[:offset] + self.binary[offset+size:]

    def save(self, output):
        with open(output, 'wb') as f:
            f.write(self.binary)

class Symbol:
    def __init__(self, name, info, value, size, shname, shndx=-1):
        self.name   = name
        self.info   = info
        self.value  = value
        self.size   = size
        self.shname = shname
        self.shndx  = shndx

    def __str__(self):
        return "%s;%s;%s;%s;%s" % (self.name, self.value, self.size, 
            self.info, self.shname)

def log(msg=''):
    print("[%s] %s" % (PLUG_NAME, msg))

def log_r2(msg=''):
    print("%s" % msg)


def write_symbols(input_file, output_file, symbols):
    try:        
        with open(input_file, 'rb') as f:
            bin = ELF(f.read())

        if len(symbols) < 1:
            log("No symbols to export")
            return

        log("Exporting symbols to ELF...")
        bin.strip_symbols()

        # raw strtab
        strtab_raw = b"\x00" + b"\x00".join([sym.name.encode('ascii') for sym in symbols]) + b"\x00"

        symtab = {
            "name"      : SHN_UNDEF,
            "type"      : SHTypes.SHT_SYMTAB,
            "flags"     : 0,
            "addr"      : 0,
            "offset"    : len(bin.binary) + (bin.sizeof_sh() * (bin.ElfHeader.e_shnum + 2)),
            "size"      : (len(symbols) + 1) * bin.sizeof_sym(),
            "link"      : bin.ElfHeader.e_shnum + 1, # index of SHT_STRTAB
            "info"      : 1,
            "addralign" : 4,
            "entsize"   : bin.sizeof_sym()
        }

        off_strtab = (len(bin.binary) + (bin.sizeof_sh() * (bin.ElfHeader.e_shnum + 2)) + (bin.sizeof_sym() * (len(symbols) + 1)))

        strtab = {
            "name"      : SHN_UNDEF,
            "type"      : SHTypes.SHT_STRTAB,
            "flags"     : 0,
            "addr"      : 0,
            "offset"    : off_strtab,
            "size"      : len(strtab_raw),
            "link"      : 0,
            "info"      : 0,
            "addralign" : 1,
            "entsize"   : 0
        }

        shdrs = bin.binary[bin.ElfHeader.e_shoff:bin.ElfHeader.e_shoff + (bin.sizeof_sh() * bin.ElfHeader.e_shnum)]
        bin.ElfHeader.e_shnum += 2
        bin.ElfHeader.e_shoff = len(bin.binary)
        bin.write(0, bin.ElfHeader)
        bin.binary.extend(shdrs)

        base = bin.binary[bin.ElfHeader.e_shoff:]
        _off = bin.ElfHeader.e_shoff
        for i in range(bin.ElfHeader.e_shnum - 2):
            if bin.getArchMode() == 32:
                if   bin.ei_data == ELFFlags.ELFDATA2LSB: shdr = Elf32_Shdr_LSB.from_buffer_copy(base)
                elif bin.ei_data == ELFFlags.ELFDATA2MSB: shdr = Elf32_Shdr_MSB.from_buffer_copy(base)
            elif bin.getArchMode() == 64:
                if   bin.ei_data == ELFFlags.ELFDATA2LSB: shdr = Elf64_Shdr_LSB.from_buffer_copy(base)
                elif bin.ei_data == ELFFlags.ELFDATA2MSB: shdr = Elf64_Shdr_MSB.from_buffer_copy(base)
            base = base[bin.ElfHeader.e_shentsize:]
            bin.write(_off, shdr)
            _off += bin.sizeof_sh()
        bin.append_section_header(symtab)
        bin.append_section_header(strtab)

        # Local symbol - separator
        sym = {
            "name"  : 0,
            "value" : 0,
            "size"  : 0,
            "info"  : SymFlags.STB_LOCAL,
            "other" : 0,
            "shndx" : 0 
        }
        bin.append_symbol(sym)

        # add symbols  
        for s in symbols:

            sh_idx = bin.get_section_id(s.shname)
            if not sh_idx:
                log("ERROR: Section ID for '%s' not found" % s.shname)
                continue

            sym = {
                "name"  : strtab_raw.index(s.name.encode('ascii')),
                "value" : s.value,
                "size"  : s.size,
                "info"  : s.info,
                "other" : 0,
                "shndx" : sh_idx
            }

            #log("0x%08x - 0x%08x - %s - %d/%d - %d" % (s.value, s.size, s.name, strtab_raw.index(s.name), len(strtab_raw), s.info))
            bin.append_symbol(sym)

        # add symbol strings
        bin.binary.extend(strtab_raw)

        log("ELF saved to: %s" % output_file)
        bin.save(output_file)

    except:
        log(traceback.format_exc())

def ida_fcn_filter(func_ea):
    if idc.get_segm_name(func_ea) not in ("extern", ".plt"):
        return True
    return False 

def get_ida_symbols():
    symbols = []

    for f in filter(ida_fcn_filter, Functions()):
        func     = get_func(f)
        seg_name = idc.get_segm_name(f)

        fn_name = idc.get_func_name(f)
        symbols.append(Symbol(fn_name, STB_GLOBAL_FUNC, 
            int(func.start_ea), int(func.size()), seg_name))

    return symbols

def get_section(addr):
  for idx, s in enumerate(r2p.cmdj("iSj")):
    if s['vaddr'] <= addr < (s['vaddr'] + s['size']):
      return (idx, s['name'])
  return None

def r2_fcn_filter(fcn):
    if fcn['name'].startswith(("sym.imp", "loc.imp", "section")):
        return False
    return True

def get_r2_symbols():
    symbols = []

    for fnc in filter(r2_fcn_filter, r2p.cmdj("aflj")):
        sh_idx, sh_name = get_section(fnc['offset'])
        fnc_name = fnc['name']

        if fnc_name.startswith(("sym","fcn")):
            fnc_name = fnc_name[4:]
            if fnc_name.startswith("go."):
                fnc_name = fnc_name[3:]

        symbols.append(Symbol(fnc_name, STB_GLOBAL_FUNC, 
            fnc['offset'], int(fnc['size']), sh_name))

    return symbols

if USE_IDA:

    import idc
    import ida_nalt
    import ida_kernwin
    from idaapi import *
    from idautils import *

    class Syms2Elf(Form):
        def __init__(self):
            Form.__init__(self, r"""syms2elf
        {formChangeCb}
        <#Output file#Output ~f~ile:{txtFile}>
        """, {
            'formChangeCb' : Form.FormChangeCb(self.OnFormChange),
            'txtFile' : Form.FileInput(save=True, swidth=50)
        })

            self.input_elf = ida_nalt.get_input_file_path()

        def OnFormChange(self, fid):
            if fid == self.txtFile.id:
                o_file = self.GetControlValue(self.txtFile)

                if os.path.exists(o_file):
                    s = "Output file already exists\n" \
                        "The output file already exists. " \
                        "Do you want to overwrite it?"

                    if bool(ida_kernwin.ask_form(s, "1")):
                        os.remove(o_file)
                    else:
                        self.SetControlValue(self.txtFile, '')
            return 1

        def Show(self):
            if not self.Compiled():
                self.Compile()

            if "ELF" in get_file_type_name():
                ok = self.Execute()
            else:
                warning("The input file is not a ELF")
                ok = 0
            return ok

    class Syms2Elf_t(plugin_t):
        flags = PLUGIN_UNL
        comment = ""
        help = ""
        wanted_name = PLUG_NAME
        wanted_hotkey = ""

        def init(self):
            return PLUGIN_OK

        def run(self, arg=0):
            f = Syms2Elf()

            if f.Show():
                symbols = get_ida_symbols()
                output_file = f.txtFile.value
                write_symbols(f.input_elf, output_file, symbols)
                f.Free()

        def term(self):
            pass

    def PLUGIN_ENTRY():
        return Syms2Elf_t()

elif USE_R2:

    import r2pipe
    r2p = r2pipe.open()
    log = log_r2

    if len(sys.argv) < 2:
        log("Usage: $syms2elf <output file>")
        sys.exit(0)

    file_info = r2p.cmdj("ij").get("core")

    if file_info['format'].lower() in ('elf','elf64'):
        symbols = get_r2_symbols()  
        write_symbols(file_info['file'], sys.argv[1], symbols)
    else:
        log("The input file is not a ELF")




repository: naim94a/rebasedcomment
file: rebased_comment.py
description: Comment rebasing for IDA Pro
wanted_name: rebasedcomment
comment: Updates offsets in comments when rebased
created_at: 2020-04-05 20:18:20+00:00
pushed_at: 2020-04-06 21:24:23+00:00
url:
https://github.com/naim94a/rebasedcomment/blob/8c7ea3ae0f01565fd63eb644bfb217a413fd6050/rebased_comment.py

content:
"""
comment rebaser plugin for IDA
Copyright (C) 2020 Naim Abda <naim@abda.nl>

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 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 <https://www.gnu.org/licenses/>.

"""
__version__ = "0.1.0"

import re
import ida_idaapi
import idaapi
import ida_idp
import ida_hexrays
from ida_idaapi import plugin_t
from idaapi import BADADDR, get_cmt

RE_COMMENT = re.compile(r'(?<!\w)((0x[A-f\d]+)|([A-f\d]+h))(?!\w)')

def get_all_comments(seg=None):
    start = 0 if not seg else seg.start_ea
    end = -1 if not seg else seg.end_ea
    while True:
        res = idaapi.next_that(start, end, lambda f: f & idaapi.FF_COMM)
        if res == BADADDR:
            return
        start = res + 1
        yield res

def is_addr(segs, ea):
    for seg in segs:
        s = seg.start_ea
        e = seg.end_ea
        if e > ea >= s:
            return True
    return False

def get_segs():
    seg_count = idaapi.get_segm_qty()
    ranges = []

    for i in range(0, seg_count):
        seg = idaapi.getnseg(i)
        ranges.append(seg)

    return ranges

def rebase_comment(segs, delta, comment):
    res = RE_COMMENT.finditer(comment)
    if not res:
        return None

    spans = []

    for match in res:
        span = match.span()
        s_addr = match.string[span[0]:span[1]]
        h_at = s_addr[0] == '@'
        if h_at:
            s_addr = s_addr[1:]

        h_tail = s_addr[-1] == 'h'
        if h_tail:
            s_addr = s_addr[:-1]
        else:
            s_addr = s_addr[2:]

        addr = int(s_addr, 16)
        fixed_addr = addr + delta

        # This will happen after segments have changed...
        if not is_addr(segs, fixed_addr):
            continue

        new_str = hex(fixed_addr)[2:]
        if new_str[-1] == 'L':
            new_str = new_str[:-1]

        if h_tail:
            new_str += 'h'
        else:
            new_str = '0x' + new_str

        if h_at:
            new_str = '@' + new_str
        
        spans.append((span, new_str))

    if not len(spans):
        return None
    
    fragments = []
    pos = 0
    for (idx, (span, newval)) in enumerate(spans):
        pre = comment[:span[0]][pos:]
        pos = span[1]

        fragments.append(pre)
        fragments.append(newval)

        if idx == len(spans) - 1:
            fragments.append(comment[pos:])
    res = ''.join(fragments)
    return res

def handle_comments(delta, segs):
    for cmt_offset in get_all_comments():
        for is_repeatable in (True, False):
            cmt = idaapi.get_cmt(cmt_offset, is_repeatable)
            if not cmt:
                continue
            new_cmt = rebase_comment(segs, delta, cmt)
            if not new_cmt:
                continue
            idaapi.set_cmt(cmt_offset, new_cmt, is_repeatable)

def handle_structs(delta, segs):
    for idx in range(idaapi.get_struc_qty()):
        tid = idaapi.get_struc_by_idx(idx)
        for cmt_type in (True, False):
            cmt = idaapi.get_struc_cmt(tid, cmt_type)
            if cmt:
                new_cmt = rebase_comment(segs, delta, cmt)
                if new_cmt:
                    idaapi.set_struc_cmt(tid, new_cmt, cmt_type)
        s = idaapi.get_struc(tid)
        for midx in range(s.memqty):
            m = s.get_member(midx)
            for cmt_type in (True, False):
                cmt = idaapi.get_member_cmt(m.id, cmt_type)
                if cmt:
                    new_cmt = rebase_comment(segs, delta, cmt)
                    if new_cmt:
                        idaapi.set_member_cmt(m, new_cmt, cmt_type)


class enum_memb_visitor(idaapi.enum_member_visitor_t):
    def __init__(self, segs, delta):
        super(enum_memb_visitor, self).__init__()
        self.segs = segs
        self.delta = delta

    def visit_enum_member(self, cid, val):
        for cmt_type in (True, False):
            cmt = idaapi.get_enum_member_cmt(cid, cmt_type)
            if cmt:
                new_cmt = rebase_comment(self.segs, self.delta, cmt)
                if new_cmt:
                    idaapi.set_enum_member_cmt(cid, new_cmt, cmt_type)
        return 0

def handle_enums(delta, segs):
    for idx in range(idaapi.get_enum_qty()):
        e = idaapi.getn_enum(idx)
        for cmt_type in (True, False):
            cmt = idaapi.get_enum_cmt(e, cmt_type)
            if cmt:
                new_cmt = rebase_comment(segs, delta, cmt)
                if new_cmt:
                    idaapi.set_enum_cmt(e, new_cmt, cmt_type)
        idaapi.for_all_enum_members(e, enum_memb_visitor(segs, delta))

def get_all_funcs():
    for idx in range(idaapi.get_func_qty()):
        f = idaapi.getn_func(idx)
        yield f

def handle_function_comments(delta, segs):
    for f in get_all_funcs():
        for cmt_type in (True, False):
            cmt = idaapi.get_func_cmt(f, cmt_type)
            if cmt:
                new_cmt = rebase_comment(segs, delta, cmt)
                if new_cmt:
                    idaapi.set_func_cmt(f, cmt, cmt_type)
        cmts = idaapi.restore_user_cmts(f.start_ea)
        if not cmts:
            continue
        changed = False
        for (treeloc, citm) in cmts.items():
            if citm:
                citm = citm.c_str()
            new_cmt = rebase_comment(segs, delta, citm)
            if new_cmt:
                changed = True
                it = cmts.find(treeloc)
                cmts.erase(it)
                cmts.insert(treeloc, idaapi.citem_cmt_t(new_cmt))
        if changed:
            idaapi.save_user_cmts(f.start_ea, cmts)

def patch_comments(delta, affected_segs=None):
    segs = affected_segs if affected_segs else get_segs()

    # Process structs
    handle_structs(delta, segs)

    # Process enums
    handle_enums(delta, segs)

    # Process comments
    handle_comments(delta, segs)

    # process function comments
    handle_function_comments(delta, segs)


class Hooks(ida_idp.IDB_Hooks):
    def __init__(self):
        ida_idp.IDB_Hooks.__init__(self)
    
    def allsegs_moved(self, info):
        # TODO: check that all segments actually moved by the same delta.
        seg_info = info.at(0)
        old_base = seg_info._from
        new_base = seg_info.to

        delta = new_base - old_base
        if delta == 0:
            return 0

        patch_comments(delta)

        return 0


class RebasedComment(plugin_t):
    flags = ida_idaapi.PLUGIN_MOD | ida_idaapi.PLUGIN_UNL | ida_idaapi.PLUGIN_HIDE
    comment = "Updates offsets in comments when rebased"
    wanted_name = "rebasedcomment"
    wanted_hotkey = ""
    help = ""

    def __init__(self):
        pass
    
    def init(self):
        self._hook = Hooks()
        self._hook.hook()
        
        addon = idaapi.addon_info_t()
        addon.id = "naim94a.oss.rebasedcomment"
        addon.url = "https://abda.nl/?rebasedcomment"
        addon.version = __version__
        addon.producer = "Naim A. <naim@abda.nl>"
        addon.freeform = "Copyright (C) 2020 Naim A."
        addon.name = "RebasedComment"
        idaapi.register_addon(addon)
        return ida_idaapi.PLUGIN_KEEP

    def run(self, arg):
        return 1
    
    def term(self):
        self._hook.unhook()
        pass


def PLUGIN_ENTRY():
    return RebasedComment()


repository: soggysec/psida
file: psida/idb_push_plugin.py
wanted_name: psida's IDB Push v0.2
created_at: 2017-11-28 16:59:49+00:00
pushed_at: 2018-08-03 22:30:34+00:00
url:
https://github.com/soggysec/psida/blob/f03216e078abbf2d62b683bcf4a5134fcce210af/psida/idb_push_plugin.py

content:
import idaapi
import idc
import ida_kernwin
import socket
from PyQt5 import QtGui, QtCore, QtWidgets


class IdbPushPlugin(idaapi.plugin_t):
    flags = None
    comment = """  Press Ctrl+Shift+P to start the plugin.
                Troubleshooting:
                    - Make sure psida's directory is in your PYTHONPATH
                    - Run `import psida; psida.idb_push.configure(backend_hostname='<your backend server or IP>')`"""
    help = str(comment)
    wanted_name = "psida's IDB Push v0.2"
    wanted_hotkey = "Ctrl-Shift-P"
    imported = False
    running = False
    psida_module = None

    def init(self):
        try:
            import psida
            self.psida_module = psida
            self.imported = True

            if self.psida_module.idb_push.CONFIGURATION["debug"]:
                idaapi.msg("DEBUG - Init - Successfully initialized and imported psida\n")

            return idaapi.PLUGIN_KEEP
        except ImportError:
            psida = None
            self.psida_module = psida
            idaapi.msg("Error importing psida module. Make sure it resides in any directory that is in your PYTHONPATH\n")
            return idaapi.PLUGIN_HIDE

    def reload(self):
        try:
            self.psida_module.idb_push.stop()
            self.running = False
            idaapi.msg("DEBUG - Run - idb_push already running, stopping and reloading...\n")
            reload(self.psida_module)

            if self.psida_module.idb_push.CONFIGURATION["debug"]:
                idaapi.msg("DEBUG - Reload - Successfully reloaded psida\n")

            self.psida_module.idb_push.start()
            self.running = True
            if self.psida_module.idb_push.CONFIGURATION["debug"]:
                idaapi.msg("DEBUG - Run - Successfully started idb_push\n")
        except ImportError:
            idaapi.msg("Error reloading psida module. Make sure it resides in any directory that is in your PYTHONPATH\n")
        
    def run(self, arg):
        if self.imported:
            if not self.psida_module.idb_push.CONFIGURATION["backend_hostname"]:
                connected = False
                while not connected:
                    backend_hostname = ida_kernwin.ask_str("Hostname or IP", 0, "Backend not initialized, input your backend's name or IP:")
                    if not backend_hostname:
                        # User canceled
                        return
                    try:
                        print backend_hostname
                        self.psida_module.zmq_primitives.configure(backend_hostname=backend_hostname)
                        reload(self.psida_module)
                        reload(self.psida_module.zmq_primitives)
                        # test connectivity
                        self.psida_module.zmq_primitives.zmq_test_connectivity()
                        connected = True
                    except self.psida_module.zmq_primitives.ZMQConnectionException:
                        idaapi.msg("ERROR - Run - ZMQ Connectivity failed, make sure your server is set-up correctly.\n")
                    except socket.gaierror:
                        idaapi.msg("ERROR - Run - Could not resolve server name. Make sure it's spelled correctly, and that you get DNS responses from it\n")

            if self.running:
                # Reload only if in debug mode. Do nothing otherwise
                if self.psida_module.idb_push.CONFIGURATION["debug"]:
                    self.reload()
            else:
                reload(self.psida_module)
                self.psida_module.idb_push.start()
                if self.psida_module.idb_push.CONFIGURATION["debug"]:
                    idaapi.msg("DEBUG - Run - Successfully started idb_push\n")
                self.running = True

    def term(self):
        if self.running:
            self.psida_module.idb_push.stop()
            self.running = False
            if self.psida_module.idb_push.CONFIGURATION["debug"]:
                idaapi.msg("DEBUG - Term - Successfully stopped idb_push\n")


def PLUGIN_ENTRY():
    return IdbPushPlugin()

comment:
  Press Ctrl+Shift+P to start the plugin.
                Troubleshooting:
                    - Make sure psida's directory is in your PYTHONPATH
                    - Run `import psida; psida.idb_push.configure(backend_hostname='<your backend server or IP>')`


repository: gistable/gistable
file: all-gists/8bfa8e25f36e6f619247203d1bc0291b/snippet.py
description: GitHub gists and the dockerfiles needed to run them.
wanted_name: AddSegment
comment: Add a segment to an IDA .idb from a file.
created_at: 2018-02-01 16:19:24+00:00
pushed_at: 2023-02-11 18:20:10+00:00
url:
https://github.com/gistable/gistable/blob/665d39a2bd82543d5196555f0801ef8fd4a3ee48/all-gists/8bfa8e25f36e6f619247203d1bc0291b/snippet.py

content:
'''
IDAPython plugin that adds the contents of a file as a new segment in an existing idb.
Prompts the user for:
  - file path
  - segment name
  - segment starting offset
  
Useful for reversing engineering packed software and shellcode.

Author: Willi Ballenthin <william.ballenthin@fireeye.com>
Licence: Apache 2.0
'''
import logging
from collections import namedtuple

import idc
import idaapi
import idautils


logger = logging.getLogger(__name__)


class BadInputError(Exception):
    pass


Segment = namedtuple('SegmentBuffer', ['path', 'name', 'addr'])


def prompt_for_segment():
    ''' :returns: a Segment instance, or raises BadInputError '''
    class MyForm(idaapi.Form):
        def __init__(self):
            idaapi.Form.__init__(self, """STARTITEM 0
add segment by buffer

<##buffer path:{path}>
<##segment name:{name}>
<##segment start address:{addr}>
""",
{
    'path': idaapi.Form.FileInput(open=True),
    'name': idaapi.Form.StringInput(),
    'addr': idaapi.Form.NumericInput(tp=Form.FT_ADDR),
})

        def OnFormChange(self, fid):
            return 1

    f = MyForm()
    f.Compile()
    f.path.value = ""
    f.name.value = ""
    f.addr.value = 0x0
    ok = f.Execute()
    if ok != 1:
        raise BadInputError('user cancelled')

    path = f.path.value
    if path == "" or path is None:
        raise BadInputError('bad path provided')

    if not os.path.exists(path):
        raise BadInputError('file doesn\'t exist')

    name = f.name.value
    if name == "" or name is None:
        raise BadInputError('bad name provided')

    addr = f.addr.value
    f.Free()
    return Segment(path, name, addr)


def main(argv=None):
    if argv is None:
        argv = sys.argv[:]

    try:
        seg = prompt_for_segment()
    except BadInputError:
        logger.error('bad input, exiting...')
        return -1

    with open(seg.path, 'rb') as f:
        buf = f.read()

    seglen = len(buf)
    if seglen % 0x1000 != 0:
        seglen = seglen + (0x1000 - (seglen % 0x1000))

    if not idc.AddSeg(seg.addr, seg.addr + seglen, 0, 1, 0, idaapi.scPub):
        logger.error('failed to add segment: 0x%x', seg.addr)
        return -1

    if not idc.RenameSeg(seg.addr, seg.name):
        logger.warning('failed to rename segment: %s', seg.name)

    if not idc.SetSegClass(seg.addr, 'CODE'):
        logger.warning('failed to set segment class CODE: %s', seg.name)

    if not idc.SegAlign(seg.addr, idc.saRelPara):
        logger.warning('failed to align segment: %s', seg.name)

    idaapi.patch_many_bytes(seg.addr, buf)
    

class AddSegmentPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_KEEP
    comment = "Add a segment to an IDA .idb from a file."

    help = "Add a segment to an IDA .idb from a file."
    wanted_name = "AddSegment"
    wanted_hotkey = "Alt-F8"

    def init(self):
        return idaapi.PLUGIN_OK

    def run(self, arg):
        main()

    def term(self):
        pass


def PLUGIN_ENTRY():
    return AddSegmentPlugin()


#if __name__ == '__main__':
#    logging.basicConfig(level=logging.DEBUG)
#    main()


repository: michal-kapala/dubRE
file: plugins/funcdata_export.py
description: ML-driven function symbol extraction plugin for IDA Pro
wanted_name: dubRE Function Data Exporter
comment: dubRE Function Data Exporter
created_at: 2023-05-23 15:03:18+00:00
pushed_at: 2024-12-17 22:26:33+00:00
url:
https://github.com/michal-kapala/dubRE/blob/ce7be63623178ed55081e3b3297a4b4f3d49ac6a/plugins/funcdata_export.py

content:
import idaapi, idc, sark
import sqlite3
import os, platform


IDATA = sark.Segment(name=".idata")
"""`.idata` imports section (Sark segment)."""

RDATA = sark.Segment(name=".rdata")
"""`.rdata` data section (Sark segment)."""

def get_filename():
  """Returns the extensionless IDB file name (removes `-stripped` suffix)."""
  path = idc.get_input_file_path()
  return os.path.splitext(os.path.basename(path))[0].replace("-stripped", "")

def is_import(ea):
  """Checks if an address belongs to `.idata` section."""
  return ea >= IDATA.ea and ea < IDATA.ea + IDATA.size

def is_before_rdata(ea):
  """Checks if the address is from before `.rdata` section."""
  return ea < RDATA.ea

def is_tostring_xref(xref):
  """Checks if cross reference points to a string."""
  # IDAPython docs say `get_strlit_contents` returns empty strings on failure but it actually returns None
  # https://www.hex-rays.com/products/ida/support/idapython_docs/idc.html#idc.get_strlit_contents
  return (
    xref.type.is_data
    and idc.get_strlit_contents(xref.to) not in [None, ""]
    and not is_before_rdata(xref.to)
    and not is_import(xref.to)
  )

def get_nb_xref_to(ea):
  """Returns the number of code references to a function."""
  func = sark.Function(ea)
  unique = []

  for xref in func.xrefs_to:
    if xref.type.is_code and xref.frm not in unique:
      unique.append(xref.frm)

  return len(unique)


def get_nb_xref_from(ea):
  """Returns the number of code references from a function."""
  func = sark.Function(ea)
  unique = []
  
  for xref in func.xrefs_from:
    if xref.type.is_code and xref.to not in unique:
      unique.append(xref.to)

  return len(unique)

def get_nb_strings(ea):
  """Returns the number of function's unique string xrefs."""
  func = sark.Function(ea)
  unique = []
  
  for xref in func.xrefs_from:
    if is_tostring_xref(xref) and xref.to not in unique:
        unique.append(xref.to)

  return len(unique)

def export_func_data():
  """Exports a function data to SQLite db (`funcs` table) in the current directory."""
  filename = get_filename() + '.db'
  trailing_slash = "\\" if platform.system() == "Windows" else "/"
  print("Exporting function data to " + os.getcwd() + trailing_slash + filename)
  
  conn = sqlite3.connect(filename)
  conn.text_factory = lambda x: x.decode("utf-8")
  c = conn.cursor()

  # Fetch addresses
  c.execute("SELECT func_addr FROM funcs")
  addrs = c.fetchall()
  func_cnt = 0

  for row in addrs:
    addr = row[0]
    func = sark.Function(addr)
    func_cnt += 1
    print(str(func_cnt) + " - " + func.name)

    nb_referrers = get_nb_xref_to(addr)
    nb_strings = get_nb_strings(addr)
    nb_referees = get_nb_xref_from(addr)
    nb_lines = len(list(func.lines))

    c.execute('''UPDATE funcs SET
      nb_referrers = ?,
      nb_strings = ?,
      nb_referees = ?,
      instructions = ?
      WHERE func_addr = ?''',
      (nb_referrers, nb_strings, nb_referees, nb_lines, addr))

  print("Exported all functions.")
  conn.commit()
  conn.close()


class FuncDataExporterPlugin(idaapi.plugin_t):
  flags = idaapi.PLUGIN_PROC
  comment = "dubRE Function Data Exporter"
  help = "This plugin exports basic function data to SQLite database."
  wanted_name = "dubRE Function Data Exporter"
  wanted_hotkey = "Shift+D"

  def init(self):
    return idaapi.PLUGIN_KEEP

  def term(self):
    pass

  def run(self, arg):
    export_func_data()


def PLUGIN_ENTRY():
  return FuncDataExporterPlugin()


repository: Kwansy98/find_xrefs
file: find_xrefs.py
description: ida script for printing xrefs function call chains
url: https://github.com/Kwansy98/find_xrefs/blob/3d28e066f828f1d5ef29a16f6a24cdc8cd0a1fad/find_xrefs.py
wanted_name: Xref Finder
comment: Xref Finder Plugin
created_at: 2024-11-04 12:27:37+00:00
pushed_at: 2025-02-15 17:47:41+00:00
content:
import idaapi
import idautils
import idc
import ida_kernwin
import re
import os
import json

def get_config_path():
    # Return the config file path in the IDA user directory.
    return os.path.join(idaapi.get_user_idadir(), "xref_finder_config.json")

def load_config():
    # Load the depth settings from the config file. Default: default_from=1, default_to=0.
    config_path = get_config_path()
    if os.path.exists(config_path):
        try:
            with open(config_path, "r") as f:
                config = json.load(f)
            return config.get("default_from", 1), config.get("default_to", 0)
        except Exception:
            return 1, 0
    return 1, 0

def save_config(default_from, default_to):
    # Save the depth settings to the config file.
    config_path = get_config_path()
    config = {"default_from": default_from, "default_to": default_to}
    with open(config_path, "w") as f:
        json.dump(config, f, indent=4)

# Load default depth values.
default_from, default_to = load_config()

blacklist = [
    '__CxxThrowException',
    '__invalid_parameter',
    'nullsub_',
    '__security_check_cookie',
    '___report_gsfailure',
    '___raise_securityfailure',
    '__invoke_',
    '__SEH_prolog',
    '___acrt',
    '___CxxFrameHandler',
    '__InternalCxxFrameHandler',
    '__errno'
]

class XrefFinderPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = "Xref Finder Plugin"
    help = "print xrefs to and xrefs from"
    wanted_name = "Xref Finder"
    wanted_hotkey = "Ctrl-Shift-F"
    
    def init(self):
        return idaapi.PLUGIN_OK

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

    def term(self):
        pass

    def show_dialog(self):
        class XrefFinderForm(ida_kernwin.Form):
            def __init__(self):
                # Use default depths loaded from the config file.
                self.depth_to = default_to
                self.depth_from = default_from
                current_ea = ida_kernwin.get_screen_ea()
                # Get the function name based on the current cursor.
                default_func_name = idc.ida_funcs.get_func_name(current_ea) or ""
                ida_kernwin.Form.__init__(self, r"""STARTITEM 0
BUTTON YES* OK
Xref Finder
{FormChangeCb}
<##Function name:{iFuncName}>
<##xrefs to depth:{iDepthTo}>
<##xrefs from depth:{iDepthFrom}>
""", {
                    'iFuncName': ida_kernwin.Form.StringInput(value=default_func_name),
                    'iDepthTo': ida_kernwin.Form.NumericInput(tp=ida_kernwin.Form.FT_DEC, value=self.depth_to),
                    'iDepthFrom': ida_kernwin.Form.NumericInput(tp=ida_kernwin.Form.FT_DEC, value=self.depth_from),
                    'FormChangeCb': ida_kernwin.Form.FormChangeCb(self.OnFormChange),
                })
            def OnFormChange(self, fid):
                return 1

        form = XrefFinderForm()
        form.Compile()
        ok = form.Execute()
        if ok == 1:
            func_name = form.iFuncName.value
            depth_to = form.iDepthTo.value
            depth_from = form.iDepthFrom.value
            global default_from, default_to
            default_from, default_to = depth_from, depth_to
            save_config(default_from, default_to)
            self.find_xrefs(func_name, depth_to, depth_from)
        form.Free()

    def find_subfunctions(self, func_name, depth, current_depth=0, path=None):
        if path is None:
            path = [func_name]
        else:
            path.append(func_name)
        if current_depth > depth:
            return
        for black in blacklist:
            if black in func_name:
                return
        func_ea = idc.get_name_ea_simple(func_name)
        if func_ea == idc.BADADDR:
            path.pop()
            try:
                if func_name.startswith("0x"):
                    target_address = int(func_name, 16)
                else:
                    target_address = int(func_name)
            except ValueError:
                return
            lib_func = None
            for xref in idautils.XrefsTo(target_address, idaapi.XREF_FAR):
                is_data_xref = (xref.type & idaapi.XREF_DATA) != 0
                if not is_data_xref:
                    disasm = idc.GetDisasm(xref.frm)
                    parts = re.split(r'[^a-zA-Z0-9]', disasm)
                    lib_func = parts[-1] if parts and parts[-1] else None
                    break
            if lib_func:
                path.append(lib_func)
                lst = " -> ".join(path)
                print(lst)
            return
        called_functions = set()
        for (startea, endea) in idautils.Chunks(func_ea):
            for head in idautils.Heads(startea, endea):
                if idc.is_code(idc.get_full_flags(head)):
                    refs = idautils.CodeRefsFrom(head, 0)
                    for ref in refs:
                        ref_name = idc.get_func_name(ref)
                        if not ref_name:
                            ref_name = hex(ref)
                        if ref_name and ref_name != func_name:
                            called_functions.add((ref, ref_name))
        if current_depth == depth or not called_functions:
            print(" -> ".join(path))
        else:
            for addr, name in called_functions:
                self.find_subfunctions(name, depth, current_depth + 1, path[:])

    def find_parentfunctions(self, func_name, depth, current_depth=0, path=None):
        if path is None:
            path = [func_name]
        else:
            path.append(func_name)
        if current_depth > depth:
            return
        func_ea = idc.get_name_ea_simple(func_name)
        if func_ea == idc.BADADDR:
            print(f"Function {func_name} not found.")
            return
        calling_functions = set()
        for ref in idautils.CodeRefsTo(func_ea, 0):
            ref_name = idc.get_func_name(ref)
            if not ref_name:
                ref_name = hex(ref)
            if ref_name and ref_name != func_name:
                calling_functions.add((ref, ref_name))
        if current_depth == depth or not calling_functions:
            print(" -> ".join(reversed(path)))
        else:
            for addr, name in calling_functions:
                self.find_parentfunctions(name, depth, current_depth + 1, path[:])

    def find_xrefs(self, func_name, depth_to, depth_from):
        self.find_parentfunctions(func_name, depth_to)
        self.find_subfunctions(func_name, depth_from)

def PLUGIN_ENTRY():
    return XrefFinderPlugin()


repository: techbliss/Python_editor
file: 7.5/Add to plugins folder/Python_editor.py
description: Better CodeEditor for Ida Pro.
wanted_name: Python Editor
comment: Run me
created_at: 2015-07-04 21:14:02+00:00
pushed_at: 2020-12-25 13:54:19+00:00
url:
https://github.com/techbliss/Python_editor/blob/4c2b67bd0b1a2b9b9403f5ae31cf8f13c247c981/7.5/Add%20to%20plugins%20folder/Python_editor.py

content:
# Created by: Storm Shadow http://www.techbliss.org
import os, sys
import ida_idaapi, ida_kernwin, ida_idaapi, ida_loader
import idc
from idc import *
from ida_diskio import *
from ida_idaapi import *
from ida_loader import *
import sys
sys.path.insert(0 , ida_diskio.idadir("plugins\\Code editor\\icons"))
import ico
from ico import *

PLUGIN_VERSION = "1.6"
IDAVERISONS    = "IDA PRO 7.2+"
AUTHORS        = "Storm Shadow"
DATE           = "2019"
TWITTER        = "Twitter @zadow28"

def banner():
    banner_options = (PLUGIN_VERSION, AUTHORS, DATE, TWITTER, IDAVERISONS)
    banner_titles = "Python Editor v%s - (c) %s - %s - %s - %s" % banner_options

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

banner()

# 1) Create the handler class
class MyEditorHandler(ida_kernwin.action_handler_t):
    def __init__(self):
        ida_kernwin.action_handler_t.__init__(self)

    # Run editor when invoked.
    def activate(self, ctx):
        g = globals()
        idahome = ida_diskio.idadir("plugins\\Code editor")
        tupac = str(idahome + "\\pyeditor.py")
        ida_kernwin.load_custom_icon(":/ico/python.png")
        load_and_run_plugin(tupac, True)
    def update(self, ctx):
        return ida_kernwin.AST_ENABLE_ALWAYS

class ripeye(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_FIX
    comment = "Run me"
    help = "Python Editor"
    wanted_name = "Python Editor"
    wanted_hotkey = "" #the tooltip horkey goes away when setting it here DONT DO it! and only is shown in File/Plugins menu


    def editor_menuaction(self):
        action_desc = ida_kernwin.action_desc_t(
            'my:editoraction',  # The action name. This acts like an ID and must be unique
            'Python Editor!',  # The action text.
            MyEditorHandler(),  # The action handler.
            'Ctrl+H',  # Optional: the action shortcut DO IT  HERE!
            'Script editor',  # Optional: the action tooltip (available in menus/toolbar)
            ida_kernwin.load_custom_icon(":/ico/python.png")  # hackish load action icon , if no custom icon use number from 1-150 from internal ida
        )

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

        ida_kernwin.attach_action_to_menu(
            'Edit/Editor...',  # The relative path of where to add the action
            'my:editoraction',  # The action ID (see above)
            ida_kernwin.SETMENU_APP)  # We want to append the action after the 'Manual instruction...

        ida_kernwin.get_current_widget()
        ida_kernwin.attach_action_to_popup(form, None, "my:editoraction", None)

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

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

        return PLUGIN_KEEP


    def _install_plugin(self):
        """
        Initialize & integrate the plugin into IDA.
        """
        self.editor_menuaction()
        self._init()

    def term(self):
        pass

    def run(self, arg = 0):
        #we need the calls again if we wanna load it via File/Plugins/editor
        ida_kernwin.msg("Python Editor Loaded to menu \n use Alt+E hot key to quick load ")
        hackish = MyEditorHandler()
        hackish.activate(self)

def PLUGIN_ENTRY():
    return ripeye()


repository: doronz88/ida-scripts
file: objc_method_params.py
wanted_name: RenameObjCParams
comment: Rename Objective-C method parameters and function name according to the selector
created_at: 2023-05-03 05:44:52+00:00
pushed_at: 2025-01-06 21:38:34+00:00
url:
https://github.com/doronz88/ida-scripts/blob/7db238438638a2faab5699dc3109a08cb9d4d769/objc_method_params.py

content:
import dataclasses
import re

import idaapi
import idc
import ida_hexrays
HOTKEY = 'Ctrl+0'


@dataclasses.dataclass
class Arg:
    a_type: str = 'id'
    a_name: str = 'no_name'

    def __str__(self):
        return f'{self.a_type} {self.a_name}'


class Function:
    def __init__(self, objc_name: str, sig: str):
        self.start = None
        self.objc_class = None
        self.args = []
        self.sig_name = objc_name.replace('[', '__').replace(']', '_').replace(':', '_').replace(' ', '_').strip(
            '+').strip('-')
        self._parse_signature(sig)
        self._parse_objc_name(objc_name)

    def _parse_objc_name(self, objc_name: str):
        pattern = re.compile(r'\[(?P<class_name>\S+)\s+(?P<selector>[^]]+)]')
        key_words = ['With', 'From', 'At', 'Using', 'Writing', 'set']
        name_parts = pattern.search(objc_name)

        self.objc_class = name_parts.group('class_name')
        selector = name_parts.group('selector')

        parts = [p for p in selector.split(':') if p]
        args = []
        for k in parts:
            special_case_match = re.search(fr'(?i)({"|".join(key_words)})([A-Z][a-z]*)', k)
            if special_case_match:
                args.append(special_case_match.group(2))
            else:
                args.append(k)

        for i in range(len(self.args)):
            self.args[i].a_name = camel_to_snake(args[i])

    def _parse_signature(self, sig: str):
        pattern = re.compile(r'^(?P<start>[^\(]+)\((?P<args>.*)\)$')
        sig_parts = pattern.search(sig)
        self.start = sig_parts.group('start')
        self.args = [Arg(a_type=t.strip(' ')) for t in sig_parts.group('args').split(',')[2:]]

    def __str__(self):
        args = ', ' + ', '.join(str(a) for a in self.args) if self.args else ''
        return f'{self.start} {self.sig_name}({self.objc_class} *self, SEL selector{args})'


def camel_to_snake(name: str):
    s1 = re.sub(r'(.)([A-Z][a-z]+)', r'\1_\2', name)
    return re.sub(r'([a-z0-9])([A-Z])', r'\1_\2', s1).lower()


def rename_objc_method_params():
    objc_method_pattern = re.compile(r'\[(?P<class_name>\S+)\s+(?P<selector>[^\]]+)\]')

    ea = idaapi.get_screen_ea()
    if ea == idaapi.BADADDR:
        print('No function is currently displayed on the screen.')
        return

    function_name = idc.get_func_name(ea)
    match = objc_method_pattern.search(function_name)
    if not match:
        print(f'Function {function_name} does not match the Objective-C pattern.')
        return

    func_type = idc.get_type(ea)
    if not func_type:
        print(f'No type information for function {function_name}.')
        return

    func = Function(function_name, func_type)

    idc.SetType(ea, str(func))
    print(f'Renamed {function_name} to {func}')
    vdui = ida_hexrays.open_pseudocode(ea, 0)
    if vdui:
        vdui.refresh_view(True)


class RenameObjCParamsPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = 'Rename Objective-C method parameters and function name according to the selector'
    help = 'Rename Objective-C method parameters and function name according to the selector'
    wanted_name = 'RenameObjCParams'
    wanted_hotkey = HOTKEY

    def init(self):
        idaapi.msg('RenameObjCParams plugin initialized.\n')
        return idaapi.PLUGIN_OK

    def run(self, arg):
        rename_objc_method_params()

    def term(self):
        idaapi.msg('RenameObjCParams plugin terminated.\n')


def PLUGIN_ENTRY():
    return RenameObjCParamsPlugin()


repository: cseagle/ida_clemency
file: clemency_dump.py
description: IDA cLEMENCy Tools
wanted_name: Dump Clemency Database to Clemency binary
created_at: 2017-07-31 02:51:32+00:00
pushed_at: 2017-07-31 02:53:00+00:00
url:
https://github.com/cseagle/ida_clemency/blob/7101ce142df676c7b62b176a96f43f8df01ce3d6/clemency_dump.py

content:
from idaapi import *

'''
Author: Chris Eagle
Name: Clemency database dump plugin defcon 25
What: Dumps the current .text section back to a packed 9-bit disk file
How: Install into <idadir>/plugins
     Activate  using Alt-9
'''

class clemency_dump_plugin_t(plugin_t):
   flags = 0
   wanted_name = "Dump Clemency Database to Clemency binary"
   wanted_hotkey = "Alt-9"
   comment = ""
   help = ""

   def init(self):
      return PLUGIN_OK

   def term(self):
      pass

   pats = [
      "0000", "0001", "0010", "0011", "0100", "0101", "0110", "0111",
      "1000", "1001", "1010", "1011", "1100", "1101", "1110", "1111"
   ]

   def run(self, arg):
      fname = AskFile(1, "*.*", "Choose save file")
      if fname is not None and len(fname) > 0:
         with open(fname, "wb") as f:
            text = get_segm_by_name(".text")
            bits = ''
            for i in range(text.startEA, text.endEA):
               b = Byte(i)
               if b & 0x100:
                  bits += '1'
               else:
                  bits += '0'
               bits += self.pats[(b >> 4) & 0xf]
               bits += self.pats[b & 0xf]

               while len(bits) >= 8:
                  f.write(chr(int(bits[0:8], 2)))
                  bits = bits[8:]

            if len(bits):
               while len(bits) % 8:
                  bits += '0'
               f.write(chr(int(bits[0:8], 2)))      

def PLUGIN_ENTRY():
   return clemency_dump_plugin_t()


repository: grayhatacademy/ida
file: plugins/mipsrop/mipsrop.py
wanted_name: MIPS ROP Finder
comment: MIPS ROP Finder
created_at: 2013-10-16 18:52:21+00:00
pushed_at: 2023-10-23 10:38:19+00:00
url:
https://github.com/grayhatacademy/ida/blob/faf13a13c4f5070fe8f0666c1508f928ffe40e23/plugins/mipsrop/mipsrop.py

content:
# IDA plugin for identifying ROP gadgets in Linux MIPS binaries.
#
# Return Oriented Programming in Linux MIPS is more like Jump Oriented
# Programming; the idea is to control enough of the stack/registers in order to
# control various jumps. Since all instructions in MIPS must be 4-byte aligned,
# You cannot "create" new instructions by returning into the middle
# of existing instructions, as is possible with some other architectures.
#
# In any given MIPS function, various registers are saved onto the stack by
# necessity:
#
#    o $s0 - $s7
#    o $fp
#    o $ra
#
# These values are restored from the stack before the function returns, thus,
# during a stack overflow one can control some or all of these register values.
# The subroutine registers ($s*) are of particular interest, as they are
# commonly used by the compiler to store function pointers. By convention, gcc
# will move function pointers into the $t9 register, then call the function
# using jalr:
#
#     move $t9, $s0  <-- If we control $s0, we control where the jump is taken
#    jalr $t9
#
# While there are other jumps that are of use, and which this plugin searches
# for, the premise is the same: control the stack/registers, and you control
# various jumps allowing you to chain various blocks of code together.
#
# With a list of controllable jumps such as these, we then just need to search
# the surrounding instructions to see if they perform some operation which may
# be useful. For example, let's say we need to load the value 1 into the $a0
# register; in this case, we would want to look for a controllable jump such as
# this:
#
#    move $t9, $s1
#    jalr $t9
#    li $a0, 1    <-- Remember MIPS has jump delay slots, so this instruction
#    is executed with the jump
#
# If we return to this piece of code (and if we control $s1), we can pre-load
# $s1 with the address of the next ROP gadget; thus, $a0 will be loaded with
# the value 1 and we can chain this block of code with other gadgets in order
# to perform more complex operations.
#
# This plugin finds all potentially controllable jumps, and then allows you to
# search for desired instructions surrounding these controllable jumps. Example:
#
#    Python> mipsrop.find("li $a0, 1")
#    -------------------------------------------------
#    |  Address     |  Action      |  Control Jump   |
#    -------------------------------------------------
#    |  0x0002F0F8  |  li $a0,1    |  jalr  $s4      |
#    |  0x00057E50  |  li $a0,1    |  jalr  $s1      |
#    -------------------------------------------------
#
# The output shows the offset of each ROP gadget, the instruction within the
# gadget that your search matched, and the effective register that is jumped to
# after that instruction is executed.
#
# The specified instruction can be a full instruction, such as the example
# above, or a partial instruction. Regex is supported for any of the
# instruction mnemonics or operands; for convenience, the dollar signs in front
# of register names are automatically escaped.
#
# If .set_base(int) is called with a non-zero value then the output will
# include base, offset and address.
# Example
# -------------------------------------------------------------------------
# |  Base       + Offset     =  Address     |  Action  |  Control Jump    |
# -------------------------------------------------------------------------
#
# Craig Heffner
# Tactical Network Solutions

import re
import idc
import sys
import idaapi
import idautils

from shims import ida_shims

# Global instance of MIPSROPFinder
mipsrop = None


def add_to_namespace(namespace, name, variable):
    '''
    Add a variable to a different namespace, likely __main__.
    '''
    import importlib

    importer_module = sys.modules[namespace]
    if name in list(sys.modules.keys()):
        if not (sys.version_info.major == 3 and sys.version_info.minor >= 4):
            import imp
            imp.reload(sys.modules[name])
        else:
            importlib.reload(sys.modules[name])
    else:
        m = importlib.import_module(name, None)
        sys.modules[name] = m

    setattr(importer_module, name, variable)


class MIPSInstruction(object):
    '''
    Class for storing info about a specific instruction.
    '''

    def __init__(self, mnem, opnd0=None, opnd1=None, opnd2=None,
                 ea=idc.BADADDR):
        self.mnem = mnem
        self.operands = [opnd0, opnd1, opnd2]
        self.opnd0 = opnd0
        self.opnd1 = opnd1
        self.opnd2 = opnd2
        self.ea = ea

    def __str__(self):
        string = self.mnem + " "

        for op in self.operands:
            if op:
                string += "%s," % op
            else:
                break

        return string[:-1]


class ROPGadget(object):
    '''
    Class for storing information about a specific ROP gadget.
    '''
    def __init__(self, control, jump, operation=None, description="ROP gaget",
                 base=0):
        self.control = control
        self.exit = jump
        self.operation = operation
        self.description = description
        self.base = base
        self.h = '-' * 112
        if self.base != 0:
            self.h += '-' * 27

        if self.control.opnd1:
            self.control.register = self.control.opnd1
        else:
            self.control.register = self.control.opnd0

        if self.exit.opnd1:
            self.exit.register = self.exit.opnd1
        else:
            self.exit.register = self.exit.opnd0

        if self.operation:
            if self.operation.ea < self.control.ea:
                self.entry = self.operation
            else:
                self.entry = self.control
        else:
            self.operation = self.control
            self.entry = self.control


    def header(self):
        if self.base != 0:
            return self.h + "\n|  Base       + Offset     =  Address       " \
                            "|  Action                                     " \
                            "         |  Control Jump                      " \
                            "    |\n" + self.h
        else:
            return self.h + "\n|  Address     |  Action                     " \
                            "                         |  Control Jump       " \
                            "                   |\n" + self.h

    def footer(self):
        return self.h

    def __str__(self):
        if self.base != 0:
            return "|  0x%.8X + 0x%.8X = 0x%.8X     |  %-50s  | " \
                   " %-5s %-30s  |" % \
                   (self.base, self.entry.ea, self.entry.ea + self.base,
                    str(self.operation), self.exit.mnem, self.control.register)
        else:
            return "|  0x%.8X  |  %-50s  |  %-5s %-30s  |" % \
                   (self.entry.ea, str(self.operation), self.exit.mnem,
                    self.control.register)


class BowcasterBuilder(object):
    '''
    Class to generate bowcaster code from a list of selected ROP gadgets. WIP.
    '''
    INSIZE = 4
    SEARCH_DEPTH = 25

    def __init__(self, gadgets):
        self.code = []
        self.gadgets = gadgets

    def build_code(self):
        keys = list(self.gadgets.keys())
        keys.sort()

        for key in keys[::-1]:
            last_instruction = False
            ea = self.gadgets[key]
            end_ea = ea + self.SEARCH_DEPTH

            while ea <= end_ea:
                mnem = idc.print_insn_mnem(ea)
                if mnem in ['jr', 'jalr']:
                    last_instruction = True
                ea += self.INSIZE

    def print_code(self):
        for line in self.code:
            print(line)


class MIPSROPFinder(object):
    '''
    Primary ROP finder class.
    '''
    CODE = 2
    DATA = 3
    INSIZE = 4
    SEARCH_DEPTH = 25

    def __init__(self):
        self.base = 0

        self._initial_find()

        if self.controllable_jumps or self.system_calls:
            print("MIPS ROP Finder activated, found %d controllable jumps " \
                  "between 0x%.8X and 0x%.8X" % (len(self.controllable_jumps),
                                                 self.start, self.end))
        else:
            print("No ROP gadgets found!")

    def _initial_find(self):
        self.start = idc.BADADDR
        self.end = idc.BADADDR
        self.system_calls = []
        self.double_jumps = []
        self.controllable_jumps = []
        start = 0
        end = 0

        for (start, end) in self._get_segments(self.CODE):
            self.controllable_jumps += self._find_controllable_jumps(start, end)
            self.system_calls += self._find_system_calls(start, end)
            self.double_jumps += self._find_double_jumps(start, end)
            if self.start == idc.BADADDR:
                self.start = start
        self.end = end

    def _get_segments(self, attr):
        segments = []
        seg = ida_shims.get_first_seg()

        while seg != idc.BADADDR:
            if ida_shims.get_segm_attr(seg, idc.SEGATTR_TYPE) == attr:
                start = ida_shims.get_segm_start(seg)
                end = ida_shims.get_segm_end(seg)
                segments.append((start, end))
            seg = ida_shims.get_next_seg(seg)

        return segments

    def _get_instruction(self, ea):
        return MIPSInstruction(ida_shims.print_insn_mnem(ea),
                               ida_shims.print_operand(ea, 0),
                               ida_shims.print_operand(ea, 1),
                               ida_shims.print_operand(ea, 2), ea)

    def _does_instruction_match(self, ea, instruction, regex=False):
        i = 0
        op_cnt = 0
        op_ok_cnt = 0
        match = False
        insn = ida_shims.decode_insn(ea)
        mnem = ida_shims.print_insn_mnem(ea)
        if (not instruction.mnem) or (instruction.mnem == mnem) or \
                (regex and re.match(instruction.mnem, mnem)):
            for operand in instruction.operands:
                if operand:
                    op_cnt += 1
                    op = ida_shims.print_operand(ea, i)

                    if regex:
                        if re.match(operand, op):
                            op_ok_cnt += 1
                    elif operand == op:
                        op_ok_cnt += 1
                i += 1

            if op_cnt == op_ok_cnt:
                match = True
        return match

    def _is_bad_instruction(self, ea, bad_instructions=['j', 'b'],
                            no_clobber=[]):
        bad = False
        mnem = ida_shims.print_insn_mnem(ea)

        if mnem and mnem[0] in bad_instructions:
            bad = True
        else:
            insn = ida_shims.decode_insn(ea)
            feature = ida_shims.get_canon_feature(insn)
            for register in no_clobber:
                if (feature & idaapi.CF_CHG1) == idaapi.CF_CHG1:
                    if ida_shims.print_operand(ea, 0) == register:
                        bad = True

        return bad

    def _contains_bad_instruction(self, start_ea, end_ea,
                                  bad_instructions=['j', 'b'], no_clobber=[]):
        ea = start_ea

        while ea <= end_ea:
            if self._is_bad_instruction(ea, bad_instructions, no_clobber):
                return True
            else:
                ea += self.INSIZE

        return False

    def _find_prev_instruction_ea(self, start_ea, instruction, end_ea=0,
                                  no_baddies=True, regex=False,
                                  dont_overwrite=[]):
        instruction_ea = idc.BADADDR
        ea = start_ea

        while ea >= end_ea:
            if self._does_instruction_match(ea, instruction, regex):
                instruction_ea = ea
                break
            elif no_baddies and self._is_bad_instruction(
                    ea, no_clobber=dont_overwrite):
                break

            ea -= self.INSIZE

        return instruction_ea

    def _find_next_instruction_ea(self, start_ea, instruction,
                                  end_ea=idc.BADADDR, no_baddies=False,
                                  regex=False, dont_overwrite=[]):
        instruction_ea = idc.BADADDR
        ea = start_ea

        while ea <= end_ea:
            if self._does_instruction_match(ea, instruction, regex):
                instruction_ea = ea
                break
            elif no_baddies and self._is_bad_instruction(
                    ea, no_clobber=dont_overwrite):
                break

            ea += self.INSIZE

        return instruction_ea

    def _find_controllable_jumps(self, start_ea, end_ea):
        controllable_jumps = []
        t9_controls = [
            MIPSInstruction("move", "\$t9"),
            MIPSInstruction("addiu", "\$t9", "^\$"),
        ]
        t9_jumps = [
            MIPSInstruction("jr", "\$t9"),
        ]
        ra_controls = [
            MIPSInstruction("lw", "\$ra"),
        ]
        ra_jumps = [
            # TODO: Search for jumps to registers other than $ra.
            MIPSInstruction("jr", "\$ra"),
        ]
        t9_musnt_clobber = ["$t9"]
        ra_musnt_clobber = ["$ra"]

        if idaapi.IDA_SDK_VERSION >= 740:
            t9_jumps.append(MIPSInstruction("jalr", "\$ra"))
        else:
            t9_jumps.append(MIPSInstruction("jalr", "\$t9"))

        for possible_control_instruction in t9_controls+ra_controls:
            ea = start_ea

            if possible_control_instruction in t9_controls:
                jumps = t9_jumps
                musnt_clobber = t9_musnt_clobber
            else:
                jumps = ra_jumps
                musnt_clobber = ra_musnt_clobber

            while ea <= end_ea:
                ea = self._find_next_instruction_ea(
                    ea, possible_control_instruction, end_ea, regex=True)
                if ea != idc.BADADDR:
                    insn = ida_shims.decode_insn(ea)

                    control_instruction = self._get_instruction(ea)
                    control_register = control_instruction.operands[1]

                    if control_register:
                        for jump in jumps:
                            jump_ea = self._find_next_instruction_ea(
                                ea+insn.size, jump, end_ea, no_baddies=True,
                                regex=True, dont_overwrite=musnt_clobber)

                            if jump_ea != idc.BADADDR:
                                jump_instruction = self._get_instruction(
                                    jump_ea)
                                controllable_jumps.append(
                                    ROPGadget(
                                        control_instruction, jump_instruction,
                                        description="Controllable Jump",
                                        base=self.base))
                                ea = jump_ea

                    ea += insn.size

        return controllable_jumps

    def _find_system_calls(self, start_ea, end_ea):
        system_calls = []
        system_load = MIPSInstruction("la", "$t9", "system")
        stack_arg_zero = MIPSInstruction("addiu", "$a0", "$sp")

        for xref in idautils.XrefsTo(ida_shims.get_name_ea_simple('system')):
            ea = xref.frm
            mnem = ida_shims.print_insn_mnem(ea)
            if mnem and ea >= start_ea and \
                    ea <= end_ea and mnem[0] in ['j', 'b']:
                a0_ea = self._find_next_instruction_ea(
                    ea+self.INSIZE, stack_arg_zero, ea+self.INSIZE)

                if a0_ea == idc.BADADDR:
                    a0_ea = self._find_prev_instruction_ea(
                        ea, stack_arg_zero, ea-(self.SEARCH_DEPTH*self.INSIZE))

                if a0_ea != idc.BADADDR:
                    control_ea = self._find_prev_instruction_ea(
                        ea-self.INSIZE, system_load,
                        ea-(self.SEARCH_DEPTH*self.INSIZE))

                    if control_ea != idc.BADADDR:
                        system_calls.append(
                            ROPGadget(
                                self._get_instruction(control_ea),
                                self._get_instruction(ea),
                                self._get_instruction(a0_ea),
                                description="System call", base=self.base))

                ea += self.INSIZE
            else:
                break

        return system_calls

    def _find_double_jumps(self, start_ea, end_ea):
        double_jumps = []

        for i in range(0, len(self.controllable_jumps)):
            g1 = self.controllable_jumps[i]
            if g1.exit.mnem != 'jalr':
                continue

            for j in range(i+1, len(self.controllable_jumps)):
                g2 = self.controllable_jumps[j]
                distance = (g2.entry.ea - g1.exit.ea)

                if 0 < distance <= (self.SEARCH_DEPTH * self.INSIZE):
                    if g1.control.register != g2.control.register:
                        if not self._contains_bad_instruction(
                                g1.exit.ea+self.INSIZE,
                                g2.control.ea-self.INSIZE,
                                no_clobber=[g2.control.register]):
                            double_jumps.append(g1)
                            break

        return double_jumps

    def _find_rop_gadgets(self, gadget):
        gadget_list = []

        for controllable_jump in self.controllable_jumps:
            gadget_ea = idc.BADADDR

            ea = self._find_next_instruction_ea(
                controllable_jump.entry.ea, gadget,
                controllable_jump.exit.ea+self.INSIZE, regex=True)

            if ea != idc.BADADDR:
                gadget_ea = ea
            else:
                ea = self._find_prev_instruction_ea(
                    controllable_jump.entry.ea, gadget,
                    controllable_jump.entry.ea-(self.SEARCH_DEPTH*self.INSIZE),
                    no_baddies=True, regex=True,
                    dont_overwrite=[controllable_jump.entry.opnd1])

                if ea != idc.BADADDR:
                    gadget_ea = ea

            if gadget_ea != idc.BADADDR:
                gadget_list.append(
                    ROPGadget(
                        controllable_jump.entry, controllable_jump.exit,
                        self._get_instruction(gadget_ea), base=self.base))

        return gadget_list

    def _print_gadgets(self, gadgets):
        if gadgets:
            print(gadgets[0].header())

        for gadget in gadgets:
            print(str(gadget))

        if gadgets:
            print(gadgets[0].footer())

        print("Found %d matching gadgets" % (len(gadgets)))

    def _get_marked_gadgets(self):
        rop_gadgets = {}

        for i in range(0 if idaapi.IDA_SDK_VERSION >= 690 else 1, 1024):
            marked_pos = ida_shims.get_bookmark(i)
            if marked_pos != idc.BADADDR:
                marked_comment = ida_shims.get_bookmark_desc(i)
                if marked_comment and marked_comment.lower().startswith("rop"):
                    rop_gadgets[marked_comment] = marked_pos
            else:
                break

        return rop_gadgets

    def double(self):
        self.doubles()

    def doubles(self):
        '''
        Prints a list of all "double jump" gadgets (useful for function calls).
        '''
        self._print_gadgets(self.double_jumps)

    def stackfinder(self):
        self.stackfinders()

    def stackfinders(self):
        '''
        Prints a list of all gadgets that put a stack address into a register.
        '''
        self.find("addiu .*, $sp")

    def lia0(self):
        '''
        Prints a list of all gadgets that load an immediate value number into
        $a0 (useful for setting up the argument to sleep).
        '''
        self.find("li $a0")

    def tail(self):
        return self.tails()

    def tails(self):
        '''
        Prints a lits of all tail call gadgets (useful for function calls).
        '''
        return self.iret()

    def iret(self):
        '''
        Prints a lits of all tail gadgets (useful for function calls).
        '''
        tail_gadgets = []

        for gadget in self._find_rop_gadgets(MIPSInstruction("move", "\$t9")):
            if gadget.exit.mnem == 'jr' and gadget.exit.register == '$t9':
                tail_gadgets.append(gadget)

        self._print_gadgets(tail_gadgets)

    def system(self):
        '''
        Prints a list of gadgets that may be used to call system().
        '''
        sys_gadgets = self.system_calls + self._find_rop_gadgets(
            MIPSInstruction("addiu", "\$a0", "\$sp"))
        self._print_gadgets(sys_gadgets)

    def find(self, *args):
        '''
        Locates all potential ROP gadgets that contain the specified
        instruction.

        @instruction_string - The instruction you need executed. This can be
        either a:

                    o Full instruction    - "li $a0, 1"
                    o Partial instruction - "li $a0"
                    o Regex instruction   - "li $a0, .*"
        '''
        count = 0
        good_gadgets = {}
        final_gadgets = []
        registers = ['$v', '$s', '$a', '$t', '$k', '$pc', '$fp', '$ra', '$gp',
                     '$at', '$zero']

        for instruction_string in args:
            comma_split = instruction_string.split(',')
            instruction_parts = comma_split[0].split()
            if len(comma_split) > 1:
                instruction_parts += comma_split[1:]

            for i in range(0, 4):
                if i > len(instruction_parts) - 1:
                    instruction_parts.append(None)
                else:
                    instruction_parts[i] = instruction_parts[i].strip().\
                        strip(',').strip()
                    for reg in registers:
                        instruction_parts[i] = instruction_parts[i].replace(
                            reg, "\\%s" % reg)

            instruction = MIPSInstruction(
                instruction_parts[0], instruction_parts[1],
                instruction_parts[2], instruction_parts[3])

            gadgets = self._find_rop_gadgets(instruction)
            for gadget in gadgets:
                if count == 0:
                    good_gadgets[gadget.control.ea] = {'gadget': gadget,
                                                       'count': 1}
                elif gadget.control.ea in good_gadgets:
                    good_gadgets[gadget.control.ea]['count'] += 1
                    break
            count += 1

        if good_gadgets:
            for (ea, info) in good_gadgets.items():
                if info['count'] == count:
                    final_gadgets.append(info['gadget'])
                else:
                    print("Skipping gadget at 0x%X" % ea)
        if final_gadgets:
            self._print_gadgets(gadgets)
        else:
            print("No ROP gadgets found!")

    def summary(self):
        '''
        Prints a summary of your currently marked ROP gadgets, in alphabetical
        order by the marked name. To mark a location as a ROP gadget, simply
        mark the position in IDA (Alt+M) with any name that starts with "ROP".
        '''
        rop_gadgets = self._get_marked_gadgets()
        summaries = []
        delim_char = "-"
        headings = {
            'name': "Gadget Name",
            'offset': "Gadget Offset",
            'summary': "Gadget Summary"
        }
        if self.base != 0:
            headings['offset'] = "Gadget Base + Offset = Address       "
        lengths = {
            'name': len(headings['name']),
            'offset': len(headings['offset']),
            'summary': len(headings['summary']),
        }
        total_length = (3 * len(headings)) + 1

        if rop_gadgets:
            gadget_keys = list(rop_gadgets.keys())
            gadget_keys.sort()

            for marked_comment in gadget_keys:
                if len(marked_comment) > lengths['name']:
                    lengths['name'] = len(marked_comment)

                summary = []
                ea = rop_gadgets[marked_comment]
                end_ea = ea + (self.SEARCH_DEPTH * self.INSIZE)

                while ea <= end_ea:
                    summary.append(idc.GetDisasm(ea))
                    mnem = ida_shims.print_insn_mnem(ea)
                    if len(mnem) > 0 and mnem[0].lower() in ['j', 'b']:
                        summary.append(idc.GetDisasm(ea+self.INSIZE))
                        break

                    ea += self.INSIZE

                if len(summary) == 0:
                    summary.append('')

                for line in summary:
                    if len(line) > lengths['summary']:
                        lengths['summary'] = len(line)

                summaries.append(summary)

            for (heading, size) in lengths.items():
                total_length += size

            delim = delim_char * total_length
            line_fmt = "| %%-%ds | %%-%ds | %%-%ds |" % \
                       (lengths['name'], lengths['offset'], lengths['summary'])

            print(delim)
            print(line_fmt % \
                  (headings['name'], headings['offset'], headings['summary']))
            print(delim)

            for i in range(0, len(gadget_keys)):
                line_count = 0
                marked_comment = gadget_keys[i]
                if self.base != 0:
                    offset = "0x%.8X + 0x%.8X = 0x%.8X" % \
                             (self.base, rop_gadgets[marked_comment],
                              self.base + rop_gadgets[marked_comment])
                else:
                    offset = "0x%.8X" % rop_gadgets[marked_comment]
                summary = summaries[i]

                for line in summary:
                    if line_count == 0:
                        print(line_fmt % (marked_comment, offset, line))
                    else:
                        print(line_fmt % ('', '', line))

                    line_count += 1

                print(delim)

    def build(self):
        '''
        WIP.
        '''
        gadgets = self._get_marked_gadgets()
        bc = BowcasterBuilder(gadgets)
        bc.build_code()

    def set_base(self, base=0):
        '''
        Set base address used for display
        '''
        self.base = base
        self._initial_find()

    def help(self):
        '''
        Show help info.
        '''
        delim = "-" * 140

        print("")
        print("mipsrop.find(instruction_string)")
        print(delim)
        print(self.find.__doc__)

        print("")
        print("mipsrop.system()")
        print(delim)
        print(self.system.__doc__)

        print("")
        print("mipsrop.doubles()")
        print(delim)
        print(self.doubles.__doc__)

        print("")
        print("mipsrop.stackfinders()")
        print(delim)
        print(self.stackfinders.__doc__)

        print("")
        print("mipsrop.tails()")
        print(delim)
        print(self.tails.__doc__)

        print("")
        print("mipsrop.set_base()")
        print(delim)
        print(self.set_base.__doc__)

        print("")
        print("mipsrop.summary()")
        print(delim)
        print(self.summary.__doc__)


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

        def activate(self, ctx):
            global mipsrop
            mipsrop = MIPSROPFinder()
            add_to_namespace('__main__', 'mipsrop', mipsrop)
            return 1

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


class mipsropfinder_t(idaapi.plugin_t):
    flags = 0
    comment = 'MIPS ROP Finder'
    help = ''
    action_name = 'mipsrop:action'
    wanted_name = "MIPS ROP Finder"
    wanted_hotkey = ""
    wanted_tooltip = "MIPS ROP Gadget Finder"
    root_tab = 'Search/'
    menu_name = 'mips rop gadgets'
    menu_context = None

    def init(self):
        if idaapi.IDA_SDK_VERSION >= 700:
            # 199 is a default icon.
            action_desc = idaapi.action_desc_t(self.action_name,
                                               self.menu_name,
                                               MipsRopHandler(),
                                               self.wanted_hotkey,
                                               self.wanted_tooltip,
                                               199)

            idaapi.register_action(action_desc)
            idaapi.attach_action_to_menu(
                self.root_tab, self.action_name, idaapi.SETMENU_APP)
        else:
            self.menu_context = idaapi.add_menu_item(
                "Search/", self.menu_name, "", 0, self.run, (None,))
        return idaapi.PLUGIN_KEEP

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

    def run(self, arg):
        global mipsrop
        mipsrop = MIPSROPFinder()


def PLUGIN_ENTRY():
        return mipsropfinder_t()


repository: matthewrhannon/IDAPython-Automation-Scripts
file: fix_cleanup.py
created_at: 2024-09-24 12:53:57+00:00
pushed_at: 2024-09-24 13:32:30+00:00
description:
Various scripts written in IDAPython to perform various tasks to augment reverse engineering using IDA Pro. Most of the scripts perform some form of automation for signatures, deobfuscation, cleaning sections such as .pdata and RUNTIME_FUNCTION SEH exception handling, LLM Auto-Reversing Application (functions, variables, structures, etc..)

url:
https://github.com/matthewrhannon/IDAPython-Automation-Scripts/blob/cad9d3a5bb5a92ac65736b674a7166348a7ebcc6/fix_cleanup.py

content:
import idc
import ida_range
import ida_ida
import ida_ua

try:
    from ida_idaapi import idaapi
except Exception:
    print("WARNING: NOTICE - from ida_idaapi import idaapi threw exception!")
    import idaapi
    pass
    
import ida_search
import ida_funcs
import ida_bytes
import ida_kernwin
import ida_ida
import ida_nalt
import idautils

Happy = False

print("Script fix_cleanup.py 5/28/2024")

def intro():

    info = idaapi.get_inf_structure()
    filename = idaapi.get_input_file_path()
    entrypoint = info.start_ip
    imagebase = ida_nalt.get_imagebase()
    is_64bits = info.is_64bit()
    is_dll = info.is_dll()
    proc_name = ida_ida.inf_get_procname()
    
    print("")
    print("")
    print("info=" + str(info))
    print("filename=" + str(filename))    
    print("entrypoint=" + str(entrypoint))
    print("imagebase=" + str(imagebase))
    print("is_64bits=" + str(is_64bits))
    print("is_dll=" + str(is_dll))
    print("proc_name=" + str(proc_name))
    print("")
    print("")
    
    return True
    
def wait_for_one_sec():
    idc.qsleep(1000)
    return 1
    
def getInstructionBytes(offset):
        ins = idautils.DecodeInstruction(offset)
        ins_bytes = ida_bytes.get_bytes(offset, ins.size)
        return ins_bytes 

'''

...
...
...

.text:0000000000410108 008 58                                      pop     rax
.text:0000000000410109 000 C3                                      retn
.text:0000000000410109                             ; } // starts at 4100B0
.text:0000000000410109
.text:0000000000410109                             sub_004100B0    endp
.text:0000000000410109
.text:0000000000410109                             ; ---------------------------------------------------------------------------
.text:000000000041010A 66                                          db  66h ; f
.text:000000000041010B 0F                                          db  0Fh
.text:000000000041010C 1F                                          db  1Fh
.text:000000000041010D 44                                          db  44h ; D
.text:000000000041010E 00                                          db    0
.text:000000000041010F 00                                          db    0
.text:0000000000410110                             ; ---------------------------------------------------------------------------
.text:0000000000410110
.text:0000000000410110                             loc_00410110:                           ; CODE XREF: init+49↓p
.text:0000000000410110                                                                     ; DATA XREF: .init_array:funcs_2EEF479↓o
.text:0000000000410110                             ; __unwind {                            ; 2 2 2
.text:0000000000410110 50                                          push    rax
.text:0000000000410111 BF 74 D9 44 04                              mov     edi, offset dword_0444D974 ; 2 2 2 2
.text:0000000000410116 E8 E5 C1 24 00                              call    sub_0065C300
.text:0000000000410116
.text:000000000041011B BF C0 C6 65 00                              mov     edi, offset func ; lpfunc
.text:0000000000410120 BE 74 D9 44 04                              mov     esi, offset dword_0444D974 ; obj
.text:0000000000410125 BA 98 C9 35 04                              mov     edx, offset qword_0435C998 ; lpdso_handle
.text:000000000041012A E8 61 58 FF FF                              call    ___cxa_atexit
.text:000000000041012A
.text:000000000041012F C7 05 3F D8 03 04 00 00                     mov     cs:dword_0444D978, 0
.text:000000000041012F 00 00
.text:0000000000410139 C7 05 39 D8 03 04 00 00                     mov     cs:dword_0444D97C, 0
.text:0000000000410139 00 00
.text:0000000000410143 C7 05 33 D8 03 04 00 00                     mov     cs:dword_0444D980, 0
.text:0000000000410143 00 00
.text:000000000041014D 48 B8 00 00 80 3F 00 00                     mov     rax, 3F8000003F800000h
.text:000000000041014D 80 3F
.text:0000000000410157 48 89 05 26 D8 03 04                        mov     cs:qword_0444D984, rax
.text:000000000041015E C7 05 24 D8 03 04 00 00                     mov     cs:dword_0444D98C, 3F800000h ; <suspicious>
.text:000000000041015E 80 3F                         ; <suspicious>
.text:0000000000410168 58                                          pop     rax
.text:0000000000410169 C3                                          retn
.text:0000000000410169                             ; } // starts at 410110
.text:0000000000410169
.text:0000000000410169                             ; ---------------------------------------------------------------------------
.text:000000000041016A 66                                          db  66h ; f
.text:000000000041016B 0F                                          db  0Fh
.text:000000000041016C 1F                                          db  1Fh
.text:000000000041016D 44                                          db  44h ; D
.text:000000000041016E 00                                          db    0
.text:000000000041016F 00                                          db    0
.text:0000000000410170                             ; ---------------------------------------------------------------------------
.text:0000000000410170
.text:0000000000410170                             loc_00410170:                           ; CODE XREF: init+49↓p
.text:0000000000410170                                                                     ; DATA XREF: .init_array:funcs_2EEF479↓o
.text:0000000000410170                             ; __unwind {                            ; 2 2 2
.text:0000000000410170 50                                          push    rax
.text:0000000000410171 BF 90 D9 44 04                              mov     edi, offset qword_0444D990 ; 2 2 2 2
.text:0000000000410176 E8 85 C1 24 00                              call    sub_0065C300
.text:0000000000410176
.text:000000000041017B BF C0 C6 65 00                              mov     edi, offset func ; lpfunc
.text:0000000000410180 BE 90 D9 44 04                              mov     esi, offset qword_0444D990 ; obj
.text:0000000000410185 BA 98 C9 35 04                              mov     edx, offset qword_0435C998 ; lpdso_handle
.text:000000000041018A E8 01 58 FF FF                              call    ___cxa_atexit
.text:000000000041018A
.text:000000000041018F 58                                          pop     rax
.text:0000000000410190 C3                                          retn
.text:0000000000410190                             ; } // starts at 410170
.text:0000000000410190
.text:0000000000410190                             ; ---------------------------------------------------------------------------
.text:0000000000410191 66                                          db  66h ; f
.text:0000000000410192 2E                                          db  2Eh ; .
.text:0000000000410193 0F                                          db  0Fh
.text:0000000000410194 1F                                          db  1Fh
.text:0000000000410195 84                                          db  84h
.text:0000000000410196 00                                          db    0
.text:0000000000410197 00                                          db    0
.text:0000000000410198 00                                          db    0
.text:0000000000410199 00                                          db    0
.text:000000000041019A 00                                          db    0
.text:000000000041019B 0F                                          db  0Fh
.text:000000000041019C 1F                                          db  1Fh
.text:000000000041019D 44                                          db  44h ; D
.text:000000000041019E 00                                          db    0
.text:000000000041019F 00                                          db    0
.text:00000000004101A0                             ; ---------------------------------------------------------------------------
.text:00000000004101A0
.text:00000000004101A0                             loc_004101A0:                           ; CODE XREF: init+49↓p
.text:00000000004101A0                                                                     ; DATA XREF: .init_array:funcs_2EEF479↓o
.text:00000000004101A0                             ; __unwind {                            ; 2 2 2
.text:00000000004101A0 50                                          push    rax
.text:00000000004101A1 BF A0 D9 44 04                              mov     edi, offset off_0444D9A0 ; 2 2 2 2
.text:00000000004101A6 E8 55 C1 24 00                              call    sub_0065C300

...
...
...
'''


def fix_cleanup_funcs():

    #current_ea = 0
    count = 0
    
    #.text:00000000004117A0                             sub_004117A0    proc near               ; CODE XREF: init+49↓p                  <--- reasonable place to start testing and stopping at ed of bad section for testing
    start_effective_address = 0x4117A0
    
    #.text:0000000000432920                             sub_00432920    proc near               ; CODE XREF: init+49↓p                  <--- good place to stop early testing after major back block
    stop_effective_address = 0x432920
    
    #Flow:
    #-iterate downward starting at our effective_start_addresss (for now)
    #    -are you already apart of a fuction?
    #       -is the function you are apart of valid?
    #          -if not - destroy function and undo dissasembly
    #       -it is a valid function? - cool move on
    #    -you say you are NOT apart of a function?
    #       -okay... lets see if we can find you a function to be apart of by:
    #           -iterate upwards until you find something with xrefs (naive with many assumptions, but I think it will work in this special case)
    #           -is this place a good and valid place to create a fuction assumming it is not already a function (in that case - something went wrong - and for now log the error and move onn past the instructionn we were trying to find a home
    #           -if it is a good place to make a function - make a function there (and manaual instructions if required) - verify we found our friend a nice home and if so move on. If not - log it - and move on for now... this might be the policy we follow ultimately for success
    #    -iterate to next instruction (assemblying one if there isn't one)
    #- ??? succcss ???
    
    curr_addr_under_examination = start_effective_address
    
    #FIXME: make sure we are only processsing executable sections
    while curr_addr_under_examination < stop_effective_address:
        
        ea = here()
        
        #Make sure we have a valid instruction
        ea_flags = ida_bytes.get_flags(ea)

        insn_len = 0
        if not ida_bytes.is_code(ea_flags):
            insn_len = ida_ua.create_insn(ea)
            print("[DEBUG] Created in structionn at ea=" + hex(ea) + " and length=" + hex(insn_len))
        
        #Check if we are in a function
        curr_possible_func = ida_funcs.get_func(ea)
        if curr_possible_func:
            #Make sure we are at the start of fuction before we carry on  with the sanity logic right below this sentence
            if ida_bytes.is_func(ea) == True:
                if int(ea, 16) % 0x08 != 0:
                    print("[WARNING] Found bad function @ " + hex(ea))
                   
                    start_ea = idc.get_func_attr(curr_possible_func, idc.FUNCATTR_START)
                    end_ea = idc.get_func_attr(curr_possible_func, idc.FUNCATTR_END)
                
                    func_length = end_ea - start_ea
                    
                    ida_funcs.del_func(ea)
                    ida_bytes.del_items(ea, 0, func_length)
                    
                    continue
        else:
            #So you say you are not apart of a function
            
            

        for function_ea in idautils.Functions():
        
            if ida_kernwin.user_cancelled():
                print("WARNING: ida_kernwin.user_cancelled exit route!")
                Happy = False
                break
                
            start_ea = idc.get_func_attr(function_ea, idc.FUNCATTR_START)
            end_ea = idc.get_func_attr(function_ea, idc.FUNCATTR_END)
            
            func_length = end_ea - start_ea
           
            #current_ea = start_ea
            
            print("start_ea=" + hex(start_ea) + " --> end_ea=" + hex(end_ea) + " size=" + str(func_length))
            
            #start_ea = idc.to_ea(start_ea)
            ida_kernwin.jumpto(start_ea)
            
            #sanitycheck
            #if start_ea % 0x08 != 0:
            if start_ea % 0x04 != 0:
                print("[SANITY FAILURE] start_ea=" + hex(start_ea)+ " and end_ea=" + hex(end_ea) + "... continuing ... and destroying instructions inbetween...")
                num_bytes_to_un_encode = func_length
                
                ida_funcs.del_func(start_ea)
                ida_bytes.del_items(start_ea, 0, num_bytes_to_un_encode)
                
                continue
                
            #ida_funcs.del_func(function_ea)
            #ida_ua.create_insn(insn, start_ea)

            ida_funcs.del_func(function_ea)
            ida_ua.create_insn(start_ea)
            ida_funcs.add_func(start_ea)
            
            if not ida_bytes.has_xref(end_ea):
                if ida_bytes.is_func(end_ea):
                    #ida_kernwin.jumpto(end_ea)
                    ida_bytes.del_items(end_ea, 0, 1)
                    #insn = ida_ua.insn_t()
                    ida_funcs.del_func(end_ea)
            elif ida_bytes.has_xref(end_ea):
                ida_ua.create_insn(end_ea)
                ida_funcs.add_func(end_ea)
            else:
                print("!!! WTF !!! end_ea=" + hex(end_ea))
                
            #current_ea = idc.next_head(current_ea)
            
            if count % 500 == 0:
                print("count % 500 == 0! count=" + str(count)) #+ "/" + str(len(idautils.Functions)))
            
            count = count + 1

        Happy = True
        return count
    
    
def PLUGIN_ENTRY():

    intro()
    ida_kernwin.show_wait_box("Processing")
    
    try:
        print("Number functions fixed = " + str(fix_cleanup_funcs()))
    except Exception as e:
        print("WARNING: exception happened! str(e)=" + str(e))
        goto .bar_begin
    finally:
        if Happy == False:
            print(hex(ea) + ": " + str(name))
            print("WARNING: finally exit route from fix_cleanup.py!")
            ida_kernwin.hide_wait_box()
            return False
        else:
            print("SUCCESS: HAPPY fix_cleanup.py run!")
            ida_kernwin.hide_wait_box()
            return True
            
    print("!!! WTF !!!")
    return False


repository: Comsecuris/ida_strcluster
file: strcluster.py
description: extending IDA's string navigation capabilities
wanted_name: Comsecuris StringCluster
created_at: 2017-07-19 16:23:32+00:00
pushed_at: 2020-10-25 09:30:04+00:00
url:
https://github.com/Comsecuris/ida_strcluster/blob/ba29eb55ee82f168d5d2ef3a5ae17ed421e3ca55/strcluster.py

content:
# (C) Copyright 2017 Comsecuris UG
# --------------------------------
# This plugin extends IDA Pro's capabilities to display strings within the binary
# by clustering found strings on a per-function basis. Strings that don't belong
# to a function are grouped into the 0_sub pseudo-function. This allows to quickly
# identify intersting functionality as well as strings that are not part of a function
# and quickly navigate through the results and filter them, hopefully making manual
# analysis more effectively. Requires IDA Pro >= 6.9.

import re

import idautils
import idc
import idaapi
from idaapi import PluginForm

import time

from PyQt5 import QtCore
from PyQt5 import QtWidgets
from PyQt5 import QtGui

# Change this flag if you prefer to not filter while you type.  I personally don't
# like this as it removes results that I may look at while creating the next
# input, but also typing becomes slow on a large number of strings.
LIVE_SEARCH = True

DEBUG          = False
PROFILE        = False
NO_FUNC        = '0_sub'
MATCH_COLOR    = QtGui.QColor(0xC7,0xF2, 0xCF)
NOMATCH_COLOR  = QtGui.QColor('white')

FUN_COLUMN     = 0
XREF_COLUMN    = 1
STR_COLUMN     = 2

if PROFILE == True:
	start_time = time.time()
	import cProfile, pstats, StringIO
	pr = cProfile.Profile()
	pr.enable()

def dprint(s):
	if DEBUG == True:
		print(s)

class StringItem(QtGui.QStandardItem):
	def __init__(self, s, ea):
		super(QtGui.QStandardItem, self).__init__(s)
		self.ea = ea

class IdaString():
	def __init__(self, s, ea, xref):
		self.s = s
		self.ea = ea
		self.xref = xref

class StringClusterMap(PluginForm):
	def getIcon(self):
		icon = (
			'0000010001001010000001002000680400001600000028000000100000002000000001'
			'0020000000000000040000130b0000130b00000000000000000000000000000773e600'
			'0577df000877dd060477df490277e0a70277e0e30277e0fb0277e0fb0277e0e30277e0'
			'a70377e0490577e0060377e1000175f00000000000000000000377e0000577df180377'
			'e0920277e0ed0177e0ff0177e0ff0177e0ff0177e0ff0177e0ff0177e0ff0277e0ed02'
			'77e0920377e1180277e100000000000577df000577df180277e0b10177e0ff0177e0ff'
			'0177e0ff0177e0ff0177e0ff0177e0ff0177e0ff0177e0ff0177e0ff0177e0ff0277e0'
			'b10377e1180377e1000174de070176df920076e0ff0077e0ff0177e0ff0177e0ff0177'
			'e0ff0076e0ff0076e0ff0177e0ff0177e0ff0177e0ff0077e0ff0076e0ff0176e09202'
			'76e107127fe0481983e2ee1f87e3ff0d7de1ff0076e0ff0077e0ff067ae1ff1d86e3ff'
			'1a84e3ff077ae1ff0177e0ff0076e0ff0e7ee1ff1e87e3ff1581e2ee0a7be1483592e4'
			'a759a6e9ff4fa0e8ff66adeaff1e86e3ff0b7ce1ff60a9e9ff57a5e9ff56a4e9ff459b'
			'e7ff0277e0ff288ce4ff68aeeaff51a1e8ff56a4e8ff2389e3a70578e0e40177e0ff00'
			'72dfff499de7ff4fa1e8ff3c96e6ff53a3e8ff0074dfff0075e0ff0579e0ff0478e0ff'
			'6cb0ebff268be4ff0075e0ff0378e0ff0478e0e40176e0fb1481e2ff439ae7ff7bb8ec'
			'ff2a8ce4ff5da8eaff63abeaff3793e5ff3894e6ff3392e5ff1481e2ff73b4edff0b7c'
			'e1ff0177e0ff0177e0ff0277e0fb2c8de4fb76b5ecff50a1e8ff1e86e3ff0075e0ff59'
			'a6e9ff63abeaff3692e5ff3793e5ff76b5ecff2389e3ff73b4ecff0d7de1ff0077e0ff'
			'0077e0ff0177e0fb5ea8e9e455a4e8ff0075e0ff0b7ce1ff0679e0ff3090e5ff59a6e9'
			'ff0377e0ff1380e2ff62abe9ff0c7ce1ff65aceaff3693e5ff0478e0ff0d7de1ff077a'
			'e1e42489e2a75ba7e9ff59a5e9ff5aa6e9ff1983e2ff0578e0ff489de7ff5da8eaff5f'
			'a9eaff2c8ee4ff0075e0ff1a84e2ff60aaeaff5ba6e9ff57a4e9ff1f86e3a70075df48'
			'0478e0ee0e7ee1ff087be1ff0177e0ff0177e0ff0177e0ff0c7de1ff087be1ff0076e0'
			'ff0177e0ff0076e0ff0479e0ff0e7ee1ff087ae1ee0377e0480777de070377e0920177'
			'e0ff0177e0ff0177e0ff0177e0ff0177e0ff0177e0ff0177e0ff0177e0ff0177e0ff01'
			'77e0ff0177e0ff0177e0ff0277e0920477e1070577df000577df180277e0b10177e0ff'
			'0177e0ff0177e0ff0177e0ff0177e0ff0177e0ff0177e0ff0177e0ff0177e0ff0177e0'
			'ff0277e0b10377e1180377e100000000000377e0000577df180377e0920277e0ed0177'
			'e0ff0177e0ff0177e0ff0177e0ff0177e0ff0177e0ff0277e0ed0277e0920377e11802'
			'77e10000000000000000000773e6000577df000877dd060477df490277e0a70277e0e3'
			'0277e0fb0277e0fb0277e0e30277e0a70377e0490676e0060377e1000174f300000000'
			'00e0070000c00300008001000000000000000000000000000000000000000000000000'
			'00000000000000000000000000000000000080010000c0030000e0070000')
		image = QtGui.QImage()
		image.loadFromData(QtCore.QByteArray.fromHex(icon))
		pixmap = QtGui.QPixmap()
		pixmap.convertFromImage(image)
		return QtGui.QIcon(pixmap)

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

	def xrefsTo(self, ea):
		s = []
		map(lambda x: s.append(x.frm), idautils.XrefsTo(ea))
		return s
	
	def funXrefs(self):
		res = {}
		for s in idautils.Strings():
			s_ea = s.ea
			s_v = str(s).rstrip()
			dprint("checking %x - %s" %(s_ea, s_v))
			s_xrefs_eas = self.xrefsTo(s_ea)
			if not s_xrefs_eas:
				dprint("no xref found for %s" %(s_v))
				s_xrefs_eas = [s_ea]
	
			# same string can be xref'ed by more than one function
			for fs_ea in s_xrefs_eas:
				dprint("looking for function of %x" %(fs_ea))

				f_name = idc.GetFunctionName(fs_ea)
				f_ea = idc.GetFunctionAttr(fs_ea, idc.FUNCATTR_START)
				if not f_name or f_name == '': f_name = NO_FUNC
				if f_ea in res:
					res[f_ea]['strings'][s_v] = IdaString(s_v, s_ea, fs_ea)
				else:
					res[f_ea] = dict({
						'name' : f_name,
						'strings' : { s_v : IdaString(s_v, s_ea, fs_ea) }
					})

		return res

	def hideItem(self, item, search_text):
		if search_text ==  "":
			item.setBackground(NOMATCH_COLOR)
			return False

		if search_text != "" and self.filter_regex != None:
			res = self.filter_regex.search(item.text())
			if res:
				item.setBackground(MATCH_COLOR)
				return False
			else:
				item.setBackground(NOMATCH_COLOR)
				return True

		if search_text != "" and search_text.lower() in item.text().lower():
			item.setBackground(MATCH_COLOR)
			return False
		else:
			item.setBackground(NOMATCH_COLOR)
			return True

	def liveSearchCheckBox(self, event):
		self.live_search = not self.live_search

	def checkBoxEvent(self, event):
		self.filterEvent()

	def filterEvent(self, event = None):
		if event != None:
			QtWidgets.QLineEdit.keyReleaseEvent(self.filter_line, event)

		if event and (self.live_search == False and event.key() != QtCore.Qt.Key_Enter and event.key() != QtCore.Qt.Key_Return):
			return

		search_text = self.filter_line.text()

		if search_text != "" and self.regexckb.isChecked():
			self.filter_regex = re.compile(search_text)
		else:
			self.filter_regex = None

		res_strs = 0
		for idx in xrange(self.model.rowCount()):
			item = self.model.item(idx)
			n_strings = item.rowCount()
			hide = True if search_text != "" else False

			for c_idx in xrange(n_strings):
				hide_row = self.hidecheckb.isChecked()
				c_item1 = item.child(c_idx, XREF_COLUMN)
				c_item2 = item.child(c_idx, STR_COLUMN)

				for i in [item, c_item1, c_item2]:
					if not self.hideItem(i, search_text):
						hide = False
						hide_row = False
						if i is c_item2: # we only want to count strs
							res_strs += 1

				self.view.setRowHidden(c_idx, self.model.indexFromItem(item), hide_row)

			self.view.setRowHidden(idx, QtCore.QModelIndex(), hide)
			if item.text().startswith(NO_FUNC):
				if self.hidenscheckb.isChecked():
					self.view.collapse(self.model.indexFromItem(item))
				else:
					self.view.expand(self.model.indexFromItem(item))

		self.results.setText('%d strings' %(res_strs))

	def doubleClickEvent(self, event):
		index = event.pos()
		try:
			item = self.model.itemFromIndex(self.view.indexAt(event.pos()))
			column = item.column()
			ea = item.ea
		except:
			return

		if ea != -1:
			idc.Jump(ea)

	def PopulateForm(self):
		self.parent.setWindowIcon(self.getIcon())
		# gather data
		if not self.items:
			self.items = self.funXrefs()

		if PROFILE == True:
			print("--- %s seconds ---" % (time.time() - start_time))

		# create layout
		self.layout = QtWidgets.QVBoxLayout()
		self.layout.setContentsMargins(0,0,0,0)
		self.layout.setSpacing(0)
		self.view = QtWidgets.QTreeView()
		self.view.setSelectionBehavior(QtWidgets.QAbstractItemView.SelectRows)

		self.model = QtGui.QStandardItemModel()
		self.model.setHorizontalHeaderLabels(['Function', 'Xref EA', 'String'])

		self.view.setModel(self.model)
		self.view.setUniformRowHeights(True)
		res_strs = 0
		parent = None
		for idx, (f_ea, fi) in enumerate(self.items.iteritems()):
			fun_ea = f_ea
			fun_name = fi['name']
			fun_strs = len(fi['strings'])
			res_strs += fun_strs

			parent = StringItem('%s (%d)' %(fun_name, fun_strs), fun_ea)
			parent.setEditable(False)

			for (s_v, s) in fi['strings'].iteritems():
				ilist = [StringItem('', -1), StringItem('%x' %(s.xref), s.xref), StringItem(s_v, s.ea)]
				for i in ilist: i.setEditable(False)
				parent.appendRow(ilist)

			self.model.appendRow(parent)
			# make sure the first column is as large as the function name
			#self.view.setFirstColumnSpanned(idx, self.view.rootIndex(), True)

		index = self.model.indexFromItem(parent)
		self.view.expandAll()
		self.view.resizeColumnToContents(0)

		self.layout.addWidget(self.view)

		self.filterbox = QtWidgets.QGroupBox()
		self.ctrlbox = QtWidgets.QGroupBox()
		self.filterlayout = QtWidgets.QGridLayout()
		self.ctrllayout = QtWidgets.QHBoxLayout()
		self.ctrllayout.setContentsMargins(0,0,0,0)
		self.filterlayout.setContentsMargins(0,1,0,0)

		self.filter_line = QtWidgets.QLineEdit() #FilterLine()
		self.results = QtWidgets.QLabel('%d strings' %(res_strs))
		self.hidecheckb = QtWidgets.QCheckBox('Hide no match')
		self.hidenscheckb = QtWidgets.QCheckBox('Collapse ' + NO_FUNC)
		self.regexckb = QtWidgets.QCheckBox('Regex')
		self.live_searchcb = QtWidgets.QCheckBox('Live search')
		self.hidecheckb.setChecked(True)
		self.hidenscheckb.setChecked(False)
		self.regexckb.setChecked(False)
		self.live_searchcb.setChecked(LIVE_SEARCH)
		self.live_search = LIVE_SEARCH
		self.filter_regex = None

		self.live_searchcb.stateChanged.connect(self.liveSearchCheckBox)
		self.hidecheckb.stateChanged.connect(self.checkBoxEvent)
		self.hidenscheckb.stateChanged.connect(self.checkBoxEvent)
		self.regexckb.stateChanged.connect(self.checkBoxEvent)
		self.filter_line.keyReleaseEvent = self.filterEvent
		self.view.mouseDoubleClickEvent = self.doubleClickEvent
		self.filterlayout.addWidget(self.filter_line, 1, 1)
		self.filterlayout.addWidget(self.results, 1, 2)

		self.ctrllayout.addWidget(self.hidecheckb)
		self.ctrllayout.addWidget(self.regexckb)
		self.ctrllayout.addWidget(self.hidenscheckb)
		self.ctrllayout.addWidget(self.live_searchcb)
		self.ctrllayout.addStretch()

		self.layout.addWidget(self.filterbox)
		self.layout.addWidget(self.ctrlbox)
		self.parent.setLayout(self.layout)
		self.filterbox.setLayout(self.filterlayout)
		self.ctrlbox.setLayout(self.ctrllayout)
		# this is mostly useful for the first column, which forms the parent.
		# other columns can also be sorted, but due to the tree nature remain
		# sorted within their respective parents, which can be confusing at first.
		self.view.setSortingEnabled(True)

		w_width = self.parent.size().width()
		c0_width = self.view.columnWidth(0)
		max_width = 20 * w_width / 100
		if c0_width >= max_width:
			self.view.setColumnWidth(0, max_width)

		if PROFILE == True:
			print("--- %s seconds ---" % (time.time() - start_time))
			s = StringIO.StringIO()
			sortby = 'cumulative'
			ps = pstats.Stats(pr, stream=s).sort_stats(sortby)
			ps.print_stats()
			print(s.getvalue())

	def OnClose(self, form):
		pass

class ida_string_cluster_plugin(idaapi.plugin_t):
	flags = idaapi.PLUGIN_OK
	comment = ""
	help = ""
	wanted_name = "Comsecuris StringCluster"
	wanted_hotkey = "Alt-s"

	def init(self):
		return idaapi.PLUGIN_OK
	
	def run(self, arg):
		plg = StringClusterMap()
		plg.Show("StringCluster")
		return

	def term(self):
		pass


def PLUGIN_ENTRY():
	return ida_string_cluster_plugin()


repository: czs108/Microsoft-Malware-Classification
file: ida-plugin/classifier.py
wanted_name: 2015 Microsoft Malware Classifier
created_at: 2021-07-17 20:11:32+00:00
pushed_at: 2024-02-13 19:35:41+00:00
description:
🔍 "2015 Microsoft Malware Classification Challenge" - Using machine learning to classify malware into different families based on Windows PE structures, disassembly scripts and machine code.（“2015微软恶意软件分类挑战赛” - 使用机器学习基于Windows PE结构、反汇编和机器代码对恶意软件进行分类。）

url:
https://github.com/czs108/Microsoft-Malware-Classification/blob/5a6343b727ca2d946e6ebf9460e130039a3b8d8a/ida-plugin/classifier.py

content:
"""
-- INTRODUCTION --
An IDA Pro plug-in to classify a malware into nine families defined by the 2015 Microsoft Malware Classification Challenge.

-- USAGE --
This script must be located in the same folder as `generator.py`.

-- HELP --
See: https://github.com/czs108/Microsoft-Malware-Classification
"""

from pathlib import Path
import zlib
import joblib
import re

import pandas as pd
import numpy as np

import idaapi
import ida_loader
import ida_diskio

from generator import Generator


class Reader:
    def __init__(self, file: str) -> None:
        self._file: str = file

    def bytes(self) -> list[str]:
        with Path(self._file + ".bytes").open() as file:
            data = file.read()
        items = data.split()
        byte_list = []
        for item in items:
            if len(item) == 2 and item != "??":
                byte_list.append(item)
        return byte_list

    def asm(self) -> str:
        with Path(self._file + ".asm").open(encoding="utf-8", errors="ignore") as file:
            return file.read()

    def asm_lines(self) -> list[str]:
        with Path(self._file + ".asm").open(encoding="utf-8", errors="ignore") as file:
            return file.readlines()


class Section:
    _sctn_rgx: re.Pattern = re.compile(r"([A-Za-z]+):[\dA-F]{8}\s+;\s+Virtual\s+size\s+:\s+[\dA-F]{8}\s+\(\s*(\d+)\.[\s\S]+?Section\s+size\s+in\s+file\s+:\s+[\dA-F]{8}\s+\(\s*(\d+)\.[\s\S]+?;\s+Flags\s+[\dA-F]{8}:[\s\S]+?\s+((?:Executable|Readable|Writable|\s)+)", flags=re.IGNORECASE)

    def __init__(self) -> None:
        self.name: str = ""
        self.virtual_size: int = 0
        self.raw_size: int = 0
        self.executable: bool = False
        self.writable: bool = False

    @classmethod
    def load(cls, asm: str) -> list[Section]:
        sctns = []
        for attr in cls._sctn_rgx.findall(asm):
            sctn = Section()
            sctn.name, sctn.virtual_size, sctn.raw_size, access = \
                attr[0].lower(), int(attr[1]), int(attr[2]), set(attr[3].split())
            if "Executable" in access:
                sctn.executable = True
            if "Writable" in access:
                sctn.writable = True
            sctns.append(sctn)
        return sctns


class Extractor:
    _sctn_cols: tuple[str] = ("idata-Raw", "text-Raw", "data-Virtual", "data-Ratio", "text-Virtual",
                                "data-Raw", "idata-Virtual", "text-Ratio", "idata-Ratio", "bss-Virtual",
                                "bss-Ratio", "rdata-Ratio", "rdata-Virtual", "rdata-Raw", "bss-Raw",
                                "code-Raw", "code-Ratio", "tls-Raw", "code-Virtual", "tls-Ratio",
                                "tls-Virtual", "rata-Raw", "rata-Virtual", "rata-Ratio", "zenc-Ratio")

    def __init__(self, file: str) -> None:
        self._file: str = file
        self._id: str = Path(file).name

    def extract(self) -> pd.DataFrame:
        size, rwe = self._section_attributions()
        selected_size = pd.DataFrame(index=size.index, columns=self._sctn_cols)
        for col in size.columns:
            if col in selected_size.columns:
                selected_size.at[self._id, col] = size.at[self._id, col]
        return pd.concat([selected_size.fillna(0), rwe, self._content_complexity()], axis="columns")

    def _content_complexity(self) -> pd.DataFrame:
        df = pd.DataFrame(
            columns=["ID", "Asm-Len", "Zip-Asm-Len", "Asm-Zip-Ratio", "Byte-Len", "Zip-Byte-Len", "Byte-Zip-Ratio"],
            dtype=float).set_index("ID")

        reader = Reader(self._file)
        asm = reader.asm().encode("utf-8")
        bytes = " ".join([str(byte) for byte in reader.bytes()]).encode("utf-8")
        df.at[self._id, "Asm-Len"] = len(asm)
        df.at[self._id, "Zip-Asm-Len"] = len(zlib.compress(asm))
        df.at[self._id, "Byte-Len"] = len(bytes)
        df.at[self._id, "Zip-Byte-Len"] = len(zlib.compress(bytes))

        df[["Asm-Len", "Zip-Asm-Len", "Byte-Len", "Zip-Byte-Len"]] = \
            df[["Asm-Len", "Zip-Asm-Len", "Byte-Len", "Zip-Byte-Len"]].astype(int)
        df["Asm-Zip-Ratio"] = (df["Asm-Len"] / df["Zip-Asm-Len"]).round(5)
        df["Byte-Zip-Ratio"] = (df["Byte-Len"] / df["Zip-Byte-Len"]).round(5)
        return df

    def _section_attributions(self) -> tuple[pd.DataFrame, pd.DataFrame]:
        rwe_cols = [f"{i}-{j}" for i in ["Executable", "Writable", "Readable"] for j in ["Virtual", "Raw", "Ratio"]]
        rwe = pd.DataFrame(columns=["ID"] + rwe_cols, dtype=float).set_index("ID")
        sctn = pd.DataFrame(columns=["ID"], dtype=float).set_index("ID")

        sctn_int_cols = set()

        def save_by_sections(attr: Section) -> None:
            sctn.at[self._id, f"{attr.name}-Virtual"] = attr.virtual_size
            sctn.at[self._id, f"{attr.name}-Raw",] = attr.raw_size
            sctn.at[self._id, f"{attr.name}-Ratio",] = round(attr.raw_size / attr.virtual_size, 5)
            sctn_int_cols.update([f"{attr.name}-Virtual", f"{attr.name}-Raw"])

        def save_by_access(attr: Section) -> None:
            rwe.at[self._id, "Readable-Virtual"] += attr.virtual_size
            rwe.at[self._id, "Readable-Raw"] += attr.raw_size
            if attr.writable:
                rwe.at[self._id, "Writable-Virtual"] += attr.virtual_size
                rwe.at[self._id, "Writable-Raw"] += attr.raw_size
            if attr.executable:
                rwe.at[self._id, "Executable-Virtual"] += attr.virtual_size
                rwe.at[self._id, "Executable-Raw"] += attr.raw_size

        rwe.loc[self._id, :] = 0
        sctn.loc[self._id, :] = 0
        for attr in Section.load(Reader(self._file).asm()):
            save_by_sections(attr)
            save_by_access(attr)

        sctn.fillna(0, inplace=True)
        sctn_int_cols = list(sctn_int_cols)
        sctn[sctn_int_cols] = sctn[sctn_int_cols].astype(int)

        rwe.fillna(0, inplace=True)
        rwe_int_cols = [i for i in rwe.columns if "-Virtual" in i or "-Raw" in i]
        rwe[rwe_int_cols] = rwe[rwe_int_cols].astype(int)

        rwe["Readable-Ratio"] = (rwe["Readable-Raw"] / rwe["Readable-Virtual"]).round(5)
        rwe["Writable-Ratio"] = (rwe["Writable-Raw"] / rwe["Writable-Virtual"]).round(5)
        rwe["Executable-Ratio"] = (rwe["Executable-Raw"] / rwe["Executable-Virtual"]).round(5)
        return sctn, rwe.fillna(0)


class MalwareClassifier(idaapi.plugin_t):
    flags: int = idaapi.PLUGIN_UNL
    comment: str = "This plug-in can calculate the probability that it belongs to each malware family defined by the 2015 Microsoft Malware Classification Challenge."
    help: str = "See: https://github.com/czs108/Microsoft-Malware-Classification"
    wanted_name: str = "2015 Microsoft Malware Classifier"
    wanted_hotkey: str = "Alt-F8"

    def init(self) -> int:
        self._file: str = ida_loader.get_path(ida_loader.PATH_TYPE_IDB)[:-len(".idb")]
        self._classes: tuple[str] = \
            ("Ramnit", "Lollipop", "Kelihos_ver3", "Vundo", "Simda", "Tracur", "Kelihos_ver1", "Obfuscator.ACY", "Gatak")
        return idaapi.PLUGIN_OK

    def run(self, arg: int) -> None:
        self._gene_data()
        plugin_folder = Path(ida_diskio.idadir("plugins"))
        X = Extractor(self._file).extract()
        classifier = joblib.load(plugin_folder.joinpath("tpot_exported_ida_clf.joblib"))
        proba = classifier.predict_proba(X)[0]
        for i in np.argsort(proba)[::-1]:
            idaapi.msg("{:04.2f} -> {}\n".format(proba[i], self._classes[i]))

    def term(self) -> None:
        pass

    def _gene_data(self) -> None:
        generator = Generator(self._file)
        generator.asm()
        generator.bytes()


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

comment:
This plug-in can calculate the probability that it belongs to each malware family defined by the 2015 Microsoft Malware Classification Challenge.


repository: hoangcuongflp/Stuff
file: ida-master/plugins/mipslocalvars/mipslocalvars.py
description: because I've orderliness obsession :(
wanted_name: Names MIPS registers saved on the stack
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/mipslocalvars/mipslocalvars.py

content:
# IDA plugin to name stack variables that are simply used to store register values until a function returns ($ra, $s0-$s7, $fp, $gp).
#
# Invoke by going to Options->Name saved registers, or by using the Alt+4 hotkey.
#
# Craig Heffner
# Tactical Network Solutions

import idc
import idaapi
import idautils

class NameMIPSSavedRegisters(object):

	INSIZE = 4
	SEARCH_DEPTH = 25

	ARCH = {
			'arguments'	: ['$a0', '$a1', '$a2', '$a3'],
			'savedregs'	: ['$s0', '$s1', '$s2', '$s3', '$s4', '$s5', '$s6', '$s7', '$fp', '$gp', '$ra'],
	}

	def __init__(self):
		print "Naming saved register locations...",

		for ea in idautils.Functions():
			mea = ea
			named_regs = []
			last_iteration = False

			while mea < (ea + (self.INSIZE * self.SEARCH_DEPTH)):
				mnem = idc.GetMnem(mea)

				if mnem in ['sw', 'sd']:
					reg = idc.GetOpnd(mea, 0)
					dst = idc.GetOpnd(mea, 1)
	
					if reg in self.ARCH['savedregs'] and reg not in named_regs and dst.endswith('($sp)') and 'var_' in dst:
						offset = int(dst.split('var_')[1].split('(')[0], 16)
						idc.MakeLocal(ea, idc.FindFuncEnd(ea), "[sp-%d]" % offset, "saved_%s" % reg[1:])
						named_regs.append(reg)
				
				if last_iteration:
					break
				elif mnem.startswith('j') or mnem.startswith('b'):
					last_iteration = True

				mea += self.INSIZE

		print "done."


class mips_saved_registers_t(idaapi.plugin_t):
	flags = 0
	comment = ""
	help = ""
	wanted_name = "Names MIPS registers saved on the stack"
	wanted_hotkey = ""

	def init(self):
		self.menu_context = idaapi.add_menu_item("Options/", "Name saved registers", "Alt-4", 0, self.name_saved_registers, (None,))
		return idaapi.PLUGIN_KEEP

	def term(self):
		idaapi.del_menu_item(self.menu_context)
		return None

	def run(self, arg):
		pass

	def name_saved_registers(self, arg):
		NameMIPSSavedRegisters()

def PLUGIN_ENTRY():
	return mips_saved_registers_t()


repository: FBLeee/IDA_GPT
file: Model_xunfei/Gepetto_free.py
description: 集成GPT或者deepseek免费API到IDA插件中，解释IDA反汇编代码，提高逆向效率，感谢免费开源项目。
wanted_name: Gepetto
created_at: 2023-04-26 07:44:04+00:00
pushed_at: 2025-03-05 00:54:34+00:00
url:
https://github.com/FBLeee/IDA_GPT/blob/366a2ecdc53744b444ddde75ddb7d9e76743e656/Model_xunfei/Gepetto_free.py

content:
import functools
import json
import os
import re
import textwrap
import threading
import gettext

import requests
import sys
import urllib3
import time

import idaapi
import ida_hexrays
import ida_kernwin
import idc
# import openai

import SparkApi




#以下密钥信息从控制台获取
appid = "XXXXXXXX"     #填写控制台中获取的 APPID 信息
api_secret = "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX"   #填写控制台中获取的 APISecret 信息
api_key ="XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX"    #填写控制台中获取的 APIKey 信息

#用于配置大模型版本，默认“general/generalv2”
domain = "general"   # v1.5版本
# domain = "generalv2"    # v2.0版本
#云端环境的服务地址
Spark_url = "ws://spark-api.xf-yun.com/v1.1/chat"  # v1.5环境的地址
# Spark_url = "ws://spark-api.xf-yun.com/v2.1/chat"  # v2.0环境的地址










# =============================================================================
# EDIT VARIABLES IN THIS SECTION
# =============================================================================


# Specify the program language. It can be "fr_FR", "zh_CN", or any folder in gepetto-locales.
# Defaults to English.
language = "zh_CN"

# =============================================================================
# END
# =============================================================================

# Set up translations
translate = gettext.translation('gepetto',
                                os.path.join(os.path.abspath(
                                    os.path.dirname(__file__)), "gepetto-locales"),
                                fallback=True,
                                languages=[language])
_ = translate.gettext

# =============================================================================
# Setup the context menu and hotkey in IDA
# =============================================================================


class GepettoXunfeiPlugin(idaapi.plugin_t):
    flags = 0
    explain_action_name = "gepettoxunfei:explain_function"
    explain_menu_path = "Edit/Gepetto_xunfei/" + _("Explain function")
    rename_action_name = "gepettoxunfei:rename_function"
    rename_menu_path = "Edit/Gepetto_xunfei/" + _("Rename variables")
    wanted_name = 'Gepetto'
    wanted_hotkey = ''
    comment = _("Uses gpt-free to enrich the decompiler's output")
    help = _("See usage instructions on GitHub")
    menu = None

    def init(self):
        # Check whether the decompiler is available
        if not ida_hexrays.init_hexrays_plugin():
            return idaapi.PLUGIN_SKIP

        # Function explaining action
        explain_action = idaapi.action_desc_t(self.explain_action_name,
                                              _('Explain function'),
                                              ExplainHandler(),
                                              "Ctrl+Alt+1",
                                              _('Use gpt-free to explain the currently selected function'),
                                              199)
        idaapi.register_action(explain_action)
        idaapi.attach_action_to_menu(
            self.explain_menu_path, self.explain_action_name, idaapi.SETMENU_APP)

        # Variable renaming action
        rename_action = idaapi.action_desc_t(self.rename_action_name,
                                             _('Rename variables'),
                                             RenameHandler(),
                                             "Ctrl+Alt+2",
                                             _("Use gpt-free to rename this function's variables"),
                                             199)
        idaapi.register_action(rename_action)
        idaapi.attach_action_to_menu(
            self.rename_menu_path, self.rename_action_name, idaapi.SETMENU_APP)

        # Register context menu actions
        self.menu = ContextMenuHooks()
        self.menu.hook()

        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        pass

    def term(self):
        idaapi.detach_action_from_menu(
            self.explain_menu_path, self.explain_action_name)
        idaapi.detach_action_from_menu(
            self.rename_menu_path, self.rename_action_name)
        if self.menu:
            self.menu.unhook()
        return

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


class ContextMenuHooks(idaapi.UI_Hooks):
    def finish_populating_widget_popup(self, form, popup):
        # Add actions to the context menu of the Pseudocode view
        if idaapi.get_widget_type(form) == idaapi.BWN_PSEUDOCODE:
            idaapi.attach_action_to_popup(
                form, popup, GepettoXunfeiPlugin.explain_action_name, "Gepetto_xunfei/")
            idaapi.attach_action_to_popup(
                form, popup, GepettoXunfeiPlugin.rename_action_name, "Gepetto_xunfei/")

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


def comment_callback(address, view, response):
    """
    在给定地址设置评论的回调。
     :param address: 要注释的函数地址
     :param view: 反编译器窗口的句柄
     :param response: 要添加的评论
    """
    response = "\n".join(textwrap.wrap(response, 80, replace_whitespace=False))

    # Add the response as a comment in IDA, but preserve any existing non-Gepetto comment
    comment = idc.get_func_cmt(address, 0)
    comment = re.sub(r'----- ' + _("Comment generated by Gepetto") + ' -----.*?----------------------------------------',
                     r"",
                     comment,
                     flags=re.DOTALL)

    idc.set_func_cmt(address, '----- ' + _("Comment generated by Gepetto") +
                     f" -----\n\n"
                     f"{response.strip()}\n\n"
                     f"----------------------------------------\n\n"
                     f"{comment.strip()}", 0)
    # Refresh the window so the comment is displayed properly
    if view:
        view.refresh_view(False)
    print(_("gpt-free query finished!"))


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

class ExplainHandler(idaapi.action_handler_t):
    """
    This handler is tasked with querying gpt-3.5-turbo for an explanation of the
    given function. Once the reply is received, it is added as a function
    comment.
    """

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

    def activate(self, ctx):
        decompiler_output = ida_hexrays.decompile(idaapi.get_screen_ea())
        v = ida_hexrays.get_widget_vdui(ctx.widget)
        query_model_async(_("Can you explain what the following C function does and suggest a better name for it?\n"
                            "{decompiler_output}").format(decompiler_output=str(decompiler_output)),
                          functools.partial(comment_callback, address=idaapi.get_screen_ea(), view=v))
        return 1

    # This action is always available.
    def update(self, ctx):
        return idaapi.AST_ENABLE_ALWAYS

# -----------------------------------------------------------------------------
def rename_lvar(ea, src, dst):
    def make_unique_name(name, taken):
        if name not in taken:
            return name
        fmt = "%s_%%i" % name
        for i in range(3, 1024):
            tmpName = fmt % i
            if tmpName not in taken:
                return tmpName
        return "i_give_up"

    #  if you want to use an existing view:
    #      widget = ida_kernwin.find_widget('Pseudocode-Y')
    #      vu = ida_hexrays.get_widget_vdui(widget)
    func = idaapi.get_func(ea)
    if func:
        ea = func.start_ea
        vu = idaapi.open_pseudocode(ea, 0)
        names = [n.name for n in vu.cfunc.lvars]
        if dst in names:
            dst = make_unique_name(dst, names)
        lvars = [n for n in vu.cfunc.lvars if n.name == src]
        if len(lvars) == 1:
            print("renaming {} to {}".format(lvars[0].name, dst))
            vu.rename_lvar(lvars[0], dst, 1)
            # how to close the view without a widget object?
            #     idautils.close_pseudocode (nope)
            #     ida_kerwin.close_widget   (nope)
        else:
            print("couldn't find var {}".format(src))

def rename_callback(address, view, response, retries=0):
    """
    Callback that extracts a JSON array of old names and new names from the
    response and sets them in the pseudocode.
    :param address: The address of the function to work on
    :param view: A handle to the decompiler window
    :param response: The response from gpt-3.5-turbo
    :param retries: The number of times that we received invalid JSON
    """
    j = re.search(r"\{[^}]*?\}", response)
    if not j:
        if retries >= 3:  # Give up obtaining the JSON after 3 times.
            print(_("Could not obtain valid data from the model, giving up. Dumping the response for manual import:"))
            print(response)
            return
        print(
            _("Cannot extract valid JSON from the response. Asking the model to fix it..."))
        query_model_async(_("The JSON document provided in this response is invalid. Can you fix it?\n"
                            "{response}").format(response=response),
                          functools.partial(rename_callback,
                                            address=address,
                                            view=view,
                                            retries=retries + 1))
        return
    try:
        names = json.loads(j.group(0))
    except json.decoder.JSONDecodeError:
        if retries >= 3:  # Give up fixing the JSON after 3 times.
            print(_("Could not obtain valid data from the model, giving up. Dumping the response for manual import:"))
            print(response)
            return
        print(_("The JSON document returned is invalid. Asking the model to fix it..."))
        query_model_async(_("Please fix the following JSON document:\n{json}").format(json=j.group(0)),
                          functools.partial(rename_callback,
                                            address=address,
                                            view=view,
                                            retries=retries + 1))
        return

    # The rename function needs the start address of the function
    function_addr = idaapi.get_func(address).start_ea

    replaced = []
    for n in names:
        if rename_lvar(function_addr, n, names[n]):
            replaced.append(n)

    # Update possible names left in the function comment
    comment = idc.get_func_cmt(address, 0)
    if comment and len(replaced) > 0:
        for n in replaced:
            comment = re.sub(r'\b%s\b' % n, names[n], comment)
        idc.set_func_cmt(address, comment, 0)

    # Refresh the window to show the new names
    if view:
        view.refresh_view(True)
    print(
        _("gpt-3.5-turbo query finished! {replaced} variable(s) renamed.").format(replaced=len(replaced)))

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


class RenameHandler(idaapi.action_handler_t):
    """
    This handler requests new variable names from gpt-3.5-turbo and updates the
    decompiler's output.
    """

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

    def activate(self, ctx):
        decompiler_output = ida_hexrays.decompile(idaapi.get_screen_ea())
        v = ida_hexrays.get_widget_vdui(ctx.widget)
        query_model_async(_("Analyze the following C function:\n{decompiler_output}"
                            "\nSuggest better variable names, reply with a JSON array where keys are the original names "
                            "and values are the proposed names. Do not explain anything, only print the JSON "
                            "dictionary.").format(decompiler_output=str(decompiler_output)),
                          functools.partial(rename_callback, address=idaapi.get_screen_ea(), view=v))
        return 1

    # This action is always available.
    def update(self, ctx):
        return idaapi.AST_ENABLE_ALWAYS

# =============================================================================
# gpt-3.5-turbo interaction
# =============================================================================



 
text =[]

def getText(role,content):
    jsoncon = {}
    jsoncon["role"] = role
    jsoncon["content"] = content
    text.append(jsoncon)
    return text

def getlength(text):
    length = 0
    for content in text:
        temp = content["content"]
        leng = len(temp)
        length += leng
    return length

def checklen(text):
    while (getlength(text) > 8000):
        del text[0]
    return text
    


def updateImage(query):

    text.clear
    Input = query
    question = checklen(getText("user",Input))
    SparkApi.answer =""
    #print("星火:",end = "")
    return SparkApi.main(appid,api_key,api_secret,Spark_url,domain,question)



def query_model(query, cb, max_tokens=2500):
    """
    Function which sends a query to gpt-3.5-turbo and calls a callback when the response is available.
    Blocks until the response is received
    :param query: The request to send to gpt-3.5-turbo
    :param cb: Tu function to which the response will be passed to.
    """
    try:
        response = updateImage(query)

        ida_kernwin.execute_sync(functools.partial(cb, response=response),
                                 ida_kernwin.MFF_WRITE)
    except Exception as e:
        print(_("General exception encountered while running the query: {error}").format(
            error=str(e)))

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


def query_model_async(query, cb):
    """
    Function which sends a query to gpt-3.5-turbo and calls a callback when the response is available.
    :param query: The request to send to gpt-3.5-turbo
    :param cb: Tu function to which the response will be passed to.
    """
    print(_("Request to gpt-free sent..."))
    t = threading.Thread(target=query_model, args=[query, cb])
    t.start()

# =============================================================================
# Main
# =============================================================================


def PLUGIN_ENTRY():
    # if not openai.api_key:
    #     openai.api_key = os.getenv("OPENAI_API_KEY")
    #     if not openai.api_key:
    #         print(_("Please edit this script to insert your OpenAI API key!"))
    #         raise ValueError("No valid OpenAI API key found")

    return GepettoXunfeiPlugin()


repository: cellebrite-labs/FunctionInliner
file: functioninliner.py
wanted_name: FunctionInliner
comment: inlines functions that were outlined
created_at: 2021-09-13 12:04:02+00:00
pushed_at: 2024-12-31 13:21:07+00:00
description:
An IDA plugin that eases reversing of binaries that have been code-size-optimized with function outlining

url:
https://github.com/cellebrite-labs/FunctionInliner/blob/cd0ca68f9b1a5675032c24892d6dcd1e22297a08/functioninliner.py

content:
import collections
import contextlib
import functools
import itertools
import logging
import pickle
import re
import struct
import time
import types

import ida_auto
import ida_bytes
import ida_funcs
import ida_frame
import ida_idaapi
import ida_idp
import ida_kernwin
import ida_segment
import ida_ua
import ida_xref
import idc

import keypatch
import netnode
import parse
import sark
import tqdm
import wrapt


# DEFINITIONS


INLINED_FUNCTION_PREFIX = "inlined_"
CLONE_NAME_FMT = "inlined_0x{func_ea:x}_for_0x{src_ea:x}"
TRACE = False


# LOGGING


class LoggerWithTrace(logging.getLoggerClass()):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        if TRACE:
            # add TRACE level
            logging.TRACE = 5
            logging.addLevelName(logging.TRACE, "TRACE")

    def trace(self, msg, *args, **kwargs):
        if TRACE:
            self.log(logging.TRACE, msg, *args, **kwargs)


logger = LoggerWithTrace("FunctionInliner")


# EXCEPTIONS


class FunctionInlinerException(Exception):
    pass


class FunctionInlinerUnsupportedException(FunctionInlinerException):
    pass


class FunctionInlinerUnknownFlowException(FunctionInlinerException):
    pass


# HELPERS


@contextlib.contextmanager
def autoanalysis(enabled):
    ida_auto.enable_auto(enabled)
    try:
        yield None
    finally:
        ida_auto.enable_auto(not enabled)


def with_autoanalysis(enabled):
    @wrapt.decorator
    def decorator(wrapped, instance, args, kwargs):
        with autoanalysis(enabled):
            return wrapped(*args, **kwargs)
    return decorator


@contextlib.contextmanager
def wait_box(msg, hide_cancel=False):
    prefix = "HIDECANCEL\n" if hide_cancel else ""
    ida_kernwin.show_wait_box(prefix + msg)
    try:
        yield None
    finally:
        ida_kernwin.hide_wait_box()


def get_function_under_cursor():
    line = sark.Line()

    # abort on unmapped addresses
    if not ida_bytes.is_mapped(line.ea):
        return None

    # if we're on a call -> return its target
    for xref in line.xrefs_from:
        if xref.type.is_jump or xref.type.is_call:
            try:
                f = sark.Function(xref.to)
                if xref.to == f.ea:
                    return f
            except sark.exceptions.SarkNoFunction:
                return None

    # if we're on the start of a function -> return it
    try:
        func = sark.Function()
        if func.start_ea == line.ea:
            return func
    except sark.exceptions.SarkNoFunction:
        return None

    return None


def align_downwards(ea, alignment):
    return ea & ~(alignment - 1)


def align_upwards(ea, alignment):
    return align_downwards(ea + (alignment - 1), alignment)


def reanalyze_line(line):
    ida_auto.plan_range(line.ea, line.end_ea)


def reanalyze_program():
    """ we used to not reanalyze the entire program, but for some reason when we surgically marked
    for reanalysis only the stuff that we've changed, sometimes the auto analysis didn't recursively
    go through to everything """
    ida_auto.plan_range(0, ida_idaapi.BADADDR)


def is_conditional_insn(insn):
    # is having a condition suffix
    return insn._insn.segpref != 0xe  # see module/arm/arm.hpp in the IDA SDK


def is_chunked_function(func):
    return len(list(function_chunk_eas(func))) > 1


def function_chunk_eas(func):
    ea = idc.first_func_chunk(func.ea)
    while ea != ida_idaapi.BADADDR:
        yield ea
        ea = idc.next_func_chunk(func.ea, ea)


def function_chunk_lines(ea):
    start_ea = idc.get_fchunk_attr(ea, idc.FUNCATTR_START)
    end_ea = idc.get_fchunk_attr(ea, idc.FUNCATTR_END)

    ea = start_ea
    l = sark.Line(start_ea)
    while l.ea < end_ea:
        yield l
        l = l.next


def function_chunk_crefs(ea, ret_ea=None):
    start_ea = idc.get_fchunk_attr(ea, idc.FUNCATTR_START)
    end_ea = idc.get_fchunk_attr(ea, idc.FUNCATTR_END)

    for l in function_chunk_lines(ea):
        try:
            if l.insn.mnem == "RET" and ret_ea is not None:
                external_cref_eas = (ret_ea,)
            else:
                external_cref_eas = (c for c in l.crefs_from if c < start_ea or c >= end_ea)

            for target_ea in external_cref_eas:
                yield (l.ea - start_ea, target_ea)
        except sark.exceptions.SarkNoInstruction:
            if not list(l.crefs_to):
                pass  # some times there's non-code inside function chunks, e.g. jumptables
            else:
                raise  # but if there's a flow cref into it, this is bad


def function_chunk_parent_eas(ea):
    fchunk = ida_funcs.get_fchunk(ea)
    fpi = ida_funcs.func_parent_iterator_t(fchunk)
    if not fpi.first():
        return

    while True:
        yield fpi.parent()
        if not fpi.next():
            break


def containing_funcs(line):
    funcs = set()
    for parent_ea in function_chunk_parent_eas(line.ea):
        try:
            funcs.add(sark.Function(parent_ea))
        except sark.exceptions.SarkNoFunction:
            pass
    try:
        funcs.add(sark.Function(line.ea))
    except sark.exceptions.SarkNoFunction:
        pass
    return funcs


def unreachable_function_chunks_eas(func):
    # map all chunks in our function
    remaining_chunks = set()
    for start_ea in function_chunk_eas(func):
        end_ea = idc.get_fchunk_attr(start_ea, idc.FUNCATTR_END)
        remaining_chunks.add((start_ea, end_ea))

    if len(remaining_chunks) == 1:
        return

    # discard reachable chunks
    def discard_reachable_chunks(chunk):
        remaining_chunks.discard(chunk)

        for src_off, target in function_chunk_crefs(chunk[0]):
            for other_chunk in remaining_chunks:
                start_ea, end_ea = other_chunk
                if start_ea <= target < end_ea:
                    discard_reachable_chunks(other_chunk)
                    break
            if not remaining_chunks:
                break

    main_chunk = (func.start_ea, func.end_ea)
    assert main_chunk in remaining_chunks

    discard_reachable_chunks(main_chunk)

    yield from (c[0] for c in remaining_chunks)


def function_crefs(func, ret_ea=None):
    chunk_eas = list(function_chunk_eas(func))
    for chunk_ea in chunk_eas:
        for off, target_ea in function_chunk_crefs(chunk_ea, ret_ea):
            if ida_funcs.func_contains(func._func, target_ea):
                continue

            src_ea = chunk_ea + off
            src = sark.Line(src_ea)
            if src.insn.mnem == "BL" and target_ea == src.end_ea:
                # we have a flow xref going from a BL which is the last instruction of a function
                # chunk. this can be one of two cases:
                # 1. IDA didn't recognize that the target function is a NORET function, hence there
                #    shouldn't be a flow-xref from the BL
                # 2. after BL-ing, the source function implicitly tail-calls to the next instruction
                #    which should be another function
                # because we cannot know which of the cases is the right one, and in case it's (1),
                # the next line might be data, we'll let a heuristic decide
                if is_data_heuristic(sark.Line(target_ea)):
                    continue

            yield src_ea, target_ea


def has_function_flow_xref(line):
    for xref in line.xrefs_to:
        # we're only interested in flow xrefs
        if not xref.type.is_flow:
            continue

        # sometimes there are useless NOPs before real code, so we discard xrefs from code
        # which isn't in a function
        try:
            sark.Function(xref.frm)
        except sark.exceptions.SarkNoFunction:
            continue

        return True

    return False


def external_callers(line, functions_only=False, include_flow=False):
    funcs = containing_funcs(line)

    for xref in line.xrefs_to:
        caller = sark.Line(xref.frm)

        # skip non-code xrefs
        if not xref.type.is_code:
            continue

        # skip flow xrefs
        if not include_flow and xref.type.is_flow:
            continue

        # skip recursive calls
        caller_funcs = containing_funcs(caller)

        if funcs and caller_funcs == funcs:
            continue
        elif not caller_funcs and functions_only:
            continue

        yield caller


def linegroups(n):
    iters = [sark.lines() for _ in range(n)]
    for i, it in enumerate(iters):
        for _ in range(i):
            next(it)
    return zip(*iters)


def register_parts(r):
    w = r[0]
    n = r[1:]

    families = (
        ("W", "X"),
        ("B", "H", "S", "D", "Q")
    )

    for f in families:
        if w in f:
            return (ww + n for ww in f)
    else:
        raise FunctionInlinerException(f"encountered unknown register: {r}")


def add_comment(line, cmt, prepend=True):
    if line.comments.regular is None:
        line.comments.regular = cmt
    elif prepend:
        line.comments.regular = cmt + "\n" + line.comments.regular
    else:
        line.comments.regular = line.comments.regular + "\n" + cmt


def get_branch_target_ea(line):
    crefs_from = set()
    for xref in line.xrefs_from:
        if not xref.iscode:
            continue
        if xref.type.is_flow:
            continue
        crefs_from.add(xref.to)
    assert len(crefs_from) == 1
    return crefs_from.pop()


def drefs_from_eas(line):
    # IDA marks enum refs as drefs with top address byte set to 0xff, so we filter out
    # drefs that are not actually mapped
    return (ea for ea in line.drefs_from if ida_bytes.is_mapped(ea))


# NETNODE


class PickleNetnode(netnode.Netnode):
    @staticmethod
    def _encode(data):
        return pickle.dumps(data)

    @staticmethod
    def _decode(data):
        return pickle.loads(data)


class RenamesStorage(PickleNetnode):
    NETNODE = "$ FunctionInliner.renames"

    RenameInfo = collections.namedtuple("RenameInfo", ("orig_name", "new_name"))

    def __init__(self):
        super().__init__(self.NETNODE)

    def __setitem__(self, func_ea, rename_info):
        if not isinstance(rename_info, RenamesStorage.RenameInfo):
            raise ValueError("value must be of type RenamesStorage.RenameInfo")
        super().__setitem__(func_ea, tuple(rename_info))

    def __getitem__(self, func_ea):
        v = super().__getitem__(func_ea)
        return RenamesStorage.RenameInfo(*v)


class SingletonUserDict(type(collections.UserDict)):
    _instance = None

    def __call__(cls, *args, **kwargs):
        if not cls._instance:
            cls._instance = super(SingletonUserDict, cls).__call__(*args, **kwargs)
        return cls._instance


class ClonesStorage(collections.UserDict, metaclass=SingletonUserDict):
    NETNODE = "$ functioninliner.clones"

    CloneInfo = collections.namedtuple("CloneInfo", ("clone_ea", "orig_bytes"))

    class InlinedFunctionInfo(collections.UserDict):
        def __init__(self, update_callback):
            super().__init__()
            self._update_callback = update_callback

        def __setitem__(self, src_ea, clone):
            if not isinstance(clone, ClonesStorage.CloneInfo):
                raise ValueError("value must be of type ClonesStorage.CloneInfo")
            self.data[src_ea] = clone
            self._update_callback(self, src_ea)

        def __delitem__(self, src_ea):
            del self.data[src_ea]
            self._update_callback(self, src_ea)

    def __init__(self):
        super().__init__()
        self.netnode = PickleNetnode(self.NETNODE)
        self.update_from_storage()

    def update_from_storage(self):
        self.data.clear()

        for k, v in self.netnode.items():
            parts = ClonesStorage.parse_storage_key(k)
            if not parts:
                continue

            func_ea = parts["func_ea"]
            src_ea = parts["src_ea"]

            func_storage = self[func_ea]
            clone_info = ClonesStorage.CloneInfo(*v)
            func_storage[src_ea] = clone_info

    @staticmethod
    def storage_key(func_ea, src_ea):
        return CLONE_NAME_FMT.format(func_ea=func_ea, src_ea=src_ea)

    @staticmethod
    def parse_storage_key(k):
        return parse.parse(CLONE_NAME_FMT, k)

    def write_to_storage(self, func_ea, func_storage, src_ea):
        key = ClonesStorage.storage_key(func_ea, src_ea)
        clone_info = func_storage.get(src_ea, None)
        if clone_info is None:  # delete
            del self.netnode[key]

            # if there are no more outlined
            if not func_storage and func_ea in self.data:
                del self.data[func_ea]
        else:  # set
            self.netnode[key] = tuple(clone_info)
            self.data[func_ea] = func_storage

    def __getitem__(self, func_ea):
        if func_ea in self.data:
            return self.data[func_ea]
        else:
            update_callback = functools.partial(self.write_to_storage, func_ea)
            func_storage = ClonesStorage.InlinedFunctionInfo(update_callback)
            return func_storage

    def __setitem__(self, k, v):
        raise RuntimeError("Do not try setting a value directly, but rather get its value as with "
                           "a defaultdict")

    def __delitem__(self, func_ea):
        for src_ea, _ in self[func_ea].items():
            key = ClonesStorage.storage_key(func_ea, src_ea)
            del self.netnode[key]

        if func_ea in self.data:
            del self.data[func_ea]


# FUNCTION INLINING


def get_cloned_function(ea):
    seg = sark.Segment(ea)
    if not seg.name:
        return None

    parts = parse.parse(CLONE_NAME_FMT, seg.name)
    if not parts:
        return None

    return sark.Function(parts["func_ea"])


def is_originally_chunked_function(func):
    for chunk_ea in function_chunk_eas(func):
        # dismiss the "main" chunk
        if chunk_ea == func.ea:
            continue

        # dismiss chunks which are inlined clones of outlined functions
        if get_cloned_function(chunk_ea):
            continue

        # we found a "real" chunk
        return True

    return False


def create_code_segment(name, size, close_to=None, page_align=False):
    if page_align:
        alignment = 0x1000
    else:
        alignment = 0x4

    size = align_upwards(size, alignment)

    segs = list(sorted(sark.segments(), key=lambda s: s.ea))

    # delete a previously cloned segment if such exists
    for s in segs:
        if s.name == name:
            ida_segment.del_segm(s.start_ea, ida_segment.SEGMOD_KILL)

    # map the holes between existing segments
    holes = []
    holes.append((0, segs[0].start_ea))
    for s, next_s in zip(segs, segs[1:]):
        holes.append((s.end_ea, next_s.start_ea))
    holes.append((segs[-1].end_ea, ida_idaapi.BADADDR))

    # align the start and end of each hole
    holes = [(align_upwards(h[0], alignment), align_downwards(h[1], alignment)) for h in holes]

    # filter-out holes which are too small
    holes = [h for h in holes if h[1] - h[0] >= size]

    # find the hole nearest to our caller
    if close_to is None:
        hole = holes[-1]
    else:
        def hole_dist(h):
            start, end = h
            if start > close_to:
                return start - close_to
            else:
                return close_to - end - size
        hole = min(holes, key=hole_dist)

    # create the segment
    if hole[0] > close_to:
        start_ea = hole[0]
        end_ea = hole[0] + size
    else:
        start_ea = hole[1] - size
        end_ea = hole[1]

    seg_t = ida_segment.segment_t()
    seg_t.start_ea = start_ea
    seg_t.end_ea = end_ea
    seg_t.align = ida_segment.saRelDble
    seg_t.comb = ida_segment.scPub
    seg_t.perm = ida_segment.SEGPERM_EXEC | ida_segment.SEGPERM_READ
    seg_t.bitness = 2  # 64 bits
    seg_t.sel = ida_segment.setup_selector(0)
    seg_t.type = ida_segment.SEG_CODE
    seg_t.color = idc.DEFCOLOR

    flags = ida_segment.ADDSEG_NOSREG | ida_segment.ADDSEG_QUIET | ida_segment.ADDSEG_NOAA

    ida_segment.add_segm_ex(seg_t, name, "CODE", flags)

    return sark.Segment(start_ea)


def validate_branch_displacements(func, src_ea, clone_ea, ret_ea):
    # we only go over the first function chunk since that's the one we're cloning

    max_b_displ = 0x8000000

    def b_displ(src, target):
        return abs(target - (src + 4))

    # this isn't the most accurate condition since the clone offsets may move because of
    # our translation, but we discard that. if anything, the assembly later will fail
    clone_cref_displs = (b_displ(clone_ea + src_off, target_ea) for src_off, target_ea in
                         function_chunk_crefs(func.ea, ret_ea))
    if b_displ(src_ea, clone_ea) >= max_b_displ or \
            any(displ >= max_b_displ for displ in clone_cref_displs):

        # TBH we were greedy when choosing when to create our segment, but we don't expect
        # this to fail with +-128MB of max displacement, so we didn't bother implementing
        # a better algorithm
        raise FunctionInlinerException("created clone segment is not close enough to its "
                                       "caller or one of its call targets")


def fix_outlined_function_call(src, clone_ea, clone_end_ea, func_ea, kp_asm=None):
    if kp_asm is None:
        kp_asm = keypatch.Keypatch_Asm()

    # unfortunately, we've seen cases where IDA creates a function out of our clone instead of a
    # function chunk, and this also happens sometimes when calling auto_apply_tail.
    # therefore, we first ida_funcs.append_func_tail and only then patch and plan to reanalyze the
    # caller
    ida_funcs.append_func_tail(ida_funcs.get_func(src.ea), clone_ea, clone_end_ea)

    # replace the source instruction with a B to our clone
    mnem = src.insn.mnem
    if mnem == "BL":
        asm = f"B #{clone_ea:#x}"  # we drop PAC flags
        code = bytes(kp_asm.assemble(asm, src.ea)[0])
        assert len(code) == src.size
        ida_bytes.patch_bytes(src.ea, code)
    else:  # is_jump
        fix_cloned_branch(kp_asm, src.ea, func_ea, clone_ea)

    reanalyze_line(src)

    # delete the original xref
    ida_xref.del_cref(src.ea, func_ea, 0)

    # wait for analysis of the source instruction and the clone
    assert ida_auto.auto_wait_range(src.ea, src.end_ea) >= 0
    assert ida_auto.auto_wait_range(clone_ea, clone_end_ea) >= 0

    # recalculate SP delta for the instructions in the clone. we do it for every instruction in
    # case the clone contains more than one basic block
    pfn = ida_funcs.get_func(src.ea)
    ea = clone_ea
    while ea < clone_end_ea:
        ida_frame.recalc_spd_for_basic_block(pfn, ea)
        ea += ida_bytes.get_item_size(ea)

def inline_function_call(src, func, kp_asm=None):
    storage = ClonesStorage()
    func_storage = storage[func.ea]

    if kp_asm is None:
        kp_asm = keypatch.Keypatch_Asm()

    # verify that the function isn't chunked
    if is_chunked_function(func):
        raise FunctionInlinerUnsupportedException("chunked functions are currently unsupported")

    # clone and inline the function for each of its callers
    size = func.end_ea - func.start_ea

    # create a segment for the cloned function
    seg_size = size * 2  # we put a factor of 2 here for our ADR->ADRP+ADD fixups
    seg_name = CLONE_NAME_FMT.format(func_ea=func.ea, src_ea=src.ea)
    seg = create_code_segment(seg_name, seg_size, src.ea)
    clone_ea = seg.ea

    try:
        # analyze the caller
        if src.insn.mnem == "B" and not is_conditional_insn(src.insn):  # tail-call
            ret_ea = None
        elif src.insn.mnem in ("B", "BL", "CBNZ", "CBZ", "TBNZ", "TBZ"):
            ret_ea = src.end_ea
        else:
            raise FunctionInlinerException(f"unexpected call opcode: {src.insn.mnem}")

        # validate that the created segment is close enough for the required branches
        validate_branch_displacements(func, src.ea, clone_ea, ret_ea)

        # clone the function
        logger.debug(f"cloning to {clone_ea:#x}")
        clone_end_ea = clone_function(func, clone_ea, ret_ea, kp_asm)

        # replace the source opcode with a branch to our clone
        orig_bytes = src.bytes
        fix_outlined_function_call(src, clone_ea, clone_end_ea, func.ea, kp_asm)

        # add clone info to storage
        clone_info = ClonesStorage.CloneInfo(clone_ea, orig_bytes)
        func_storage[src.ea] = clone_info

    except:  # noqa: E722
        # remove from storage if it's already been added
        func_storage.pop(clone_ea, None)

        # undo the source patch if it's already been done
        ida_bytes.patch_bytes(src.ea, src.bytes)
        reanalyze_line(src)

        # remove the created segment
        ida_segment.del_segm(clone_ea, ida_segment.SEGMOD_KILL)

        logger.error(f"unhandled exception was raised while inlining call to {func.name} from {src.ea:#x}")
        raise

    return clone_ea


def function_chunk_inlined_functions(ea):
    start_ea = idc.get_fchunk_attr(ea, idc.FUNCATTR_START)

    for off, target_ea in function_chunk_crefs(ea):
        func = get_cloned_function(target_ea)

        if func:
            src = sark.Line(start_ea + off)
            yield (src, func)


def rename_outlined_function(func):
    storage = RenamesStorage()

    rename_info = storage.get(func.ea)
    if rename_info and func.name == rename_info.new_name:
        return

    new_name = f"outlined_{func.name}"
    rename_info = RenamesStorage.RenameInfo(func.name, new_name)

    func.name = new_name
    storage[func.ea] = rename_info


def undo_rename_outlined_function(func):
    storage = RenamesStorage()

    rename_info = storage.get(func.ea)
    if not rename_info:
        return

    if func.name == rename_info.new_name:
        func.name = rename_info.orig_name

    del storage[func.ea]


def inline_function(func, kp_asm=None):
    # verify that the function doesn't have any chunks which aren't inlined clones of outlined
    # functions
    if is_originally_chunked_function(func):
        raise FunctionInlinerUnsupportedException("chunked functions are currently unsupported")

    # find functions that we've inlined into this function
    inlined_function_calls = list(function_chunk_inlined_functions(func.ea))

    # temporarily undo inlining into our function
    for src, inlined_func in inlined_function_calls:
        logger.debug(f"temporarily undoing inlining of {inlined_func.name} into {src.ea:#x}")
        undo_inline_function_call(src, inlined_func)

    # wait for analysis of our function
    assert ida_auto.auto_wait_range(func.ea, func.end_ea) >= 0

    # assert that the function is now unchunked
    assert not is_chunked_function(func)

    func_has_outgoing_crefs = bool(list(function_chunk_crefs(func.ea)))

    # inline our function into its callers
    for src in external_callers(func):
        logger.debug(f"inlining call to {func.name} from {src.ea:#x}")
        clone_ea = inline_function_call(src, func, kp_asm)

        # wait for analysis of our clone (optimization: only if needed)
        # this is both for the case we need to redo inlining into it in a bit, and also for the case
        # we're inlining all outlined functions in the IDB and we want the clones' outgoing crefs to
        # be indexed as well
        if func_has_outgoing_crefs:
            seg = sark.Segment(clone_ea)
            assert ida_auto.auto_wait_range(seg.start_ea, seg.end_ea) >= 0

    # redo inlining into our function clones
    for _, inlined_func in inlined_function_calls:
        logger.debug(f"redoing inlining of {inlined_func.name} into the cloned functions")
        inline_function(inlined_func)

    # if there are no more xrefs to this function, rename it
    if not list(external_callers(func)):
        rename_outlined_function(func)


@with_autoanalysis(False)
def inline_all_functions(extra=False):
    logger.info("inlining all outlined functions...")

    failed_analysis = 0
    inlined = 0
    skipped = 0
    erroronous = 0
    kp_asm = keypatch.Keypatch_Asm()

    all_funcs = list(sark.functions())

    with wait_box("finding outlined functions..."):
        outlined_funcs = []
        for func in tqdm.tqdm(all_funcs, desc="analyzing", ncols=80, unit="func"):
            ida_auto.show_auto(func.ea)

            if ida_kernwin.user_cancelled():
                return False

            logger.debug(f"analyzing {func.name}")
            try:
                if is_function_outlined(func, extra=extra):
                    outlined_funcs.append(func)
            except Exception:
                logger.exception(f"unhandled exception raised when trying to analyze {func.name}:")
                failed_analysis += 1

    logger.info(f"found {len(outlined_funcs)} outlined functions")
    if failed_analysis:
        logger.error(f"failed analysing {failed_analysis} functions")

    retval = True
    with wait_box("inlining all outlined functions... (this may take a few minutes)"):
        start_time = time.time()

        for func in tqdm.tqdm(outlined_funcs, desc="inlining", ncols=80, unit="func"):
            ida_auto.show_auto(func.ea)

            if ida_kernwin.user_cancelled():
                retval = False
                break

            logger.debug(f"inlining {func.name}")
            try:
                inline_function(func, kp_asm)
                inlined += 1
            except FunctionInlinerUnsupportedException:
                skipped += 1  # skip functions that we can't inline
            except Exception:
                logger.exception(f"unhandled exception raised when trying to inline {func.name}:")
                erroronous += 1

        elapsed_time = int(time.time() - start_time)

    logger.info(f"inlined a total of {inlined} functions in {elapsed_time} seconds")

    if skipped:
        logger.warning(f"skipped {skipped} unsupported functions")

    if erroronous:
        logger.error(f"failed inlining {erroronous} functions")

    return retval


def clone_insn_ret(kp_asm, line, dst_ea, ret_ea):
    assert line.insn.mnem == "RET" and not is_conditional_insn(line.insn)

    asm = f"B #{ret_ea:#x}"  # we drop PAC flags
    code = bytes(kp_asm.assemble(asm, dst_ea)[0])

    logger.trace(f"   translated to: {asm}")
    return code


def clone_insn_branch(kp_asm, line, dst_ea, func, ret_ea):
    mnem = line.insn.mnem

    # resolve the the branch target
    if mnem == "BR":
        target_ea = None
    else:
        target_ea = get_branch_target_ea(line)

    if target_ea and func.start_ea <= target_ea < func.end_ea:  # local target -> copy as-is
        logger.trace("   local target -> copied as-is")
        return line.bytes, target_ea - func.ea

    else:  # external target -> fix it
        if is_conditional_insn(line.insn):
            raise FunctionInlinerUnsupportedException("translating conditioned tail-calls is "
                                                      "currently unsupported")

        if mnem == "BR":
            # target is the first arg even if it's an authenticated BR
            target_reg = line.insn.operands[0].text
            asm = f"BLR {target_reg}"  # we drop PAC flags
        else:
            assert mnem in ("BL", "B")
            asm = f"BL #{target_ea:#x}"  # we drop PAC flags

        if mnem in ("B", "BR"):  # tail-call -> also add a following B back or RET
            if ret_ea:
                asm += f"\nB #{ret_ea:#x}"
            else:
                asm += "\nRET"

        code = bytes(kp_asm.assemble(asm, dst_ea)[0])

        logger.trace(f"   translated to: {asm}")
        return code, None


def clone_insn_mem(kp_asm, line, dst_ea):
    drefs_from = set(drefs_from_eas(line))
    insn = line.insn

    if len(drefs_from) == 1:
        target_ea = drefs_from.pop()
    else:  # this may happen with LDR when the target contains another address
        assert len(drefs_from) == 2 and insn.mnem == "LDR"
        xref = [x for x in line.xrefs_from if x.type.is_read][0]
        target_ea = xref.to

    target_page = target_ea & ~0xfff
    target_offset = target_ea & 0xfff

    if is_conditional_insn(line.insn):
        raise FunctionInlinerUnsupportedException("translating conditional mem instructions is "
                                                  "currently unsupported")

    # full_mnem should be the same as insn.mnem, but compare the full one just to be on the safe side
    full_mnem = ida_ua.print_insn_mnem(line.ea)

    if full_mnem in ("ADR", "ADRL"):
        # usually ADR is followed by NOP, because the compiler doesn't know if it'll be ADR or
        # ADRP + ADD (ADRL), but we don't want to rely on it so we'll translate both alternatives
        # with two instructions

        reg = line.insn.regs.pop()
        asm = f"""
            ADRP {reg}, #{target_page:#x}
            ADD {reg}, {reg}, #{target_offset:#x}
        """
        code = bytes(kp_asm.assemble(asm, dst_ea)[0])

        logger.trace(f"   {full_mnem} -> translated to: {asm}")
        return code

    elif full_mnem == "ADRP":
        reg = line.insn.regs.pop()
        asm = f"ADRP {reg}, #{target_page:#x}"
        code = bytes(kp_asm.assemble(asm, dst_ea)[0])

        logger.trace(f"   ADRP -> translated to: {asm}")
        return code

    elif full_mnem == "MOV":
        logger.trace("   MOV -> copied as-is")
        return line.bytes

    else:
        # we expect the rest to be e.g. LDR/STR/ADD using an address or a PAGEOFF

        # IDA won't always show "@PAGEOFF" in the disassembly, so we have to check for this case
        # manually
        found_displ = 0
        for op in line.insn.operands:
            if op.type.is_mem:
                found_displ = -1
                break
            if op.type.is_displ:
                found_displ += 1
        pageoff_flow = found_displ == 1

        if pageoff_flow:
            # the PAGEOFF shouldn't change, so we can copy as-is
            logger.trace("   PAGEOFF -> copied as-is")
            return line.bytes

        else:  # direct memory access flow
            # try finding and replacing the target operand with the fixed one
            new_ops = []
            ops_fixed = 0
            for op in insn.operands:
                if op.type.is_mem and op.addr == target_ea:
                    new_ops.append(f"#{target_ea:#x}")
                    ops_fixed += 1
                else:
                    new_ops.append(op.text)
            assert ops_fixed == 1

            # recreate the instruction
            asm = full_mnem + " " + ", ".join(new_ops)
            code = bytes(kp_asm.assemble(asm, dst_ea)[0])

            logger.trace(f"   direct memory reference -> translated to: {asm}")
            return code


def fix_cloned_branch(kp_asm, src_ea, current_target_ea, fixed_target_ea):
    # analyze this single instruction
    idc.set_flag(idc.INF_AF, idc.AF_CODE, 0)
    ida_ua.create_insn(src_ea)
    idc.set_flag(idc.INF_AF, idc.AF_CODE, 1)

    line = sark.Line(src_ea)
    insn = line.insn

    # we believe that our handling here won't match all cases, but we've never even encountered
    # a case where this fixup has been required so we don't want to spend too much effort around it

    # try finding and replacing the target operand with the fixed one
    new_ops = []
    ops_fixed = 0
    for op in insn.operands:
        if op.type.is_near and op.addr == current_target_ea:
            new_ops.append(f"#{fixed_target_ea:#x}")
            ops_fixed += 1
        else:
            new_ops.append(op.text)
    assert ops_fixed == 1

    # recreate the instruction
    full_mnem = ida_ua.print_insn_mnem(line.ea)
    asm = full_mnem + " " + ", ".join(new_ops)
    code = bytes(kp_asm.assemble(asm, line.ea)[0])

    logger.trace(f"  -> {asm}")

    assert line.size == len(code)

    # undo the analysis of this instruction and patch it
    ida_bytes.del_items(src_ea, ida_bytes.DELIT_SIMPLE, line.size)
    ida_bytes.patch_bytes(line.ea, code)


def clone_function(func, dst_ea, ret_ea=None, kp_asm=None):
    if kp_asm is None:
        kp_asm = keypatch.Keypatch_Asm()

    clone_ea = dst_ea

    # maps func_offset to clone_offset, for each cloned instruction
    clone_offsets = {}
    # maps source_func_offset to target_func_offset, for cloned branches which point internally
    # to the cloned function
    potential_target_fixups = {}

    # go over all instructions in the first function chunk
    for line in function_chunk_lines(func.start_ea):
        logger.trace(f"-> {line.disasm}")

        # remember where we moved it to in the clone
        func_offset = line.ea - func.ea
        clone_offset = dst_ea - clone_ea
        clone_offsets[func_offset] = clone_offset

        # clone this instruction
        try:
            assert line.is_code

            # generate some metadata about it

            is_brk = line.insn.mnem == "BRK"
            is_ret = line.insn.mnem == "RET"

            is_normal_flow = False
            for xref in line.xrefs_from:
                if not xref.iscode:
                    continue

                if not xref.type.is_flow:
                    is_normal_flow = False  # at least one non-flow code xrefs -> not normal flow
                    break
                else:
                    is_normal_flow = True  # all (at least one) flow code xrefs -> normal flow

            if any(drefs_from_eas(line)):
                has_drefs = True
            else:
                has_drefs = False

            # identify which kind of translation we should do

            if (is_normal_flow and not has_drefs) or (is_ret and not ret_ea) or is_brk:  # "simple" instruction
                logger.trace("   'simple' instruction flow -> copied as-is")
                code = line.bytes

            elif is_ret and ret_ea:  # ret (and we should translate it)
                assert not has_drefs

                logger.trace("   ret flow")
                code = clone_insn_ret(kp_asm, line, dst_ea, ret_ea)

            elif not is_normal_flow:  # conditional branch, BL, or tail-call
                assert not has_drefs

                # note: we never actually encountered any outlined code with conditional branches or BLs,
                # but tail-call handling is similiar so we're keeping the logic here anyway

                logger.trace("   conditional branch, bl or tail-call flow")
                code, target_offset = clone_insn_branch(kp_asm, line, dst_ea, func, ret_ea)

                # if the new code jumps into an address internal to our cloned function, we may need
                # to fix it up afterwards (depending on how much we'll move stuff around)
                if target_offset is not None:
                    potential_target_fixups[func_offset] = target_offset

            elif has_drefs:  # memory access
                assert is_normal_flow

                # note: we never actually encountered outlined code which did direct memory access
                # but we're keeping the logic here anyway

                logger.trace("   memory access flow")
                code = clone_insn_mem(kp_asm, line, dst_ea)

            else:
                raise FunctionInlinerException("unexpected instruction")

            # write the new code to the clone

            ida_bytes.patch_bytes(dst_ea, code)
            dst_ea += len(code)

        except Exception:
            logger.error(f"failed to clone instruction @ {line.ea:#x}: {line.disasm}")
            raise

    # fix local branch targets if required
    for src_offset, target_offset in potential_target_fixups.items():
        src_ea = clone_ea + clone_offsets[src_offset]
        logger.trace(f"analyzing local branch at {src_ea:#x}")

        current_target_ea = src_ea + (target_offset - src_offset)
        fixed_target_ea = clone_ea + clone_offsets[target_offset]

        if current_target_ea != fixed_target_ea:
            # note: we never actually gotten to this flow in any outlined code we've inlined
            # but we're keeping the logic here anyway

            logger.trace(f"  fixing target from {current_target_ea:#x} to {fixed_target_ea:#x}")
            fix_cloned_branch(kp_asm, src_ea, current_target_ea, fixed_target_ea)
        else:
            logger.trace("  no fixing is required")

    # analyze the clone as code
    ida_auto.auto_make_code(clone_ea)

    return dst_ea


# UNDO FUNCTION INLINING


def get_inlined_function_under_cursor():
    line = sark.Line()

    # abort on unmapped addresses
    if not ida_bytes.is_mapped(line.ea):
        return None

    # if we're on a branch/call -> analyze its target instead
    for xref in line.xrefs_from:
        if xref.type.is_jump or xref.type.is_call:
            line = sark.Line(xref.to)
            break

    # if we're in a cloned segment -> return the function it was cloned from
    seg = sark.Segment(line.ea)
    if seg is None:  # we're pointing at an old (now non-existant) line
        return None
    func = get_cloned_function(line.ea)
    if line.ea == seg.ea and func:
        return func

    # if we're on the beginning of an inlined function -> return it
    try:
        func = sark.Function()
    except sark.exceptions.SarkNoFunction:
        return None

    storage = ClonesStorage()
    if line.ea == func.ea and func.ea in storage:
        return func

    return None


def undo_inline_function_call(src, func):
    storage = ClonesStorage()
    func_storage = storage[func.ea]
    clone_info = func_storage[src.ea]

    logger.debug(f"undoing clone of {func.name} at {clone_info.clone_ea:#x} for caller at {src.ea:#x}")

    # delete the cloned function
    ida_segment.del_segm(clone_info.clone_ea, ida_segment.SEGMOD_KILL)

    if ida_bytes.is_mapped(src.ea):  # maybe this was into another clone that has been undone as well
        # revert the BL patch
        # we don't want to do revert_byte() here, since the patched opcode may have been
        # originally patched (e.g. inlining one outlined function into another outlined function
        # that has been inlined)
        ida_bytes.patch_bytes(src.ea, clone_info.orig_bytes)
        reanalyze_line(src)

        # remove unreachable chunks from the calling function. this may happen in case our clone had
        # function chunks (e.g. it called other outlined functions that were inlined into it)
        src_func = sark.Function(src)
        for chunk_ea in unreachable_function_chunks_eas(src_func):
            ida_funcs.remove_func_tail(src_func._func, chunk_ea)

    # remove from storage
    del func_storage[src.ea]


def undo_inline_function(func):
    storage = ClonesStorage()
    func_storage = storage[func.ea]

    # pre-iterate the generator since we're deleting items inside
    for src_ea, clone_info in list(func_storage.items()):
        src = sark.Line(src_ea)
        undo_inline_function_call(src, func)

    undo_rename_outlined_function(func)


# FUNCTION EXPLORATION


def fix_function_noret_flags():
    found = False

    # pre-iterate since we might be adding functions inside
    for func in list(sark.functions()):
        for src_ea, target_ea in function_crefs(func):
            try:
                sark.Function(target_ea)
                continue
            except sark.exceptions.SarkNoFunction:
                pass

            # the target is not inside a function, ida probably failed to make one

            logger.debug(f"found call to non-function from {src_ea:#x} to {target_ea:#x}")

            fn = ida_funcs.func_t()
            fn.start_ea = target_ea
            ret = ida_funcs.find_func_bounds(fn, ida_funcs.FIND_FUNC_NORMAL)
            if ret != ida_funcs.FIND_FUNC_UNDEF:
                logger.debug(f"  finding function bounds unexpectedly succeeded: {ret}. skipping...")
                continue

            logger.debug(f"  making function failed because of undefined instruction @ {fn.end_ea:#x}")
            tail_call = sark.Line(fn.end_ea - 1)

            # we expect that to be because the last insn was a BL to a NORET function which isn't
            # marked as such

            if tail_call.insn.mnem != "BL":
                logger.debug("  didn't find a BL in the last instruction. skipping...")
                continue

            bl_target_ea = get_branch_target_ea(tail_call)
            logger.debug(f"  previous instruction is a BL to {bl_target_ea:#x} -> setting to NORET")

            try:
                bl_target_func = sark.Function(bl_target_ea)
            except sark.exceptions.SarkNoFunction:
                logger.debug("  BL target isn't a function. skipping...")
                continue
            if bl_target_func.ea != bl_target_ea:
                logger.debug("  BL target isn't the beginning of a function. skipping...")
                continue

            flags = idc.get_func_flags(bl_target_ea)
            if not (flags & ida_funcs.FUNC_NORET):
                idc.set_func_attr(bl_target_ea, idc.FUNCATTR_FLAGS, flags | ida_funcs.FUNC_NORET)
                ida_auto.plan_range(fn.start_ea, fn.end_ea)
                found = True
            else:
                logger.debug("  BL target was already set to NORET. skipping...")

    return found


def create_missing_functions():
    found = False

    # pre-iterate since we might be adding functions inside
    for func in list(sark.functions()):
        for src_ea, target_ea in function_crefs(func):
            try:
                sark.Function(target_ea)
                continue
            except sark.exceptions.SarkNoFunction:
                pass

            # the target is not inside a function, make one

            logger.debug(f"found call to non-function from {src_ea:#x} to {target_ea:#x} "
                         "-> making function")

            if ida_funcs.add_func(target_ea):
                found = True
            else:
                logger.debug("  failed to make function")

    return found


def is_data_heuristic(line):
    if line.is_code:
        return False

    # data with refs (e.g. jumptable)
    if list(line.drefs_to):
        return True

    # all 00s (alignment data)
    if not any(line.bytes):
        return True

    return False


# IDB PREPROCESSING


def detach_chunk(chunk_ea):
    # check whether we should dechunk this
    chunk = sark.Line(chunk_ea)
    if has_function_flow_xref(chunk):
        logger.debug(f"  found flow xref to {chunk_ea:#x} -> shouldn't dechunk")
        return None

    logger.debug(f"  detaching chunk @ {chunk_ea:#x}")
    chunk_end_ea = idc.get_fchunk_attr(chunk_ea, idc.FUNCATTR_END)

    # pre-iterate the generator since we're removing fchunks inside
    parents = list(function_chunk_parent_eas(chunk_ea))

    # remove the chunk from each of its parents
    for parent_ea in parents:
        logger.trace(f"    removing from parent @ {parent_ea:#x}")
        idc.remove_fchunk(parent_ea, chunk_ea)

    # create a function out of it
    ida_funcs.add_func(chunk_ea, chunk_end_ea)
    func = sark.Function(chunk_ea)

    # remove unreachable chunks from the parents. this may happen in case our detachee has had
    # additional function chunks
    for parent_ea in parents:
        parent_func = sark.Function(parent_ea)
        for chunk_ea in unreachable_function_chunks_eas(parent_func):
            ida_funcs.remove_func_tail(parent_func._func, chunk_ea)

    return func


def dechunk_functions():
    # pre-iterate since we're adding functions inside
    functions = list(sark.functions())

    for func in functions:
        if ida_kernwin.user_cancelled():
            return False

        if not is_originally_chunked_function(func):
            continue

        logger.debug(f"dechunking {func.name}...")

        # pre-iterate the generator since we're removing fchunks inside
        for chunk_ea in list(function_chunk_eas(func)):
            if chunk_ea == func.ea:
                continue  # skip first chunk

            if get_cloned_function(chunk_ea):
                continue  # skip inlined chunks

            new_func = detach_chunk(chunk_ea)

            # add the new function for processing
            if new_func:
                functions.append(new_func)


def split_outlined_function_trampolines():
    for l in sark.lines():
        try:
            insn = l.insn
        except sark.exceptions.SarkNoInstruction:
            continue  # nothing to do here...

        if insn.mnem != "B" or is_conditional_insn(insn):
            continue  # we're looking for unconditional branches

        target_ea = insn.operands[0].addr
        if target_ea != l.end_ea:
            continue  # to the next instruction

        # check if the next instruction is already marked as a different function
        try:
            src_func = sark.Function(l)
            target_func = sark.Function(target_ea)

            if src_func != target_func:
                continue  # it's already split
        except sark.exceptions.SarkNoFunction:
            continue  # nothing to split

        # check if there are external crefs to the next instruction (i.e. someone else has it as
        # a function chunk)
        target = sark.Line(target_ea)
        if not list(external_callers(target)):
            continue  # no external cref. no benefit in splitting

        logger.debug(f"splitting trampoline from adjacent function at {target_ea:#x}")

        # split the function after this branch
        end_ea = src_func.end_ea
        ida_funcs.set_func_end(l.ea, target_ea)
        ida_funcs.add_func(target_ea, end_ea)


def make_function_chunk(line):
    to_reprocess = []

    try:
        func = sark.Function(line)

        # check whether we need to split target's chunk
        chunk_start_ea = idc.get_fchunk_attr(line.ea, idc.FUNCATTR_START)
        should_split = chunk_start_ea != line.ea

        chunk_end_ea = idc.get_fchunk_attr(line.ea, idc.FUNCATTR_END)
    except sark.exceptions.SarkNoFunction:
        func = None
        should_split = None
        chunk_end_ea = ida_idaapi.BADADDR

    if should_split:
        ida_funcs.set_func_end(chunk_start_ea, line.ea)
        reanalyze_line(line.prev)
        to_reprocess.append(func)

    # it's important to first add it to the other callers and not the original function, otherwise
    # IDA will automatically merge it with the chunk we've just removed it from
    for caller in external_callers(line, functions_only=True, include_flow=True):
        for caller_func in containing_funcs(caller):
            if func != caller_func:
                ida_funcs.append_func_tail(caller_func._func, line.ea, chunk_end_ea)
                to_reprocess.append(caller_func)

    if should_split:
        # re-add it to the original function and set it as the owner
        ida_funcs.append_func_tail(func._func, line.ea, chunk_end_ea)
        ida_funcs.set_tail_owner(ida_funcs.get_fchunk(line.ea), func.ea)

    ida_auto.plan_range(line.ea, chunk_end_ea)

    return to_reprocess


def split_function(line):
    chunk_start_ea = idc.get_fchunk_attr(line.ea, idc.FUNCATTR_START)
    chunk_end_ea = idc.get_fchunk_attr(line.ea, idc.FUNCATTR_END)

    # verify that there's no flow xref into this line
    if has_function_flow_xref(line):
        # try adding it as a function tail to src
        logger.debug(f"  found flow xref to {line.ea:#x} -> cannot split function. making function "
                     "chunk instead")
        return make_function_chunk(line)

    # if there's a flow xref back into *all* of our callers, create a function chunk instead
    line_after_chunk = sark.Line(chunk_end_ea).next
    if has_function_flow_xref(line_after_chunk):
        # accumulate all of the funcs calling us (incl. a marker for non-func callers)
        calling_funcs = set()
        for caller in external_callers(line, include_flow=True):
            caller_funcs = containing_funcs(caller)
            if not caller_funcs:
                calling_funcs.add(None)
            else:
                calling_funcs.update(caller_funcs)

        if calling_funcs == containing_funcs(line_after_chunk):
            logger.debug("  found flow xref back to all of the the callers -> making function chunk"
                         " instead")
            return make_function_chunk(line)

    # split the function
    logger.debug(f"  splitting function chunk at {line.ea:#x}")

    ida_funcs.set_func_end(chunk_start_ea, line.ea)
    reanalyze_line(line.prev)

    ida_funcs.add_func(line.ea, chunk_end_ea)
    ida_auto.plan_range(line.ea, chunk_end_ea)

    return (sark.Function(line.ea), sark.Function(line.prev.ea))


def split_adjacent_functions():
    # pre-iterate since we're adding functions inside
    functions = list(sark.functions())

    while True:
        to_reprocess = []
        for func in functions:
            for src_ea, target_ea in function_crefs(func):
                src = sark.Line(src_ea)
                target = sark.Line(target_ea)

                # check whether we're jumping to the middle of another function
                try:
                    target_func = sark.Function(target.ea)
                except sark.exceptions.SarkNoFunction:
                    logger.debug(f"found flow xref from {src.ea:#x} to non-function.")
                    continue

                if target.ea == target_func.ea:
                    continue

                # if so, split the target function
                logger.debug(f"found call/branch to middle of function from {src.ea:#x} (as part of "
                             f"{func.name}) to {target.ea:#x}")
                more_to_reprocess = split_function(target)

                # add the new function for processing, and the current for reprocessing
                if more_to_reprocess:
                    to_reprocess.extend(more_to_reprocess)

        if to_reprocess:
            # wait for what we've previously done to finish analysing
            if not ida_auto.auto_wait():
                return False  # auto-analysis was cancelled

            # repeat with what we need to reprocess
            functions = to_reprocess
        else:
            break


@with_autoanalysis(False)
def explore_idb():
    logger.info("exploring IDB...")

    logger.debug("reanalyzing program...")
    reanalyze_program()

    exploration_steps = {
        # some times IDA misses some trivial functions on reanalysis, I'm not sure why, so we do
        # this ourselves as well
        "creating missing functions...": create_missing_functions,
        "fixing missing NORET flags on functions...": fix_function_noret_flags,
    }

    for i in itertools.count():
        found = False
        with wait_box(f"exploring (iteration {i})..."):
            for msg, func in exploration_steps.items():
                logger.debug("waiting for auto-analysis to complete...")
                if not ida_auto.auto_wait():
                    return False  # auto-analysis was cancelled

                logger.info(f"{msg} (iteration {i})")
                found = found or func()

                if ida_kernwin.user_cancelled():
                    return False
        if found:
            reanalyze_program()
        else:
            break

    logger.info("exploration done!")
    return True


@with_autoanalysis(False)
def preprocess_idb():
    logger.info("preprocessing IDB...")

    preprocessing_steps = {
        "dechunking functions...": dechunk_functions,
        "splitting trampolines from adjacent functions...": split_outlined_function_trampolines,
        "splitting adjacent functions...": split_adjacent_functions,
    }

    with wait_box("preprocessing..."):
        for msg, func in preprocessing_steps.items():
            logger.debug("waiting for auto-analysis to complete...")
            if not ida_auto.auto_wait():
                return False  # auto-analysis was cancelled

            logger.info(msg)
            func()

            if ida_kernwin.user_cancelled():
                return False

    reanalyze_program()

    logger.info("preprocessing done!")
    return True


# OUTLINED FUNCTION FINDING


def code_flow_iterator(line, forward=True, stop=None, abort_on_calls=True, dfs=False,
                       ignore_aborts=False, _visited_eas=None):

    func = sark.Function(line)
    if stop is None:
        if forward:
            stop = sark.Line(func.end_ea - 1)
        else:
            stop = sark.Line(func.start_ea)

    if _visited_eas is None:
        _visited_eas = set()

    while True:
        if line.ea in _visited_eas:
            break
        else:
            _visited_eas.add(line.ea)

        yield line

        # if we encounter a BL or a tail-call, this may be to an actual outlined function, and
        # therefore analysis following this code flow shouldn't treat it as an ABI-complaint
        # function and "skip" it
        if abort_on_calls:
            if any(x.type.is_code and not ida_funcs.func_contains(func._func, x.to) for x in
                   line.xrefs_from):
                raise FunctionInlinerUnknownFlowException()

        if line == stop:
            break

        if forward:
            crefs = [x.to for x in line.xrefs_from if
                     x.type.is_code and ida_funcs.func_contains(func._func, x.to)]
        else:
            crefs = [x.frm for x in line.xrefs_to if
                     x.type.is_code and ida_funcs.func_contains(func._func, x.frm)]

        if len(crefs) == 0:
            break

        if len(crefs) > 1 and ignore_aborts:
            crefs = [x for x in crefs if sark.Line(x).insn.mnem != "BRK"]

        if len(crefs) > 1:
            if not dfs:
                raise FunctionInlinerUnknownFlowException()

            for cref in crefs[:-1]:
                yield from code_flow_iterator(
                    sark.Line(cref),
                    stop=stop,
                    forward=forward,
                    abort_on_calls=abort_on_calls,
                    dfs=dfs,
                    ignore_aborts=ignore_aborts,
                    _visited_eas=_visited_eas)

            crefs = crefs[-1:]

        line = sark.Line(crefs[0])


def find_function_ends(func, ignore_aborts=False):
    def is_internal(ea):
        return ida_funcs.func_contains(func._func, ea)

    def finder(ea, visited):
        end_eas = set()

        while ea not in visited:
            visited.add(ea)
            next_eas = [next_ea for next_ea in sark.Line(ea).crefs_from if is_internal(next_ea)]

            # ignore aborts if asked to
            if ignore_aborts and sark.Line(ea).insn.mnem == "BRK":
                break

            # handle ret/tail-call
            if len(next_eas) == 0:
                end_eas.add(ea)
                break

            # handle branch to self
            if len(next_eas) == 1 and next_eas[0] == ea:
                end_eas.add(ea)
                break

            for next_ea in next_eas[1:]:
                end_eas |= finder(next_ea, visited)

            ea = next_eas[0]

        return end_eas

    end_eas = finder(func.start_ea, set())
    return map(sark.Line, end_eas)


def is_function_prologue(line):
    # check LR is signed (relevant only ARMv8.3 code which uses it)
    if line.disasm == "PACIBSP":
        return True

    # check for BTI (relevant only ARMv8.5 code which uses it)
    if line.insn.mnem == "BTI" and line.insn.operands[0].text == "c":
        return True

    return False

def is_function_stack_prologue(line):
    """
    we used to have this logic as part of is_function_prologue() but nowadays stack preparation
    is often outlined as well, so we don't use it
    """

    # expect stack space to be allocated
    insn = line.insn
    ops = insn.operands

    potential_stored_regs = set((f"X{i}" for i in range(19, 31))) | set(("FP", "LR"))

    if insn.mnem == "SUB" and ops[0].text == "SP" and ops[1].text == "SP":
        # stack space is allocated explicitly
        stack_space = ops[-1].value
    elif insn.mnem in ("STP", "STR") and ops[-1].base == "SP" and insn.indexing_mode.is_pre and \
            all(o.text in potential_stored_regs for o in ops[:-1]):

        # stack space is allocated inline with the first store
        stack_space = struct.unpack("<q", struct.pack("<Q", ops[-1].displacement))[0]
    else:
        return False

    # expect stack space to be 16-bytes aligned
    if stack_space & 0xf:
        return False

    # stack space is not always fully filled (since maybe some was reserved for the local frame), so
    # we don't check for that

    # FP is not always set, so we don't check for that

    return True


def get_op_regs(op):
    # regular vector operands will be shown as Qn in op.regs, so we don't have
    # to treat them specially. multiple-registers are arch specific so they're not
    # parsed by sark and we have to resolve them on our own
    if op.type.is_special and op.text.startswith("{V"):
        m = re.fullmatch(r"\{V(\d+)\.\d+.\}", op.text)
        if m:
            reg = f"Q{m.group(1)}"
            return set((reg,))
        else:
            m = re.fullmatch(r"\{V(\d+)\.\d+.-V(\d+)\.\d+.\}", op.text)
            assert m
            start, end = m.groups()
            return set((f"Q{n}" for n in range(int(start), int(end) + 1)))
    # some times (e.g. register with optional shift), Sark can't parse the register properly
    elif op.type.is_special and any(t in op.text for t in (',ASR#', ',LSL#', ',LSR#', ',ROR#')):
        reg = sark.base.get_register_name(op.reg_id)
        return set((reg,))
    else:
        return op.regs


def is_insn_using_condition_flags(insn):
    if insn.mnem in ("ADC", "CSEL", "CSINC", "CSINV", "CSNEG", "CSET", "CSETM", "CINC", "CINV",
                     "CNEG", "SBC", "NGC", "FCSEL", "VSEL"):
        return True
    else:
        return is_conditional_insn(insn)


def is_insn_setting_condition_flags(insn):
    if insn.mnem in ("TST", "CMP", "CMN", "CCMP", "CCMN", "FCMP", "FCCMP", "FCMPE", "FCCMPE",
                     "VCMP", "VCMPE"):
        return True
    else:
        # is using the S suffix
        return insn._insn.auxpref & 1  # see module/arm/arm.hpp in the IDA SDK


def get_fake_condition_flags_ops(insn):
    # check whether this instruction uses/affects condition flags
    is_read = is_insn_using_condition_flags(insn)
    is_write = is_insn_setting_condition_flags(insn)

    if not (is_read or is_write):
        return tuple()

    # we don't bother with creating a fake operand per each read/set condition flag
    # so create an operand reading/writing a fake NZCV reg

    # create a fake IDA op_t
    fake_op_t = types.SimpleNamespace(
        n=-1,
        type=ida_ua.o_idpspec5 + 1,  # simulated as o_cond in arm.hpp, Sark will treat as "special"
        addr=0,
        value=0,
        flags=0,
        dtype=ida_ua.dt_void,
        reg=-1
    )

    # create a fake Sark Operand with fake reg/regs properties
    class OperandWithFakeReg(sark.code.instruction.Operand):
        def __init__(self, fake_reg, *args, **kwargs):
            super().__init__(*args, **kwargs)
            self._fake_reg = fake_reg

        @sark.code.instruction.Operand.reg.getter
        def reg(self):
            return self._fake_reg

        @sark.code.instruction.Operand.regs.getter
        def regs(self):
            return set((self._fake_reg,))

    op = OperandWithFakeReg("NZCV",
                            fake_op_t,
                            insn._ea,
                            insn=insn._insn,
                            write=is_write,
                            read=is_read)

    return (op,)


def get_insn_ops_with_condition_flags(insn):
    # we create fake operands for condition flags used/set by the instruction
    yield from itertools.chain(insn.operands, get_fake_condition_flags_ops(insn))


def is_function_using_uninitialized_regs(func):
    logger.trace(f"analyzing {func.name} for uninitialized reg usage")

    initialized_regs = set(("SP", "X30", "LR", "WZR", "XZR"))

    # assume arg registers are initialized
    for i in range(8):
        initialized_regs.add(f"W{i}")
        initialized_regs.add(f"X{i}")

    # callee saved regs can be STR/STP-ed but nothing else
    callee_saved_allowed_mnems = {
        "STR": "LDR",
        "STP": "LDP",
    }
    potential_callee_saved_regs = set((f"X{i}" for i in range(19, 30)))
    callee_saved_regs = dict()  # matches stored reg to (store ea, store mnem)

    # look for uninitialized reg reads
    try:
        for l in code_flow_iterator(sark.Line(func.start_ea), ignore_aborts=True):
            logger.trace(f"-> {l.disasm}")

            # regs that we'll treat as initialized once done with this instruction
            new_initialized_regs = set()

            insn = l.insn
            for op in get_insn_ops_with_condition_flags(insn):
                regs = get_op_regs(op)

                # special treatment for zeroing regs with EOR
                is_eor = insn.mnem == "EOR" and \
                    insn.operands[0].text == insn.operands[1].text == insn.operands[2].text

                for r in regs:
                    if op.type.is_displ:
                        # if we're reading an initialized reg, it's all good
                        if r in initialized_regs:
                            continue

                        logger.trace(f"   found uninitialized use of {r}")
                        return True

                    if op.is_read and not is_eor:
                        # if we're reading an initialized reg, it's all good
                        if r in initialized_regs:
                            continue

                        # if an uninitialized reg is being read it's either that or that it's actually a
                        # callee saved reg, and the we're reading it to store it somewhere and reuse that
                        # register internally
                        if r in potential_callee_saved_regs and insn.mnem in callee_saved_allowed_mnems:
                            # if it's the latter case -- we'll remember it
                            callee_saved_regs[r] = (l.ea, insn.mnem)
                            logger.trace(f"   marking {r} as callee saved")
                            continue

                        logger.trace(f"   found uninitialized use of {r}")
                        return True

                    elif op.is_write:
                        if r in initialized_regs:
                            continue

                        # special treatment for registers who aren't can't be split into parts
                        if r == "NZCV":
                            logger.trace(f"   marking {r} as initialized")
                            new_initialized_regs.add(r)
                            continue

                        # mark all of the "parts" of this register as initialized
                        parts = list(register_parts(r))
                        parts_s = ", ".join(parts)
                        logger.trace(f"   marking {{{parts_s}}} as initialized")
                        new_initialized_regs.update(parts)

            initialized_regs |= new_initialized_regs
    except FunctionInlinerUnknownFlowException:
        logger.trace("aborting because code flow cannot be followed any longer...")
        return False

    # verify that callee saved regs are actually restored in the end
    # if not -- they were really just used and not stored
    if not callee_saved_regs:
        return False

    logger.trace("finished forward pass ; starting backwards pass")

    # some functions have multiple ends. in case one of them points to a noret function,
    # the callee saved regs may not be restored, so in case there are multiple function ends
    # we'll just skip this validation
    func_ends = list(find_function_ends(func))
    if len(func_ends) > 1:
        logger.trace("aborting because more than one function end was found...")
        return False

    # we don't have to wrap this code_flow_iterator with try-except, since if the forward-pass
    # didn't find any basic blocks, the backwards pass surely won't
    for l in code_flow_iterator(func_ends[0], forward=False, ignore_aborts=True):
        logger.trace(f"-> {l.disasm}")

        insn = l.insn
        for op in get_insn_ops_with_condition_flags(insn):
            for r in op.regs:
                if r not in callee_saved_regs:
                    continue

                store_ea, store_mnem = callee_saved_regs[r]
                restore_mnem = callee_saved_allowed_mnems[store_mnem]

                # the last "use" of the callee saved reg must be to restore it, mnem must match
                # the storing mnem, and the ea must be after that of the storing ea
                if not op.is_write or insn.mnem != restore_mnem or l.ea < store_ea:
                    logger.trace(f"   last use of callee saved {r} isn't restoring")
                    return True

                # it was properly restored -> forget about it
                logger.trace(f"   found restore of callee saved {r}")
                del callee_saved_regs[r]
                if not callee_saved_regs:
                    return False

    return False


def is_function_affecting_non_result_regs(func, *, allow_multiple_result_registers=True):
    logger.trace(f"analyzing {func.name} for non-result-registers effects")

    # these special regs can also be affected in the function epilogue
    result_regs = set(("SP", "X29", "X30", "LR", "WZR", "XZR"))

    n = 8 if allow_multiple_result_registers else 1
    for i in range(n):
        result_regs.add(f"W{i}")
        result_regs.add(f"X{i}")

    # callee saved regs can be STR/STP-ed but nothing else
    callee_saved_allowed_mnems = {
        "LDR": "STR",
        "LDP": "STP",
    }
    potential_callee_saved_regs = set((f"X{i}" for i in range(19, 29)))
    callee_saved_regs = dict()  # matches restored reg to (restore ea, restore mnem)

    # some functions have multiple ends. in case one of them points to a noret function,
    # we may be in the midst of a function, and see internally registers being used, so we'll just
    # skip this validation
    func_ends = list(find_function_ends(func, ignore_aborts=True))
    if len(func_ends) > 1:
        logger.trace("aborting because more than one function end was found...")
        return False

    # look for "useless" writes into non-result regs
    try:
        for l in code_flow_iterator(func_ends[0], forward=False, ignore_aborts=True):
            logger.trace(f"-> {l.disasm}")

            # regs that we'll treat as result regs once done with this instruction
            new_result_regs = set()

            insn = l.insn

            for op in get_insn_ops_with_condition_flags(insn):
                regs = get_op_regs(op)

                for r in regs:
                    if op.is_write and not op.type.is_displ:
                        # if we're writing into a result reg, it's all good
                        if r in result_regs:
                            continue

                        # if a non-result reg is being written into it's either that or that it's actually a
                        # callee saved reg, and the we're restoring its value here
                        if r in potential_callee_saved_regs and insn.mnem in callee_saved_allowed_mnems:
                            # if it's the latter case -- we'll remember it
                            callee_saved_regs[r] = (l.ea, insn.mnem)
                            logger.trace(f"   marking {r} as callee saved")
                            continue

                        logger.trace(f"   found write into non-result reg {r}")
                        return True
                    elif op.is_read or op.type.is_displ:
                        if r in result_regs:
                            continue

                        # if we got here and won't return True on anything else about this instruction,
                        # it means that this non-result reg is read and is either stored into memory
                        # or affects only-result regs. in both of these cases, we should treat it as
                        # a result reg as well

                        # special treatment for registers who aren't can't be split into parts
                        if r == "NZCV":
                            logger.trace(f"   marking {r} as result reg")
                            new_result_regs.add(r)
                            continue

                        # mark all of the "parts" of this register as result regs
                        logger.trace(f"   marking *{r[1:]} as result regs")
                        new_result_regs.update(register_parts(r))

            result_regs |= new_result_regs

    except FunctionInlinerUnknownFlowException:
        logger.trace("aborting because code flow cannot be followed any longer...")
        return False

    # verify that callee saved regs are actually stored in the beginning
    # if not -- they were really just written to and not restored
    if not callee_saved_regs:
        return False

    logger.trace("finished backwards pass ; starting forward pass")

    try:
        for l in code_flow_iterator(sark.Line(func.start_ea), ignore_aborts=True):
            logger.trace(f"-> {l.disasm}")

            insn = l.insn
            for op in get_insn_ops_with_condition_flags(insn):
                for r in op.regs:
                    if r not in callee_saved_regs:
                        continue

                    restore_ea, restore_mnem = callee_saved_regs[r]
                    store_mnem = callee_saved_allowed_mnems[restore_mnem]

                    # the first "use" of the callee saved reg must be to store it, mnem must match
                    # the restoring mnem, and the ea must be before that of the restoring ea
                    if not op.is_read or not insn.mnem == store_mnem or l.ea > restore_ea:
                        logger.trace(f"   first use of callee saved {r} isn't storing")
                        return True

                    # it was properly stored -> forget about it
                    logger.trace(f"   found store of callee saved {r}")
                    del callee_saved_regs[r]
                    if not callee_saved_regs:
                        return False

    except FunctionInlinerUnknownFlowException:
        logger.trace("aborting because code flow cannot be followed any longer...")

    return False


def is_function_using_uninitialized_stack(func, stack_args_size=0):
    """
    note these accesses
    """

    logger.trace(f"analyzing {func.name} for stack accesses outside the stack frame")

    size = ida_frame.get_frame_size(func._func) + stack_args_size
    logger.trace(f"   frame size: {size:#x}")

    try:
        for l in code_flow_iterator(sark.Line(func.start_ea), ignore_aborts=True):
            logger.trace(f"-> {l.disasm}")

            insn = l.insn

            if not insn.has_reg("SP"):
                continue

            for i in range(len(insn.operands)):
                off = ida_frame.calc_stkvar_struc_offset(func._func, insn._insn, i)

                if off == ida_idaapi.BADADDR:
                    continue

                if off > size:
                    logger.trace(f"   accessed stack frame at offset {off:#x} -> out of bounds")
                    return True
                else:
                    logger.trace(f"   accessed stack frame at offset {off:#x}")
    except FunctionInlinerUnknownFlowException:
        logger.trace("aborting because code flow cannot be followed any longer...")

    return False


def is_function_outlined(func, *, include_inlined=False, extra=False):
    if include_inlined:
        if func.ea in ClonesStorage().items:
            return True

    # nothing to do here if we have no callers
    if not list(external_callers(func)):
        return False

    # i'm not really sure about whether outlined functions never have prologues, but we'll see
    if is_function_prologue(sark.Line(func.ea)):
        return False

    if is_function_using_uninitialized_regs(func):
        return True

    # most programs don't really return structs from functions, if "extra" is set, we will assume
    # that X1-X7 aren't really used as return registers
    if is_function_affecting_non_result_regs(func, allow_multiple_result_registers=not extra):
        return True

    if not extra:
        return False

    # this is under "extra" since we don't know how many args func gets so we can't really say
    # if a stack access is out-of-bounds or not
    #
    # however, we assume that a function with more than 8 arguments will have a proper prologue
    # or stack frame set up
    if (not is_function_stack_prologue(sark.Line(func.ea))
        and is_function_using_uninitialized_stack(func)):
        return True

    return False


def find_next_reg_use(line, reg):
    reg_parts = set(register_parts(reg))

    for l in code_flow_iterator(line, abort_on_calls=False, dfs=True):
        for op in l.insn.operands:
            if not op.regs & reg_parts:
                continue
            if op.is_write:
                return None
            else:
                return l


def apply_code_patch(start_ea, end_ea, code, kp_asm=None):
    if kp_asm is None:
        kp_asm = keypatch.Keypatch_Asm()

    size = end_ea - start_ea
    assert len(code) <= size

    if len(code) < size:
        nop = bytes(kp_asm.assemble("NOP", 0)[0])
        nop_slide_size = size - len(code)
        assert nop_slide_size % len(nop) == 0
        code += nop * (nop_slide_size // len(nop))

    assert len(code) == size
    ida_bytes.patch_bytes(start_ea, code)
    ida_auto.plan_range(start_ea, end_ea)


def patch_constant_BRs(kp_asm=None):
    """patches snippets of the form:
                ADR/L     Xn, sub_1337
                NOP/-
                BR/BLR    Xn

        to:
                B/BL      sub_1337
    """
    if kp_asm is None:
        kp_asm = keypatch.Keypatch_Asm()

    count = 0
    retval = True
    for l1, l2, l3 in linegroups(3):
        if ida_kernwin.user_cancelled():
            retval = False
            break

        # check if we're at a constant BR
        try:
            if l1.insn.mnem not in ("ADR", "ADRL"):
                continue

            if l1.insn.mnem == "ADR":
                if l2.insn.mnem != "NOP":
                    continue
            else:  # ADRL spans 8 bytes and hence no NOP
                l3 = l2  # align both cases' line numbers

            if l3.insn.mnem not in ("BR", "BLR"):
                continue
        except sark.exceptions.SarkNoInstruction:
            continue

        target_ea = l1.insn.operands[1].value

        r = l1.insn.operands[0].text
        if r != l3.insn.operands[0].text:
            continue

        logger.debug(f"found constant BR to {target_ea:#x} at {l1.ea:#x}")

        # verify that the register isn't used in the rest of the function
        l = find_next_reg_use(l3.next, r)
        if l is not None:
            logger.debug(f"  constant BR is unpatchable because there's another ref to {r} at {l.ea:#x}")
            continue

        # patch to a standard call
        call_mnem = l3.insn.mnem[:-1]
        asm = f"{call_mnem} #{target_ea:#x}"  # we drop PAC flags
        code = bytes(kp_asm.assemble(asm, l1.ea)[0])
        apply_code_patch(l1.ea, l3.end_ea, code, kp_asm)

        add_comment(l1, f"FunctionInliner: patched from constant BR using {r}")

        logger.debug("  patched")
        count += 1

    logger.info(f"patched {count} constant BRs")
    return retval


def patch_constant_tested_BRs(kp_asm=None):
    """patches snippets of the form:
                ADR/L     Xn, sub_1337
                NOP/-
                CBNZ      Xn, do_call
                B         dont_call
            do_call:
                BR/BLR    Xn
            dont_call:

        to:
                B/BL      sub_1337
    """
    if kp_asm is None:
        kp_asm = keypatch.Keypatch_Asm()

    count = 0
    retval = True
    for l1, l2, l3, l4, l5 in linegroups(5):
        if ida_kernwin.user_cancelled():
            retval = False
            break

        # check if we're at a constant BR
        try:
            if l1.insn.mnem not in ("ADR", "ADRL"):
                continue

            if l1.insn.mnem == "ADR":
                if l2.insn.mnem != "NOP":
                    continue
            else:  # ADRL spans 8 bytes and hence no NOP
                l3, l4, l5 = l2, l3, l4  # align both cases' line numbers

            target_ea = l1.insn.operands[1].value
            r = l1.insn.operands[0].text

            if (l3.insn.mnem != "CBNZ"
               or l3.insn.operands[0].text != r  # noqa: W503
               or l3.insn.operands[1].addr != l5.ea):  # noqa: W503
                continue

            if (l4.insn.mnem != "B"
               or l4.insn.operands[0].addr != l5.end_ea):  # noqa: W503
                continue

            if (l5.insn.mnem not in ("BR", "BLR")
               or l5.insn.operands[0].text != r):  # noqa: W503
                continue
        except sark.exceptions.SarkNoInstruction:
            continue

        logger.debug(f"found constant tested BR to {target_ea:#x} at {l1.ea:#x}")

        try:
            sark.Function(l5.next)
        except sark.exceptions.SarkNoFunction:
            logger.debug(f"  constant tested BR is unpatchable because this is not a function and so"
                         f" we couldn't verify whether there's another ref to {r}")
            continue

        # verify that the register isn't used in the rest of the function
        l = find_next_reg_use(l5.next, r)
        if l is not None:
            logger.debug(f"  constant tested BR is unpatchable because there's another ref to {r} at {l.ea:#x}")
            continue

        # patch to a standard call
        call_mnem = l5.insn.mnem[:-1]
        asm = f"{call_mnem} #{target_ea:#x}"  # we drop PAC flags
        code = bytes(kp_asm.assemble(asm, l1.ea)[0])
        apply_code_patch(l1.ea, l5.end_ea, code, kp_asm)

        add_comment(l1, f"FunctionInliner: patched from constant tested BR using {r}")

        logger.debug("  patched")
        count += 1

    logger.info(f"patched {count} constant tested BRs")
    return retval


def patch_constant_data_BLRs(kp_asm=None):
    """patches snippets of the form:
                NOP/ADRP
                LDR       Xn, =sub_1337
                BLR       Xn

        where the data lives in a const segment

        to:
                BL      sub_1337
    """
    if kp_asm is None:
        kp_asm = keypatch.Keypatch_Asm()

    count = 0
    retval = True
    for l1, l2, l3 in linegroups(3):
        if ida_kernwin.user_cancelled():
            retval = False
            break

        # check if we're at a constant BR
        try:
            if l2.insn.mnem != "LDR":
                continue

            if l1.insn.mnem == "NOP":
                if not l2.insn.operands[1].type.is_mem:
                    continue
            elif l1.insn.mnem == "ADRP":
                if not l2.insn.operands[1].type.is_displ:
                    continue
                if not l2.insn.operands[1].reg == l1.insn.operands[0].reg:
                    continue
            else:
                continue

            r = l2.insn.operands[0].reg

            if l3.insn.mnem != "BLR" or l3.insn.operands[0].reg != r:
                continue
        except sark.exceptions.SarkNoInstruction:
            continue

        # resolve the loaded addr
        drefs_from = set(drefs_from_eas(l2))
        if len(drefs_from) == 1:
            p_target_ea = drefs_from.pop()
        else:  # this may happen with LDR when the target contains another address
            assert len(drefs_from) == 2 and l2.insn.mnem == "LDR"
            xref = [x for x in l2.xrefs_from if x.type.is_read][0]
            p_target_ea = xref.to

        # skip pointers which aren't in __auth_ptr or in const data segments
        p_target_seg_name = sark.Segment(p_target_ea).name
        if p_target_seg_name != "__auth_ptr" and "const" not in p_target_seg_name.lower():
            continue

        target_ea = ida_bytes.get_qword(p_target_ea)

        logger.debug(f"found constant data BLR to {target_ea:#x} at {l1.ea:#x}")

        # verify that the register isn't used in the rest of the function
        l = find_next_reg_use(l3.next, r)
        if l is not None:
            logger.debug(f"  constant data BLR is unpatchable because there's another ref to {r} at {l.ea:#x}")
            continue

        # patch to a standard call
        asm = f"BL #{target_ea:#x}"  # we drop PAC flags
        code = bytes(kp_asm.assemble(asm, l1.ea)[0])
        apply_code_patch(l1.ea, l3.end_ea, code, kp_asm)

        add_comment(l1, f"FunctionInliner: patched from constant data BLR using {r}")

        logger.debug("  patched")
        count += 1

    logger.info(f"patched {count} constant data BLRs")
    return retval


# PLUGIN STUFF


class FunctionInlinerActionBase(ida_kernwin.action_handler_t):
    def __init__(self, plugin):
        super().__init__()
        self.plugin = plugin

    @property
    def name(self):
        return f"{self.plugin.wanted_name}:{self.__class__.__name__}"

    @property
    def label(self):
        return "Inline function under cursor"

    @property
    def shortcut(self):
        return None

    @property
    def tooltip(self):
        return None

    @property
    def icon(self):
        return ida_kernwin.get_action_icon("MakeFunction")[1]

    @property
    def flags(self):
        return 0

    @property
    def menu_flags(self):
        return 0

    @property
    def path(self):
        return f"Edit/Plugins/{self.plugin.wanted_name}/"

    def register(self):
        desc = ida_kernwin.action_desc_t(
            self.name,
            self.label,
            self,
            self.shortcut,
            self.tooltip,
            self.icon,
            self.flags
        )
        ida_kernwin.register_action(desc)

    def unregister(self):
        ida_kernwin.unregister_action(self.name)

    def activate(self, ctx):
        raise NotImplementedError()

    def update(self, ctx):
        raise NotImplementedError()


class FunctionInlinerInlineAction(FunctionInlinerActionBase):
    @property
    def label(self):
        return "Inline function"

    @property
    def shortcut(self):
        return "Meta-P"

    def activate(self, ctx):
        func = get_function_under_cursor()
        inline_function(func)
        return 1

    def update(self, ctx):
        f = get_function_under_cursor()
        if f and list(external_callers(f)):
            return ida_kernwin.AST_ENABLE
        else:
            return ida_kernwin.AST_DISABLE


class FunctionInlinerUndoInlineAction(FunctionInlinerActionBase):
    @property
    def label(self):
        return "Undo function inlining"

    def activate(self, ctx):
        outlined_func = get_inlined_function_under_cursor()
        undo_inline_function(outlined_func)
        return 1

    def update(self, ctx):
        if get_inlined_function_under_cursor():
            return ida_kernwin.AST_ENABLE
        else:
            return ida_kernwin.AST_DISABLE


class FunctionInlinerInlineAllAction(FunctionInlinerActionBase):
    @property
    def label(self):
        return "Inline all outlined functions"

    @property
    def menu_flags(self):
        return ida_kernwin.SETMENU_ENSURE_SEP

    def activate(self, ctx):
        if not explore_idb():
            return 1
        if not preprocess_idb():
            return 1
        if not inline_all_functions(extra=self.plugin.inline_extra):
            return 1
        reanalyze_program()
        return 1

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


class FunctionInlinerInlineAllActionNoPreprocessing(FunctionInlinerActionBase):
    @property
    def label(self):
        return "Inline all outlined functions (no preprocessing)"

    def activate(self, ctx):
        if not inline_all_functions(extra=self.plugin.inline_extra):
            return 1
        reanalyze_program()
        return 1

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


class FunctionInlinerToggleExtra(FunctionInlinerActionBase):
    @property
    def label(self):
        return "Toggle inlining of extra functions"

    @property
    def flags(self):
        flags = ida_kernwin.ADF_CHECKABLE

        if self.plugin.inline_extra:
            flags |= ida_kernwin.ADF_CHECKED

        return flags

    @property
    def menu_flags(self):
        return ida_kernwin.SETMENU_ENSURE_SEP

    @property
    def icon(self):
        return 0

    def activate(self, ctx):
        self.plugin.inline_extra = ida_kernwin.get_action_checked(self.name)[1]
        return 1

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


class FunctionInlinerPatchConstantBLRs(FunctionInlinerActionBase):
    @property
    def label(self):
        return "Patch constant register-based calls to regular calls"

    @property
    def icon(self):
        return 0

    def activate(self, ctx):
        if not explore_idb():
            return 1
        with wait_box("patching constant BRs..."):
            if not patch_constant_BRs():
                return 1
            if not patch_constant_tested_BRs():
                return 1
            if not patch_constant_data_BLRs():
                return 1
        return 1

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


class FunctionInlinerHooks(ida_kernwin.UI_Hooks):
    def __init__(self, ctx_actions, menu_actions):
        super().__init__()

        self.ctx_actions = ctx_actions
        self.menu_actions = menu_actions

    def ready_to_run(self):
        for action in self.menu_actions:
            ida_kernwin.attach_action_to_menu(
                action.path, action.name, ida_kernwin.SETMENU_APP | action.menu_flags
            )

    def finish_populating_widget_popup(self, form, popup):
        if ida_kernwin.get_widget_type(form) in (ida_kernwin.BWN_DISASM, ida_kernwin.BWN_PSEUDOCODE):
            ida_kernwin.attach_action_to_popup(form, popup, "-", None, ida_kernwin.SETMENU_FIRST)
            for action in reversed(self.ctx_actions):
                ida_kernwin.attach_action_to_popup(
                    form, popup, action.name, None, ida_kernwin.SETMENU_FIRST | action.menu_flags
                )


class FunctionInlinerPlugin(ida_idaapi.plugin_t):
    version = ida_idp.IDP_INTERFACE_VERSION
    flags = ida_idaapi.PLUGIN_MOD | ida_idaapi.PLUGIN_HIDE

    comment = "inlines functions that were outlined"
    help = ""
    wanted_name = "FunctionInliner"
    wanted_hotkey = ""

    ctx_actions_types = (FunctionInlinerInlineAction, FunctionInlinerUndoInlineAction)
    menu_actions_types = (
        FunctionInlinerPatchConstantBLRs, FunctionInlinerInlineAllAction,
        FunctionInlinerInlineAllActionNoPreprocessing,FunctionInlinerToggleExtra
    )

    NETNODE = "$ FunctionInliner.plugin"

    @staticmethod
    def init_logging():
        logger_formatter = logging.Formatter(fmt="{name}.{levelname:<5s}: {message}", style="{")

        class TqdmHandler(logging.StreamHandler):
            def emit(self, record):
                msg = self.format(record)
                tqdm.tqdm.write(msg, nolock=True)

        logger_hdlr = TqdmHandler()
        logger_hdlr.setFormatter(logger_formatter)

        logger.addHandler(logger_hdlr)

        if TRACE:
            logger.setLevel(logging.TRACE)
        else:
            logger.setLevel(logging.INFO)

    def is_compatible(self):
        info = ida_idaapi.get_inf_structure()
        return info.procname == "ARM" and info.is_64bit()

    def init(self):
        FunctionInlinerPlugin.init_logging()

        self._netnode = netnode.Netnode(self.NETNODE)

        self.ctx_actions = []
        self.menu_actions = []
        self.hooks = None

        if not self.is_compatible():
            logger.error("IDB deemed unsuitable (not an ARM64 binary). Skipping...")
            return ida_idaapi.PLUGIN_SKIP

        for t in FunctionInlinerPlugin.ctx_actions_types:
            a = t(self)
            a.register()
            self.ctx_actions.append(a)

        for t in FunctionInlinerPlugin.menu_actions_types:
            a = t(self)
            a.register()
            self.menu_actions.append(a)

        self.hooks = FunctionInlinerHooks(self.ctx_actions, self.menu_actions)
        self.hooks.hook()

        logger.info("initialized successfully")

        return ida_idaapi.PLUGIN_KEEP

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

        for a in self.ctx_actions:
            a.unregister()

        for a in self.menu_actions:
            a.unregister()

    def run(self, arg=0):
        pass

    @property
    def inline_extra(self):
        return bool(self._netnode.get("inline_extra"))

    @inline_extra.setter
    def inline_extra(self, value: bool):
        self._netnode["inline_extra"] = value


def PLUGIN_ENTRY():
    return FunctionInlinerPlugin()


repository: pfalcon/idapython
file: Scripts/msdnapihelp.py
description: Automatically exported from code.google.com/p/idapython
wanted_name: MSDN API Help
comment: Online MSDN API Help
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/msdnapihelp.py

content:
"""
User contributed script: MSDN API HELP plugin

This script fetches the API reference (from MSDN) of a given highlighted identifier
and returns the results in a new web browser page.

This script depends on the feedparser package: http://code.google.com/p/feedparser/

10/05/2010
- initial version


"""

import idaapi

# -----------------------------------------------------------------------
class msdnapihelp_plugin_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = "Online MSDN API Help"
    help = "Help me"
    wanted_name = "MSDN API Help"
    wanted_hotkey = "F3"

    def init(self):
        return idaapi.PLUGIN_OK


    @staticmethod
    def sanitize_name(name):
        t = idaapi.FUNC_IMPORT_PREFIX
        if name.startswith(t):
            return name[len(t):]
        return name


    def run(self, arg):
        # Get the highlighted identifier
        id = idaapi.get_highlighted_identifier()
        if not id:
            print "No identifier was highlighted"
            return

        import webbrowser

        try:
            import feedparser
        except:
            idaapi.warning('Feedparser package not installed')
            return

        id = self.sanitize_name(id)
        print "Looking up '%s' in MSDN online" % id
        d = feedparser.parse("http://social.msdn.microsoft.com/Search/Feed.aspx?locale=en-us&format=RSS&Query=%s" % id)
        if len(d['entries']) > 0:
            url = d['entries'][0].link
            webbrowser.open_new_tab(url)
        else:
            print "API documentation not found for: %s" % id


    def term(self):
        pass


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


repository: junron/auto-enum
file: plugin/auto_enum.py
description: IDA/Binary Ninja Plugin to automatically identify and set enums for standard functions
created_at: 2024-05-05 13:43:44+00:00
pushed_at: 2025-01-16 22:22:27+00:00
url:
https://github.com/junron/auto-enum/blob/75c2f88a1d9e1b2604d86444cdca2b3e5369c437/plugin/auto_enum.py

content:
def PLUGIN_ENTRY():
    from enumlib.ui_plugins.ida_plugin import AutoEnumPlugin
    return AutoEnumPlugin()

try:
    import idaapi
except:
    try:
        import binaryninja as bn
        from enumlib.ui_plugins.binja_plugin import auto_enum
        bn.PluginCommand.register("Auto Enum", "Automatically detect standard enums", auto_enum)
    except:
        pass


repository: elastic/HexForge
file: HexForge.py
url: https://github.com/elastic/HexForge/blob/b8695dbb06643b03e7405e6646c5bc4a469c7c2e/HexForge.py
wanted_name: HexForge
created_at: 2024-08-28 11:54:28+00:00
pushed_at: 2025-02-26 12:12:21+00:00
description:
This IDA plugin extends the functionality of the assembly and hex view. With this plugin, you can conveniently decode/decrypt/alter data directly from the IDA Pro interface. 

content:
import idaapi
import inspect

from hexforge_modules import crypto_modules, encoding_modules, misc_modules, search_modules
from hexforge_modules.search_modules import SearchVirustotalBytes, SearchVirustotalString, SearchGitHub, SearchGoogle, SearchGrepApp

CRYPTO_MODULE_PATH = "HexForge/crypto/"
ENCODING_MODULE_PATH = "HexForge/encoding/"
MISC_MODULE_PATH = "HexForge/misc/"
SEARCH_MODULE_PATH = "HexForge/search/"


ida_version = idaapi.get_kernel_version()
is_ida_9_or_later = float(ida_version) >= 9.0

g_crypto_modules = []
g_encoding_modules = []
g_misc_modules = []
g_search_modules = []

class hexforge_plugin_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_KEEP
    comment = ""
    help = ""
    wanted_name = "HexForge"

    def init(self):
        global g_crypto_modules
        global g_encoding_modules
        global g_misc_modules
        global g_search_modules

        idaapi.msg("init() called!\n")
        g_crypto_modules = self._init_modules(crypto_modules)
        g_encoding_modules = self._init_modules(encoding_modules)
        g_misc_modules = self._init_modules(misc_modules)
        g_search_modules = self._init_modules(search_modules)
        self._init_actions()
        self._init_hooks()
        return idaapi.PLUGIN_KEEP

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

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

    # --------------------------------------------------------------------------
    # Initializations
    # --------------------------------------------------------------------------
    def _init_modules(self, modules) -> None:
        initialized_modules = []
        for _, cls in inspect.getmembers(modules, inspect.isclass):
            try:
                initialized_modules.append(cls())
            except Exception as e:
                idaapi.msg(f"Failed to initialize {cls.__name__}: {e}\n")
        return initialized_modules

    def _init_actions(self) -> None:
        for module in g_crypto_modules + g_encoding_modules + g_misc_modules + g_search_modules:
            module.init_action()

    def _del_action(self) -> None:
        for module in g_crypto_modules + g_encoding_modules + g_misc_modules + g_search_modules:
            module.del_action()

    # --------------------------------------------------------------------------
    # Initialize Hooks
    # --------------------------------------------------------------------------

    def _init_hooks(self) -> None:
        """
        Install plugin hooks into IDA.
        """
        self._hooks = Hooks()
        self._hooks.hook()


# 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_actions(widget, popup, idaapi.get_widget_type(widget))
        return 0


# Prefix Wrappers


def inject_actions(form, popup, form_type) -> int:
    """
    Inject actions to popup menu(s) based on context.
    """

    if is_ida_9_or_later:
        valid_form_types = (idaapi.BWN_DISASM, idaapi.BWN_HEXVIEW, idaapi.BWN_PSEUDOCODE)
    else:
        valid_form_types = (idaapi.BWN_DISASMS, idaapi.BWN_DUMP, idaapi.BWN_PSEUDOCODE)

    if form_type in valid_form_types:
        for module in g_crypto_modules:
            idaapi.attach_action_to_popup(
                form,
                popup,
                module.ACTION_NAME,
                CRYPTO_MODULE_PATH,
                idaapi.SETMENU_APP,
            )

        for module in g_misc_modules:
            idaapi.attach_action_to_popup(
                form,
                popup,
                module.ACTION_NAME,
                MISC_MODULE_PATH,
                idaapi.SETMENU_APP,
            )

        for module in g_encoding_modules:
            idaapi.attach_action_to_popup(
                form,
                popup,
                module.ACTION_NAME,
                ENCODING_MODULE_PATH,
                idaapi.SETMENU_APP,
            )

        for module in g_search_modules:
            # Exclude byte searches in decompiler
            if form_type == idaapi.BWN_PSEUDOCODE and isinstance(module, SearchVirustotalBytes):
                continue
            
            # Exclude string searches in disassembler
            if is_ida_9_or_later:
                if form_type == idaapi.BWN_DISASM and isinstance(module, (SearchGoogle, SearchGrepApp,SearchGitHub,SearchVirustotalString)):
                    continue
                idaapi.attach_action_to_popup(
                    form,
                    popup,
                    module.ACTION_NAME,
                    SEARCH_MODULE_PATH,
                    idaapi.SETMENU_APP,
                )

            else:
                if form_type == idaapi.BWN_DISASMS and isinstance(module, (SearchGoogle, SearchGrepApp,SearchGitHub,SearchVirustotalString)):
                    continue
                idaapi.attach_action_to_popup(
                    form,
                    popup,
                    module.ACTION_NAME,
                    SEARCH_MODULE_PATH,
                    idaapi.SETMENU_APP,
                )

    return 0


# Register IDA plugin
def PLUGIN_ENTRY() -> hexforge_plugin_t:
    return hexforge_plugin_t()


PLUGIN_ENTRY()


repository: zenyard/decompai-ida-public
file: src/decompai_ida/plugin.py
wanted_name: DecompAI
created_at: 2024-12-16 07:22:36+00:00
pushed_at: 2025-02-15 12:50:53+00:00
url:
https://github.com/zenyard/decompai-ida-public/blob/cffeda3f2e14a3aded632f8b4e9170b7f23751af/src/decompai_ida/plugin.py

content:
import typing as ty
from threading import Thread

import anyio
import ida_idaapi

from decompai_ida import configuration, ida_events, ida_tasks
from decompai_ida.main import main


class StopSignal:
    callback: ty.Optional[ida_tasks.AsyncCallback] = None


class DecompaiPlugin(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_FIX | ida_idaapi.PLUGIN_MULTI  # type: ignore
    wanted_name = "DecompAI"

    def init(self):
        # IDA disables Python default handling of SIGPIPE, making the process
        # crash on socket errors.
        try:
            import signal

            signal.signal(signal.SIGPIPE, signal.SIG_IGN)
        except Exception:
            # Signals not supported (e.g. Windows).
            pass

        return DecompaiPlugmod()


class DecompaiPlugmod(ida_idaapi.plugmod_t):
    def __init__(self):
        # Install hooks immediately, so we don't miss any event.
        self._event_collector = ida_events.EventCollector()
        self._ui_hooks = ida_events.UiEventHooks(self._event_collector)
        self._ui_hooks.hook()
        self._db_hooks = ida_events.DbEventHooks(self._event_collector)
        self._db_hooks.hook()

        self._stop_signal = StopSignal()
        self._thread = Thread(
            target=main_loop, args=(self._stop_signal, self._event_collector)
        )
        self._thread.start()

    def run(self, _arg):  # type: ignore
        configuration.show_configuration_dialog.sync()

    def __del__(self):
        self._ui_hooks.unhook()
        self._db_hooks.unhook()
        if self._stop_signal.callback is not None:
            self._stop_signal.callback()
        self._thread.join()


def main_loop(
    stop_signal: StopSignal, event_collector: ida_events.EventCollector
):
    async def main_with_stop_signal():
        # Note that events cannot be signalled from outside the loop. That's why
        # we create an `AsyncCallback` from within the loop and put it in a
        # shared place.
        stop_event = anyio.Event()
        stop_signal.callback = ida_tasks.AsyncCallback(stop_event.set)

        async with anyio.create_task_group() as tg:
            tg.start_soon(main, event_collector)
            await stop_event.wait()
            tg.cancel_scope.cancel()

    anyio.run(main_with_stop_signal)


def PLUGIN_ENTRY():
    return DecompaiPlugin()


repository: malpedia/malpedia-flossed
file: plugins/malpediaflossed/FlossedIda.py
description: FLARE floss applied to all unpacked+dumped samples in Malpedia, pre-processed for further use.
wanted_name: MalpediaFlossed
comment: This plugin matches strings in the given target binary with data from malpedia-flossed.json
created_at: 2024-01-11 12:25:21+00:00
pushed_at: 2025-03-14 13:04:32+00:00
url:
https://github.com/malpedia/malpedia-flossed/blob/b6f8dfeee473e1116afd22493eb03f08c79d54d5/plugins/malpediaflossed/FlossedIda.py

content:
from plugin.apis.IdaApi import IdaApi
from plugin.gui.PluginGui import PluginGui
import config

import idaapi
import ida_nalt
import ida_kernwin


class Plugin_action_handler_t(idaapi.action_handler_t):
    def __init__(self, dummy_arg):
        idaapi.action_handler_t.__init__(self)
        return

    def activate(self, ctx):
        # some action
        return 1

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


class Hooks(ida_kernwin.UI_Hooks):
    def __init__(self):
        ida_kernwin.UI_Hooks.__init__(self)

    def finish_populating_widget_popup(self, widget, popup):
        ida_kernwin.attach_action_to_popup(widget, popup, "TODO_IDA_ACTION_NAME", None)


class PluginWidget(ida_kernwin.PluginForm):

    def OnCreate(self, form):
        self.parent = self.FormToPyQtWidget(form)
        self.api_proxy = IdaApi()
        self.plugin_gui = PluginGui(self.api_proxy, config)
        self.parent.setLayout(self.plugin_gui.layout)

        # context menu extensions / actions need to be registered here
        idaapi.register_action(
            ida_kernwin.action_desc_t(
                "TODO_IDA_ACTION_NAME",
                "TODO_ACTION_DESCRIPTION",
                Plugin_action_handler_t("DummyArg"),
                "Ctrl+Shift+S",
                "TODO_ACTION_DESCRIPTION",
            )
        ),

    def OnClose(self, form):
        idaapi.unregister_action("TODO_IDA_ACTION_NAME")
        hooks.unhook()

    def Show(self, plugin_name="MalpediaFlossedPlugin"):
        if ida_nalt.retrieve_input_file_md5() is not None:
            return ida_kernwin.PluginForm.Show(self, plugin_name, options=(ida_kernwin.PluginForm.WCLS_CLOSE_LATER | ida_kernwin.PluginForm.WOPN_RESTORE | ida_kernwin.PluginForm.WCLS_SAVE))
        return None


class TemplatePlugin(idaapi.plugin_t):
    # https://www.hex-rays.com/products/ida/support/sdkdoc/group___p_l_u_g_i_n__.html
    flags = idaapi.PLUGIN_UNL
    comment = "This plugin matches strings in the given target binary with data from malpedia-flossed.json"
    help = "help"
    wanted_name = "MalpediaFlossed"
    wanted_hotkey = "Ctrl+Shift+Y"

    def init(self):
        global hooks
        hooks = Hooks()
        hooks.hook()
        return idaapi.PLUGIN_OK

    def run(self, arg):
        plg = PluginWidget()
        plg.Show("MalpediaFlossed")

        try:
            widget_a = ida_kernwin.find_widget("IDA View-A")
            widget_template_plugin = ida_kernwin.find_widget("MalpediaFlossed")
            if widget_template_plugin and widget_a:
                ida_kernwin.set_dock_pos("MalpediaFlossed", "IDA View-A", ida_kernwin.DP_RIGHT)
        except:
            print("find_widget option is available version 7.0 or later")

    def term(self):
        pass


def PLUGIN_ENTRY():
    return TemplatePlugin()


def main():
    global MALPEDIA_STRINGS_WIDGET
    try:
        # TODO figure out why we can't reopen our form when it was closed once
        MALPEDIA_STRINGS_WIDGET = PluginWidget()
    except Exception:
        import traceback
        traceback.print_exc()
        MALPEDIA_STRINGS_WIDGET = PluginWidget()

    MALPEDIA_STRINGS_WIDGET.Show()


if __name__ == "__main__":
    main()


repository: binarly-io/idapcode
file: idapcode.py
description: IDA plugin displaying the P-Code for the current function
url: https://github.com/binarly-io/idapcode/blob/716e32750d57c99e62e1803af7e4e5218d08e8ad/idapcode.py
created_at: 2021-07-18 22:39:10+00:00
pushed_at: 2023-10-27 20:19:16+00:00
content:
import traceback
from typing import Dict, List, Optional

import ida_bytes
import ida_funcs
import ida_kernwin
import ida_lines
import ida_name
import idaapi
from pypcode import Arch, Context, PcodePrettyPrinter

NAME = "IDA P-Code"
AUTHOR = "https://github.com/yeggor"

VERSION = "0.0.1"
HELP_MSG = "Get the P-Code for the current function"
COMMENT_MSG = HELP_MSG
WANTED_KEY = "Ctrl+Alt+S"

DEBUG = False


# -----------------------------------------------------------------------
class FuncPcode:
    """Helper class for getting p-code for a function"""

    def __init__(self, addr: int) -> None:
        self._addr: int = addr
        self._func_pcode: Optional[List[str]] = None
        self._func_name: Optional[str] = None
        self._inf = idaapi.get_inf_structure()

        # adopted from
        # https://github.com/cseagle/blc/blob/b1447562a3598fd411224dfc24b970cf53ca7c94/plugin.cc#L516
        self._proc_map: Dict[Any] = dict()
        self._proc_map[idaapi.PLFM_6502] = "6502"
        self._proc_map[idaapi.PLFM_68K] = "68000"
        self._proc_map[idaapi.PLFM_6800] = "6805"
        self._proc_map[idaapi.PLFM_8051] = "8051"
        self._proc_map[idaapi.PLFM_ARM] = "ARM"
        self._proc_map[idaapi.PLFM_AVR] = "avr8"
        self._proc_map[idaapi.PLFM_CR16] = "CR16"
        self._proc_map[idaapi.PLFM_DALVIK] = "Dalvik"
        self._proc_map[idaapi.PLFM_JAVA] = "JVM"
        self._proc_map[idaapi.PLFM_MIPS] = "MIPS"
        self._proc_map[idaapi.PLFM_HPPA] = "pa-risc"
        self._proc_map[idaapi.PLFM_PIC] = "PIC"
        self._proc_map[idaapi.PLFM_PPC] = "PowerPC"
        self._proc_map[idaapi.PLFM_SPARC] = "sparc"
        self._proc_map[idaapi.PLFM_MSP430] = "TI_MSP430"
        self._proc_map[idaapi.PLFM_TRICORE] = "tricore"
        self._proc_map[idaapi.PLFM_386] = "x86"
        self._proc_map[idaapi.PLFM_Z80] = "Z80"

    def _inf_is_64bit(self) -> bool:
        return self._inf.is_64bit()

    def _inf_is_32bit(self) -> bool:
        return self._inf.is_32bit()

    def _get_app_bittness(self) -> int:
        if self._inf_is_64bit():
            return 64
        if self._inf_is_32bit():
            return 32
        return 16

    def _inf_is_be(self) -> bool:
        return self._inf.is_be()

    def _get_proc_id(self) -> int:
        return idaapi.ph_get_id()

    def _get_proc(self) -> Optional[str]:
        proc_id = self._get_proc_id()
        if proc_id not in self._proc_map:
            return None
        return self._proc_map[proc_id]

    def _get_endian(self) -> str:
        if self._inf_is_be():
            return "BE"
        return "LE"

    def _get_sleigh_id(self) -> Optional[str]:
        """Get sleigh language id string"""

        proc = self._get_proc()
        if proc is None:
            return None
        endian = self._get_endian()

        # adopted from
        # https://github.com/cseagle/blc/blob/b1447562a3598fd411224dfc24b970cf53ca7c94/plugin.cc#L637
        sleigh = f"{proc}:{endian}"
        proc_id = self._get_proc_id()
        if proc_id == idaapi.PLFM_6502:
            return f"{sleigh}:16:default"
        elif proc_id == idaapi.PLFM_68K:
            return f"{sleigh}:32:default"
        elif proc_id == idaapi.PLFM_6800:
            return f"{sleigh}:8:default"
        elif proc_id == idaapi.PLFM_8051:
            return f"{sleigh}:16:default"
        elif proc_id == idaapi.PLFM_ARM:
            if self._inf_is_64bit():
                return f"AARCH64:{endian}:64:v8A"
            return f"{sleigh}:32:v7"
        elif proc_id == idaapi.PLFM_AVR:
            if self._get_app_bittness() == 32:
                return f"avr32:{endian}:32:default"
            if self._get_app_bittness() == 16:
                return f"{sleigh}:16:default"
            return sleigh
        elif proc_id == idaapi.PLFM_CR16:
            return f"{sleigh}:16:default"
        elif proc_id == idaapi.PLFM_DALVIK:
            return f"{sleigh}:32:default"
        elif proc_id == idaapi.PLFM_JAVA:
            return f"{sleigh}:32:default"
        elif proc_id == idaapi.PLFM_MIPS:
            abi = idaapi.get_abi_name()
            if abi and ("n32" in abi):
                return f"{sleigh}:64:64-32addr"
            if self._inf_is_64bit():
                return f"{sleigh}:64:default"
            if self._inf_is_32bit():
                return f"{sleigh}:32:default"
            return sleigh
        elif proc_id == idaapi.PLFM_HPPA:
            return f"{sleigh}:32:default"
        elif proc_id == idaapi.PLFM_PIC:
            return sleigh
        elif proc_id == idaapi.PLFM_PPC:
            abi = idaapi.get_abi_name()
            if abi and ("xbox" in abi):
                return f"{sleigh}:64:A2ALT-32addr"
            if self._inf_is_64bit():
                return f"{sleigh}:64:default"
            if self._inf_is_32bit():
                return f"{sleigh}:32:default"
            return sleigh
        elif proc_id == idaapi.PLFM_SPARC:
            if self._inf_is_64bit():
                return f"{sleigh}:64:default"
            if self._inf_is_32bit():
                return f"{sleigh}:32:default"
            return sleigh
        elif proc_id == idaapi.PLFM_MSP430:
            return f"{sleigh}:16:default"
        elif proc_id == idaapi.PLFM_TRICORE:
            return f"{sleigh}:32:default"
        elif proc_id == idaapi.PLFM_386:
            bittness = self._get_app_bittness()
            sleigh += f":{str(bittness)}"
            if bittness == 16:
                return f"{sleigh}:Real Mode"
            return f"{sleigh}:default"
        elif proc_id == idaapi.PLFM_Z80:
            return sleigh
        return None

    def _get_func_name(self) -> str:
        """Get function name"""
        f = ida_funcs.get_func(self._addr)
        if f is None:
            return "unknown"
        return ida_name.get_name(f.start_ea)

    def _get_func_bytes(self) -> bytes:
        """Get function bytes"""
        f = ida_funcs.get_func(self._addr)
        if f is None:
            return bytes()
        return ida_bytes.get_bytes(f.start_ea, f.end_ea - f.start_ea)

    def _get_pcode(self) -> list:
        """Get P-Code lines"""
        code = self._get_func_bytes()
        if code is None:
            return list()
        # get sleigh id
        sleigh_id = self._get_sleigh_id()
        if sleigh_id is None:
            return list()
        if DEBUG:
            print(f"[ {NAME} ] using sleigh id: {sleigh_id}")

        # translate to P-Code
        langs = {l.id: l for arch in Arch.enumerate() for l in arch.languages}
        if sleigh_id not in langs:
            return list()
        ctx = Context(langs[sleigh_id])
        res = ctx.translate(code=code, base=0, max_inst=0, bb_terminating=False)

        pcode_lines = list()
        f = ida_funcs.get_func(self._addr)
        for insn in res.instructions:
            # append asm text
            addr = insn.address.offset + f.start_ea
            asm_prefix = ida_lines.COLSTR(f"{addr:018X}", ida_lines.SCOLOR_PREFIX)
            asm_insn = ida_lines.COLSTR(
                f"{insn.asm_mnem.lower()} {insn.asm_body.lower()}",
                ida_lines.SCOLOR_INSN,
            )
            pcode_lines.append(f"{asm_prefix}  {asm_insn}")
            # append P-Code text
            for op in insn.ops:
                pcode_lines.append(f"  {PcodePrettyPrinter.fmt_op(op)}")

            pcode_lines.append("\n")

        return pcode_lines

    @property
    def func_name(self) -> str:
        if self._func_name is None:
            self._func_name = self._get_func_name()
        return self._func_name

    @property
    def pcode(self) -> List[str]:
        if self._func_pcode is None:
            self._func_pcode = self._get_pcode()
        return self._func_pcode


# -----------------------------------------------------------------------
class pcodecv_t(ida_kernwin.simplecustviewer_t):
    def Create(self, sn=None, use_colors=True):
        fpcode = FuncPcode(idaapi.get_screen_ea())

        title = f"P-Code for {fpcode.func_name}"
        if sn:
            title += f" {sn:#d}"
        self._use_colors = use_colors

        pcode_lines = fpcode.pcode
        if not pcode_lines:
            print(f"Can't get P-Code for function {fpcode.func_name}")
            return False

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

        for pcode_line in pcode_lines:
            self.AddLine(pcode_line)

        return True

    def OnClick(self, shift):
        if DEBUG:
            print(f"OnClick, shift={shift:#d}")
        return True

    def OnDblClick(self, shift):
        if DEBUG:
            print(f"OnDblClick, shift={shift:#d}")
        return True

    def OnCursorPosChanged(self):
        if DEBUG:
            print("OnCurposChanged")

    def OnClose(self):
        if DEBUG:
            print(f"{self.title} closed")

    def OnKeydown(self, vkey, shift):
        if DEBUG:
            print(f"OnKeydown, vk={vkey:#d} shift={shift:#d}")
        return True

    def OnHint(self, lineno):
        if DEBUG:
            return (1, f"OnHint, line={lineno:#d}")
        return (0, str())

    def Show(self, *args):
        return ida_kernwin.simplecustviewer_t.Show(self, *args)


# -----------------------------------------------------------------------
# Show P-Code View
def show_pcodecv():
    pcodecv = pcodecv_t()
    if not pcodecv.Create(use_colors=True):
        print(f"[ {NAME} ] failed to create view")
        return None
    pcodecv.Show()
    return pcodecv


# -----------------------------------------------------------------------
class IdaPcode(idaapi.plugin_t):
    """IDA P-Code plugin class"""

    flags = idaapi.PLUGIN_MOD | idaapi.PLUGIN_PROC | idaapi.PLUGIN_FIX
    comment = COMMENT_MSG
    help = HELP_MSG
    wanted_name = NAME
    wanted_hotkey = WANTED_KEY

    def init(self):
        print(f"\n{NAME} ({VERSION})")
        print(f"{NAME} shortcut key is {WANTED_KEY}\n")
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        show_pcodecv()

    def term(self):
        if DEBUG:
            print(f"[ {NAME} ] terminated")


# -----------------------------------------------------------------------
def PLUGIN_ENTRY():
    try:
        return IdaPcode()
    except Exception as e:
        print(f"[ {NAME} ] {str(e)}\n{traceback.format_exc()}")


if __name__ == "__main__":
    try:
        show_pcodecv()
    except Exception as e:
        print(f"[ {NAME} ] {str(e)}\n{traceback.format_exc()}")


repository: ax330d/functions-plus
file: functions_plus.py
description: IDA Pro plugin to show functions in a tree view
wanted_name: Functions+
comment: Functions+
created_at: 2017-07-19 18:27:03+00:00
pushed_at: 2021-04-17 20:03:10+00:00
url:
https://github.com/ax330d/functions-plus/blob/a0eac3333fbd1a8df86fcdbb0eb0fbbc7a92990f/functions_plus.py

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

'''
Functions+ IDA Pro plugin -- alternative version of functions window.

Splits functions names and groups by namespaces.
'''

import re
from collections import OrderedDict

# import cProfile
# import pstats
# import StringIO

import idaapi
import idc
import idautils

from idaapi import PluginForm
from PyQt5 import QtWidgets, QtGui

__author__ = 'xxxzsx, Arthur Gerkis'
__version__ = '1.0.1'


class FunctionState(object):
    '''
    Holds the state of the current function.
    '''

    def __init__(self):
        self._args = ''
        self._flags = 0
        self._addr = 0

    @property
    def args(self):
        '''
        Returns args property.
        '''
        return self._args

    @args.setter
    def args(self, value):
        '''
        Sets args property.
        '''
        self._args = value

    @property
    def flags(self):
        '''
        Returns flags property.
        '''
        return self._flags

    @flags.setter
    def flags(self, value):
        '''
        Sets flags property.
        '''
        self._flags = value

    @property
    def addr(self):
        '''
        Returns addr property.
        '''
        return self._addr

    @addr.setter
    def addr(self, value):
        '''
        Sets addr property.
        '''
        self._addr = value


class FunctionData(object):
    '''
    Holds data of the function.
    '''

    def __init__(self, state):
        self._args = state.args
        self._flags = state.flags
        self._addr = state.addr

    @property
    def args(self):
        '''
        Returns args property.
        '''
        return self._args

    @property
    def flags(self):
        '''
        Returns flags property.
        '''
        return self._flags

    @property
    def addr(self):
        '''
        Returns ea property.
        '''
        return self._addr


class Cols(object):
    '''
    Class which is responsible for handling columns.
    '''
    def __init__(self, show_extra_fields):
        self.addr = None
        self.flags = None
        self.show_extra_fields = show_extra_fields
        self.names = [
            'Name', 'Address', 'Segment', 'Length', 'Locals', 'Arguments'
        ]

        self.handlers = {
            0: lambda: None,
            1: lambda: self.fmt(self.addr),
            2: lambda: '{}'.format(idc.get_segm_name(self.addr)),
            3: lambda: self.fmt(idc.get_func_attr(self.addr, idc.FUNCATTR_END) - self.addr),
            4: lambda: self.fmt(idc.get_func_attr(self.addr, idc.FUNCATTR_FRSIZE)),
            5: lambda: self.fmt(idc.get_func_attr(self.addr, idc.FUNCATTR_ARGSIZE))
        }

        if self.show_extra_fields:
            self.names.extend(['R', 'F', 'L', 'S', 'B', 'T', '='])
            # TODO: add Lumina column info
            self.handlers.update({
                6:  lambda: self.is_true(not self.flags & idc.FUNC_NORET, 'R'),
                7:  lambda: self.is_true(self.flags & idc.FUNC_FAR, 'F'),
                8:  lambda: self.is_true(self.flags & idc.FUNC_LIB, 'L'),
                9:  lambda: self.is_true(self.flags & idc.FUNC_STATIC, 'S'),
                10: lambda: self.is_true(self.flags & idc.FUNC_FRAME, 'B'),
                11: lambda: self.is_true(idc.get_type(self.addr), 'T'),
                12: lambda: self.is_true(self.flags & idc.FUNC_BOTTOMBP, '=')
            })

    def set_data(self, addr, flags):
        '''
        Sets data actual for the current function.
        '''
        self.addr = addr
        self.flags = flags

    def item(self, index):
        '''
        Gets the data according to requested col index.
        '''

        return self.handlers[index]()

    def is_true(self, flag, char):
        '''
        Wrapper to conform IDA default UI view.
        '''
        if flag:
            return char
        return '.'

    def fmt(self, value):
        '''
        Wrapper to conform IDA default UI view.
        '''
        return '{:08X}'.format(value)


class FunctionsTree(object):
    '''
    Builds tree of functions with all relevant information.
    '''
    def __init__(self):
        self.chunks_regexp = re.compile(r'(.*?)(?:|\((.*?)\))$')
        self.simple_regexp = re.compile(r'^[a-zA-Z0-9_]*$')

    def get(self):
        '''
        Returns functions tree.
        '''

        functions_list = self.get_list_of_functions()
        functions_tree = self.build_functions_tree(functions_list)

        return functions_tree

    def get_list_of_functions(self):
        '''
        Gets all functions list.
        '''

        functions_list = {}
        seg_ea = idc.get_segm_by_sel(idc.SEG_NORM)

        for func_ea in idautils.Functions(idc.get_segm_start(seg_ea),
                                          idc.get_segm_end(seg_ea)):
            function_name = idc.get_func_name(func_ea)
            functions_list[function_name] = func_ea

        return functions_list

    def build_functions_tree(self, functions_list):
        '''
        Builds tree of functions.
        '''

        func_state = FunctionState()
        functions_tree = OrderedDict()

        for function_name in sorted(functions_list):
            func_state.args = ''
            func_state.addr = functions_list[function_name]
            func_state.flags = \
                idc.get_func_attr(func_state.addr, idc.FUNCATTR_FLAGS)
            demangled_name = self.maybe_demangle(function_name)
            chunks = self.get_chunks(demangled_name, func_state)
            self.maybe_push(chunks, functions_tree, func_state)

        return functions_tree

    def maybe_push(self, chunks, functions_tree, func_state):
        '''
        Adds new function name or properties to the tree.
        '''

        # FIXME: handle duplicate entries properly
        if isinstance(functions_tree, FunctionData):
            return

        name = chunks.pop(0)
        if not len(name):
            return

        # If this is the last (or one) chunk
        if not len(chunks):
            functions_tree[name + func_state.args] = FunctionData(func_state)
            return

        # If this is a new namespace, create a tree
        if name not in functions_tree:
            functions_tree[name] = OrderedDict()

        return self.maybe_push(chunks, functions_tree[name], func_state)

    def get_chunks(self, func_string, func_state):
        '''
        Splits function name by namespaces.
        '''

        new_chunks = []
        matches = re.match(self.chunks_regexp, func_string)
        if not matches:
            return []

        args = ''
        if matches.group(2):
            args = '({})'.format(matches.group(2))
        func_state.args = args

        chunks = list(matches.group(1))
        if chunks[0] == '`':
            return [matches.group(1)]

        open_left_tpl = 0
        tmp_chunk = ''
        for chunk in chunks:
            if chunk == ':' and open_left_tpl == 0:
                if len(tmp_chunk):
                    new_chunks.append(tmp_chunk)
                tmp_chunk = ''
                continue
            if chunk == '<':
                open_left_tpl += 1
            if chunk == '>':
                open_left_tpl -= 1
            tmp_chunk += chunk
        new_chunks.append(tmp_chunk)
        return new_chunks

    def maybe_demangle(self, function_name):
        '''
        Demangles name of required.
        '''

        if function_name.find('@') != -1:
            function_name = self.demangle(function_name)
        return function_name

    @classmethod
    def demangle(cls, name):
        '''
        Demangles name.
        '''

        mask = idc.get_inf_attr(idc.INF_SHORT_DN)
        demangled = idc.demangle_name(name, mask)
        if demangled is None:
            return name
        return demangled


class FunctionsPlus(PluginForm):
    '''Functions+ plugin.'''

    def __init__(self):
        super(FunctionsPlus, self).__init__()
        if idc.get_inf_attr(idc.INF_PROCNAME).lower() != 'metapc':
            print('Functions+ warning: not tested in this configuration')
        self.tree = None
        self.icon = 135
        # Enable this if you want to see extra information about function
        self.show_extra_fields = False
        self.cols = Cols(self.show_extra_fields)

    def _populate_tree(self):
        '''
        Populates functions tree.
        '''
        self.tree.clear()
        self._build_tree(FunctionsTree().get(), self.tree)
        return

    def _build_tree(self, function_tree, root):
        '''
        Builds Qt Widget tree.
        '''

        if not function_tree:
            return

        if isinstance(function_tree, FunctionData):
            self._handle_function_data_instance(function_tree, root)
            return

        for name, tree in sorted(function_tree.items()):
            func_item = QtWidgets.QTreeWidgetItem(root)
            if not isinstance(tree, FunctionData):
                name = self._handle_class_name(tree, name, func_item)
            func_item.setText(0, name)
            self._build_tree(tree, func_item)

    def _handle_class_name(self, tree, name, func_item):
        '''
        Handles class name.
        '''

        tree_keys_len = len(list(tree.keys()))
        name = '{} ({} {})'.\
            format(name, tree_keys_len, self._get_word(tree_keys_len))
        font = QtGui.QFont()
        font.setBold(True)
        func_item.setFont(0, font)
        return name

    def _handle_function_data_instance(self, function_tree, root):
        '''
        Handles FunctionData instance.
        '''

        flags = int(function_tree.flags)
        addr = function_tree.addr

        self.cols.set_data(addr, flags)

        for index in range(0, len(self.cols.names)):
            if index > 0:
                root.setText(index, self.cols.item(index))
            if flags & idc.FUNC_THUNK:
                root.setBackground(index, QtGui.QColor('#E8DAEF'))
            if flags & idc.FUNC_LIB:
                root.setBackground(index, QtGui.QColor('#D1F2EB'))

    def _get_word(self, len):
        '''
        Gets proper word for number.
        '''
        word = 'items'
        if len % 10 == 1 and len % 100 != 11:
            word = 'item'
        return word

    def _dblclick(self, item):
        '''
        Handles double click event.
        '''
        try:
            idaapi.jumpto(int(item.text(1), 16))
        except:
            pass

    def OnCreate(self, form):
        '''
        Called when the plugin form is created.
        '''

        # pr = cProfile.Profile()
        # pr.enable()

        parent = self.FormToPyQtWidget(form)

        self.tree = QtWidgets.QTreeWidget()
        self.tree.setColumnCount(len(self.cols.names))
        self.tree.setHeaderLabels(self.cols.names)
        self.tree.itemDoubleClicked.connect(self._dblclick)
        # self.tree.resizeColumnToContents(True)

        layout = QtWidgets.QVBoxLayout()
        layout.addWidget(self.tree)
        layout.setSpacing(0)
        layout.setContentsMargins(0, 0, 0, 0)

        self._populate_tree()

        self.tree.setColumnWidth(0, 512)
        for index in range(6, len(self.cols.names)):
            self.tree.setColumnWidth(index, 32)
        self.tree.setAlternatingRowColors(True)

        parent.setLayout(layout)

        # pr.disable()
        # s = StringIO.StringIO()
        # ps = pstats.Stats(pr, stream=s).sort_stats('cumulative')
        # ps.print_stats()
        # print(s.getvalue())

    def OnClose(self, form):
        '''
        Called when the plugin form is closed.
        '''
        del self

    def Show(self):
        '''
        Creates the form is not created or focuses it if it was.
        '''
        return PluginForm.Show(self, 'Functions+')


class FunctionsPlusPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_KEEP
    comment = "Functions+"

    help = ""
    wanted_name = "Functions+"
    wanted_hotkey = ""

    # @staticmethod
    # def init():
    #     return idaapi.PLUGIN_KEEP

    @staticmethod
    def init():
        funp = FunctionsPlus()
        funp.Show()
        return idaapi.PLUGIN_KEEP

    @staticmethod
    def run(arg=0):
        funp = FunctionsPlus()
        funp.Show()

    @staticmethod
    def term():
        pass

def PLUGIN_ENTRY():
    return FunctionsPlusPlugin()


repository: riskeco/Match4IDA
file: match4ida.py
url: https://github.com/riskeco/Match4IDA/blob/0e26754ef69a26493936e985fdeb03eb0b14ff90/match4ida.py
wanted_name: Match4IDA
comment: Navigate to rule matched locations inside IDA.
created_at: 2023-09-07 09:08:27+00:00
pushed_at: 2023-09-14 18:02:50+00:00
content:
import json
import sys
from enum import IntEnum
from pathlib import Path

try:
    import pychrysalide
    from pychrysalide.analysis.contents import FileContent
    from pychrysalide.analysis.scan import ContentScanner
    from pychrysalide.analysis.scan import ScanOptions
    from pychrysalide.analysis.scan.patterns.backends import AcismBackend
except:
    pass

try:
    import yara
except:
    pass

try:
    import idaapi
    from PyQt5 import QtCore, QtWidgets
except:
    pass



class ScanHandler():
    """Generic handler for scan rules."""

    class ScannerType(IntEnum):
        YARA  = 0
        ROST  = 1
        GUESS = 2

    def __init__(self, rule_filename, rule_type, binary_filename):
        """Create a generic handler for scanning."""

        self._rule_filename = rule_filename
        self._rule_type = rule_type
        self._binary_filename = binary_filename


    def _guess_suitable_scanner(self):
        """Try to guess if a rule is for YARA or ROST."""

        with open(self._rule_filename, 'rb') as fd:
            content = fd.read()

        has_strings = b'strings:' in content
        has_bytes = b'bytes:' in content

        if has_strings and not(has_bytes):
            rtype = ScanHandler.ScannerType.YARA

        elif not(has_strings) and has_bytes:
            rtype = ScanHandler.ScannerType.ROST

        else:
            rtype = ScanHandler.ScannerType.GUESS

        return rtype


    def run(self):
        """Run a scan."""

        rtype = self._rule_type

        if rtype == ScanHandler.ScannerType.GUESS:
            rtype = self._guess_suitable_scanner()

        found = []

        if rtype == ScanHandler.ScannerType.YARA and 'yara' in sys.modules.keys():

            idaapi.msg('Running YARA against %s...' % self._binary_filename)

            rules = yara.compile(self._rule_filename)

            matches = rules.match(self._binary_filename)

            for m in matches:
                for s in m.strings:
                    for i in s.instances:

                        extra = {
                            'identifier': s.identifier,
                            'bytes': str(i),
                            'start': i.offset,
                            'length': i.matched_length
                        }

                        found.append(extra)

        elif rtype == ScanHandler.ScannerType.ROST and 'pychrysalide' in sys.modules.keys():

            idaapi.msg('Running ROST against %s...' % self._binary_filename)

            scanner = ContentScanner(filename=self._rule_filename)

            content = FileContent(self._binary_filename)

            options = ScanOptions()
            options.backend_for_data = AcismBackend

            ctx = scanner.analyze(options, content)

            data = scanner.convert_to_json(ctx)

            data = json.loads(data)

            for rule in data:
                for pat in rule['bytes_patterns']:
                    for m in pat['matches']:

                        extra = {
                            'identifier': pat['name'],
                            'bytes': m['content_str'].replace(r'\\', '\\'),
                            'start': m['offset'],
                            'length': m['length']
                        }

                        found.append(extra)

        return found



class MatchPanel(idaapi.PluginForm):
    """Panel for the IDA GUI."""

    def OnCreate(self, form):
        """Create a panel for the plugin activity."""

        parent = self.FormToPyQtWidget(form)

        # Create layout

        layout = QtWidgets.QGridLayout()

        parent.setLayout(layout)

        # Connection properties

        self._rule_filename = QtWidgets.QLineEdit()
        self._rule_filename.setText('/tmp/Match4IDA/Sample/APT_MAL_UNC4841_SEASPY_Jun23_1.yar')
        layout.addWidget(self._rule_filename, 0, 0)

        self._browse_button = QtWidgets.QPushButton('Browse')
        self._browse_button.clicked.connect(self._browse)
        layout.addWidget(self._browse_button, 0, 1)

        self._scanner_type = QtWidgets.QComboBox()
        self._scanner_type.addItems(['Yara', 'ROST', 'auto'])
        self._scanner_type.setCurrentIndex(2)
        layout.addWidget(self._scanner_type, 0, 2)

        self._scan_button = QtWidgets.QPushButton('Scan')
        self._scan_button.clicked.connect(self._run_scan)
        layout.addWidget(self._scan_button, 0, 3)

        # Match display

        self._rows = QtWidgets.QTableWidget()

        column_names = [ 'Identifier', 'Found bytes', 'Offset', 'Start location', 'Match size' ]

        self._rows.setColumnCount(len(column_names))
        self._rows.setHorizontalHeaderLabels(column_names)

        self._rows.setRowCount(0)
        self._rows.doubleClicked.connect(self._jump_to_match_location)

        header = self._rows.horizontalHeader()
        header.setSectionResizeMode(0, QtWidgets.QHeaderView.ResizeToContents)
        header.setSectionResizeMode(1, QtWidgets.QHeaderView.Stretch)
        header.setSectionResizeMode(2, QtWidgets.QHeaderView.ResizeToContents)
        header.setSectionResizeMode(3, QtWidgets.QHeaderView.ResizeToContents)
        header.setSectionResizeMode(4, QtWidgets.QHeaderView.ResizeToContents)

        # fromRow - fromColumn - rowSpan - columnSpan
        layout.addWidget(self._rows, 1, 0, 1, 4)


    def _browse(self):
        """Select a rule as analysis source."""

        dlg = QtWidgets.QFileDialog()
        dlg.setFileMode(QtWidgets.QFileDialog.ExistingFile)
        dlg.setNameFilter('YARA rules (*.yar);;ROST rules (*.rost)')

        if dlg.exec_():

            filename = str(Path(dlg.selectedFiles()[0]))

            self._rule_filename.setText(filename)


    def _run_scan(self):
        """Run a scan and display the results."""

        rule_filename = self._rule_filename.text()
        rule_type = self._scanner_type.currentIndex()
        binary_filename = idaapi.get_input_file_path()

        if len(rule_filename) == 0:
            return

        scanner = ScanHandler(rule_filename, rule_type, binary_filename)

        matches = scanner.run()

        self._rows.setRowCount(0)

        for m in matches:

            index = self._rows.rowCount()
            self._rows.setRowCount(index + 1)

            item = QtWidgets.QTableWidgetItem(m['identifier'])
            item.setFlags(item.flags() & ~QtCore.Qt.ItemIsEditable)

            self._rows.setItem(index, 0, item)

            item = QtWidgets.QTableWidgetItem(m['bytes'])
            item.setFlags(item.flags() & ~QtCore.Qt.ItemIsEditable)

            self._rows.setItem(index, 1, item)

            item = QtWidgets.QTableWidgetItem('0x%x' % m['start'])
            item.setFlags(item.flags() & ~QtCore.Qt.ItemIsEditable)

            self._rows.setItem(index, 2, item)

            ea = idaapi.get_fileregion_ea(m['start'])

            item = QtWidgets.QTableWidgetItem('0x%x' % ea)
            item.setFlags(item.flags() & ~QtCore.Qt.ItemIsEditable)

            self._rows.setItem(index, 3, item)

            item = QtWidgets.QTableWidgetItem('0x%x' % m['length'])
            item.setFlags(item.flags() & ~QtCore.Qt.ItemIsEditable)

            self._rows.setItem(index, 4, item)

            self._rows.update()

        idaapi.msg('Found %u match(es)\n' % self._rows.rowCount())


    def _jump_to_match_location(self, item):

        addr = int(self._rows.item(item.row(), 3).text(), 16)

        idaapi.jumpto(addr)



class Match4IDA(idaapi.plugin_t):

    flags = idaapi.PLUGIN_KEEP
    comment = 'Navigate to rule matched locations inside IDA.'

    wanted_name = 'Match4IDA'
    wanted_hotkey = ''
    help = 'Scan the current open binary against byte patterns.'


    def init(self):
        """Init the IDA plugin."""

        idaapi.msg('Starting %s\n' % self.wanted_name)

        self._form = None

        return idaapi.PLUGIN_KEEP


    def run(self, arg):
        """Run the IDA plugin."""

        idaapi.msg('Running %s\n' % self.wanted_name)

        self._form = MatchPanel()
        self._form.Show('Scan matches')


    def term(self):
        """Terminate the IDA plugin."""

        idaapi.msg('Terminating %s\n' % self.wanted_name)



def PLUGIN_ENTRY():
    return Match4IDA()



if __name__ == '__main__':
    """Script entrypoint."""

    has_yara = 'yara' in sys.modules.keys()

    if not(has_yara or False):
        print('At least one scanner is requiered')

    inside_ida = 'ida_idaapi' in sys.modules.keys()

    if not(inside_ida):

        if len(sys.argv) != 3:
            print('Usage: %s <rule> <binary>' % sys.argv[0])
            sys.exit(2)

        rule_filename = sys.argv[1]
        binary_filename = sys.argv[2]

        scanner = ScanHandler(rule_filename, ScanHandler.ScannerType.GUESS, binary_filename)

        matches = scanner.run()

        for m in matches:
            print('0x%x:%s:' % (m['start'], m['identifier']), m['bytes'])


repository: LucienMP/idapro_m6502
file: plugins/m6502plugin.py
wanted_name: m6502_stepover
created_at: 2020-11-10 04:03:52+00:00
pushed_at: 2020-11-10 04:22:35+00:00
description:
Extends existing support in IDA Pro for the m6502 by adding gdb XML support, and step-over and type information support

url:
https://github.com/LucienMP/idapro_m6502/blob/a39246343f10eeff4e82cf5890f5c5f67e15905c/plugins/m6502plugin.py

content:
import idaapi
import ida_pro
import idautils

#from pprint import pprint

# Used by the loop for function purge
from idaapi import *

# Used for printing various data
import logging

"""
    This is a sample plugin for extending gdb support for step-over for the m6502,
    and to enable type support so you can press "y" on functions.

    When you install a hook in the HT_IDP category, it will be called
    before the processor module's notify() function, so if you return
    non-zero, your results will be used.

    1) Add GDB style XML files for the CPU register descriptions to $IDA_HOME/cfg/
    2) Add to the section in $IDA_HOME/cfg/dbg_gdb.cfg that integrates that XML for GDB
        CONFIGURATIONS = { "m6502" : ...
        IDA_FEATURES = { "m6502" : ... 
        NOT NEEDED> ARCH_MAP = { "m6502" : ... 
        NOT NEEDED> FEATURE_MAP = ...
    3) Add function to handle step-over understanding around branches
        The ev_calc_step_over interprets the current instruction.
        If it is a JSR routine call then it will advance the IP to the next instruction.
        In all other cases it will advance to the next insturction by passing back BADADDR
    4) Add function to handle type info (funciton name, and paramater propagation)
        To enable type support you must set PR_TYPEINFO in ph.flag and 
        implement most of the related callbacks (see "START OF TYPEINFO 
        CALLBACKS" in idp.hpp from the IDA SDK).

    Many thanks to Ilfak Guilfanov @ Hey-Rays

    Versioning:
        2019-02-28 v1.0 (c) Lucien Murray-Pitts - Initial version to support step-over
""" 

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

    #--------------------------------------------------------------------------
    #
    # Handling of GDB related events
	
    # Calculate the address of the instruction which will be
    # executed after "step over". The kernel will put a breakpoint there.
    #
    def ev_calc_step_over(self, target, ip):
        # Supports the following instructions
        # 4E80..BF        jsr     (a0) +2
        # idaapi.BADADDR  else    BADADDR
        #
        logging.debug("JSR, checking instruction")
        ta=ida_pro.ea_pointer_frompointer(target)

        insn = idautils.DecodeInstruction(ip)
        if insn and insn.itype in [idaapi.M65_jsr]:
            over_ea = insn.ea + insn.size
            logging.debug(" - JSR STEP> +%s %s" % (insn.size, hex(over_ea)) )
            ta.assign(over_ea)
        else:
            logging.debug(" - NO STEP> "+hex(ip))
            ta.assign(idaapi.BADADDR)

        return 1

    #--------------------------------------------------------------------------
    #
    # Handling of TYPE related events
	
    # Setup default type libraries.
    # (called after loading a new file into the database). The processor module may load tils, 
    # setup memory model and perform other actions required to set up the type system. This is an optional callback.
    #
    def ev_setup_til(self, abi_names, abi_opts, comp):
        logging.debug("##########ev_setup_til")
        return
		
    # Get all possible ABI names and optional extensions for given compiler
    # abiname/option is a string entirely consisting of letters, digits and underscore
    # 0=not implemented, 1=ok
    #
    def ev_get_abi_info(self, abi_names, abi_opts, comp):
        logging.debug("##########ev_get_abi_info")
        logging.debug(abi_names)
        logging.debug(abi_opts)
        logging.debug(comp)
        return 0

    # Get maximal size of a pointer in bytes.
    #
    #
    def ev_max_ptr_size(self, target, ip):
        logging.debug("##########ev_max_ptr_size")
        return 4

    # Get default enum size.
    # returns sizeof(enum)
    #
    def ev_get_default_enum_size(self, cm):
        logging.debug("##########ev_get_default_enum_size")
        return 4

    # Get register allocation convention for given calling convention.
    # 1, or 0=not implemented
    #
    def ev_get_cc_regs(self, regs, cc):
        logging.debug("##########ev_get_cc_regs")
        return 0

    # Get offset from SP to the first stack argument.
    #
    #
    def ev_get_stkarg_offset(self):
        logging.debug("##########ev_get_stkarg_offset")
        return 0

    # Get SIMD-related types according to given attributes ant/or argument location.
    # returns number of found types, -1-error If name==NULL, initialize all SIMD types
    #
    def ev_get_simd_types(self, out, simd_attrs, argloc, create_tifs):
        logging.debug("##########ev_get_simd_types")
        return 0

    # Calculate number of purged bytes after call.
    # returns number of purged bytes (usually add sp, N)
    # 
    def ev_calc_cdecl_purged_bytes(self, ea):
        logging.debug("##########ev_calc_cdecl_purged_bytes")
        return 0

    # Calculate number of purged bytes by the given function type.
    # 1, or 0 = not implemented
    #
    def ev_calc_purged_bytes(self, p_purged_bytes, fti):
        logging.debug( "##########ev_calc_purged_bytes" )
        logging.debug( fti.size() )
        logging.debug( fti.get_call_method() )
        logging.debug( fti.is_pure() )
        logging.debug( fti.rettype.get_size() )
        
        # Information about each argument IF there are any, NOTE that if the next instruction is unlink then NO purge is done
        ##########
        # NOTES:
        #   https://reverseengineering.stackexchange.com/questions/8870/extracting-arguments-from-ida
        #   https://github.com/gsmk/hexagon/blob/master/pmbase.h
        #   https://github.com/idapython/src/blob/master/pywraps/py_idp.hpp#L608
        #   LOCATION TYPE: https://www.hex-rays.com/products/ida/support/sdkdoc/group___a_l_o_c__.html
        #
        for i in xrange(fti.size()):
                logging.debug( "Arg %d: %s (of type %s, and of location: %s)" % (i, fti[i].name, idaapi.print_tinfo('', 0, 0, idaapi.PRTYPE_1LINE, fti[i].type, '', ''), fti[i].argloc.atype()) )

        # Sort of wrong, there could be shorts or other odd sized things on the stack
        p_purged_bytes = fti.size() * 4
        
        return 1

    # Calculate return value location.
    # 0 not implemented, 1 ok, -1 error
    #
    def ev_calc_retloc(self, retloc, rettype, cc):
        logging.debug( "##########ev_calc_retloc" )
        
        logging.debug( "retloc=" + str(retloc) ) 
        logging.debug( "rettype=" + str(rettype) )
        logging.debug( "cc=" + str(strcc) )
		
        # SEE: add_argregs in emu.cpp 
        # IF just a single reg then reg1, else if a range reg2 START##########END
        retloc.set_reg1(0)
        # IF its a range reg2 START##########END
        #retloc.set_reg2(0,0)
        # IF its a mix, then create a scattered_aloc_t, and push_back multiple part argpart_t
        #retloc.consume_scattered(0)
		
        #retloc.set_stkoff(0)
        #retloc.set_ea(0)
        #retloc.consume_rrel(0) # argloc_t_consume_rrel
        #argloc_t
        return 1


    # Calculate function argument locations.
    # This callback should fill retloc, all arglocs, and stkargs.
    # This callback is never called for ::CM_CC_SPECIAL functions.
    # 0=not implemented, 1=ok, -1=error
    #
    # See: alloc_args from arc emu.cpp in SDK
    def ev_calc_arglocs(self, fti):
        logging.debug(  "##########ev_calc_arglocs" )
        logging.debug(  "is_void() = " + str(fti.rettype.is_void()) )

        # Set the location of each argument
        #    NOTE: LOCATION TYPE: https://www.hex-rays.com/products/ida/support/sdkdoc/group___a_l_o_c__.html
        #
        for i in xrange(fti.size()):
            logging.debug(  "Arg %d: %s (of type %s, and of location: %s)" % (i, fti[i].name, idaapi.print_tinfo('', 0, 0, idaapi.PRTYPE_1LINE, fti[i].type, '', ''), fti[i].argloc.atype()) )
            # Set funciton argument #i location as register number "i"
#            fti[i].argloc.set_reg1(i)
            # Set funciton argument #i location as stack at location "i"*4
            fti[i].argloc.set_stkoff(i*4)
        # Return value is D0
        fti.retloc.set_reg1(0) # D0
        # Set size of stack arguments passed
        fti.stkargs = fti.size() * 4
        return 1

    # Calculate locations of the arguments that correspond to '...'
    # 0=not implemented, -1=error, 1=ok
    # 
    def ev_calc_varglocs(self, ftd, regs, stkargs, nfixed):
        logging.debug(  "##########ev_calc_varglocs" )
        return 0

    # Adjust argloc according to its type/size and platform endianess.
    # 0=not implemented, -1=error, 1=ok
    # 
    def ev_adjust_argloc(self, argloc, optional_type, size):
        logging.debug(  "##########ev_adjust_argloc" )
        return 0

    # Get function arguments which should be converted to pointers when lowering function prototype.
    # The processor module can also modify 'fti' in order to make non-standard conversion
    # 0=not implemented, 1=argnums was filled, 2 argnums was filled and made substantial changes to fti
    def ev_lower_func_type(self, argnums, fti):
        logging.debug(  "##########ev_lower_func_type" )
        return 0

    # Are 2 register arglocs the same?
    # 1=yes, -1=no, 0=not implemented
    # 
    def ev_equal_reglocs(self, a1, a2):
        logging.debug(  "##########ev_equal_reglocs" )
        return 0

    # Use information about a stack argument.
    #  retval <=0 failed, the kernel will create a comment with the
    #  argument name or type for the instruction
    #
    def ev_use_stkarg_type(self, ea, arg):
        logging.debug(  "##########ev_use_stkarg_type" )
        return 0

    # Use information about register argument.
    # 0=not implemented, 1=???
    # 
    def ev_use_regarg_type(self, ea, rargs):
        logging.debug(  "##########ev_use_regarg_type" )
        return 0

    # Use information about callee arguments.
    # 0=not implemented, 1 (and removes handled arguments from fti and rargs)
    # 
    def ev_use_arg_types(self, ea, fti, rargs):
        logging.debug(  "##########ev_use_arg_types" )
        return 0

    # Argument address info is ready.
    # <0: do not save into idb; other values mean "ok to save"
    #
    def ev_arg_addrs_ready(self, caller, n, tif, addrs):
        logging.debug(  "##########ev_arg_addrs_ready" )
        return 1

    # Decorate/undecorate a C symbol name, unlike C the return is the demanged name
    #   ev_decorate_name(self, name, mangle, cc, optional_type) -> PyObject *
    #
    def ev_decorate_name(self, name, mangle, cc, optional_type):
        logging.debug(  "##########ev_decorate_name" )
        logging.debug(  "NAME=" + str(name) )
        logging.debug(  "MANGLE?=" + str(mangle) )
        logging.debug(  "CC=" + str(cc) )
        logging.debug(  "OptionalType=" + str(optional_type) )
        logging.debug(  "procName=" + str(idaapi.get_inf_structure().procName) )
        
        # Compiler Info "compiler_info_t cc;                   ///< Target compiler" from ida.hpp
        logging.debug(  "cc.id=" + str(idaapi.get_inf_structure().cc.id) )
        
        #get_compiler_name
        # FROM: idaapi.gen_decorate_name(name, mangle, cc, type) -> bool 
        gen_decorate_name(name,mangle,cc,optional_type)
        logging.debug( "NAME POST Decoration:" + str(name) )
        return name

        
        
#--------------------------------------------------------------------------
# CTYPE mapping for global ph structure into python
#

# NOTE: Requires MSVCR90.dll, you will need the 2008 redist from MS for this to work right
#import ctypes
from ctypes import *
from binascii import hexlify
from sys import getsizeof
import ida_idp as ida_idp

# Partial structure information for accessing PH directly
#
class structPH(Structure):
        _fields_ = [("version", c_uint),
        ("id", c_uint),   
        ("flag", c_uint)]    


#--------------------------------------------------------------------------
class m6502_stepover(idaapi.plugin_t):
    # Processor fix plugin module
    flags = idaapi.PLUGIN_PROC | idaapi.PLUGIN_HIDE
    comment = ""
    wanted_hotkey = ""
    help = "Supports Step-Over for m6502 CPU on JSR, and type information"
    wanted_name = "m6502_stepover"
       
    def init(self):
        self.prochook = None
        
        # Enable controlable logging 
        logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(filename)s:%(funcName)s@%(lineno)s: %(message)s')
       
        logging.debug('Module for m6502 class CPU Step-Overs, and type information')
        
        # Random Debug to analyze PH mapping support
        #
        #print cdll.ida.ph
        #ph = structPH.in_dll(cdll.ida, "ph")
        #ph = structPH()
        #ph.values = cast( cdll.ida.ph, POINTER(structPH) )
		
        # Map ida.dll into this Python module, and lets us change the flags to add TYPEINFO support.
        #
        # For O/S agnostic support this follow advice from: http://www.hexblog.com/?p=695
        #
        idaname = "ida64" if ida_idaapi.__EA64__ else "ida"
        try:
            logging.debug( "Platform for ida library loading: " + str(sys.platform) )
            if sys.platform == "win32":
                ida_dll = ctypes.windll[idaname + ".dll"]
            elif sys.platform == "linux2":
                ida_dll = ctypes.cdll["lib" + idaname + ".so"]
            elif sys.platform == "darwin":
                ida_dll = ctypes.cdll["lib" + idaname + ".dylib"]        
        except:
            logging.error( "Python CTYPES probably couldnt find the dll, OR the MSCVR90.dll or other system DLL is missing" )
            return idaapi.PLUGIN_SKIP
		
        # Access ph global in ida.dll using ctypes
        ph = structPH.in_dll(ida_dll, "ph")
        
        # Log Values via PH direct access
        logging.debug( "- - - - - - - - - - - - - - - - - -" )
        logging.debug( "  ph.version = %d" % ph.version )
        logging.debug( "  ph.id      = %d" % ph.id )
        logging.debug( "  ph.flag    = %d" % ph.flag )

        # Enable TYPEINFO support
        ph.flag = ph.flag | ida_idp.PR_TYPEINFO

        # Official API, confirms PH was correctly accessed
        logging.debug( "- - - - - - - - - - - - - - - - - -" )
        logging.debug( "  ph.version = %d" % idaapi.ph_get_version() )
        logging.debug( "  ph.id      = %d" % idaapi.ph_get_id() )
        logging.debug( "  ph.flag    = %d" % idaapi.ph_get_flag() )

        #print(hexlify(string_at(cdll.ida.ph, getsizeof(G))))
			
        if idaapi.ph_get_id() != idaapi.PLFM_6502: # or idaapi.cvar.inf.filetype != idaapi.f_ELF:
            logging.error("skipped this plugin because CPU being used is not 6502 derivative!")
            return idaapi.PLUGIN_SKIP

        # Connect the IDP hooks
        self.prochook = m6502_stepover_idp_hook_t()
        self.prochook.hook()

        logging.debug( "succeeded, stepover support should now work for JSR!" )
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        pass

    def term(self):
        logging.debug( "ida is terminating, removing step-over support event handling!" )
        if self.prochook:
            self.prochook.unhook()

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


repository: metaworm/worm-ida
file: plugins/wormtools.py
description: IDA实用插件(A utility plugin for IDA)
wanted_name: WormTools
comment: https://github.com/metaworm/worm-ida
created_at: 2019-09-05 08:37:28+00:00
pushed_at: 2021-10-12 02:40:39+00:00
url:
https://github.com/metaworm/worm-ida/blob/3f159397030ac1a116b11b97eca6a155a9634ede/plugins/wormtools.py

content:
# -*- coding: UTF-8 -*-

from idaapi import *
from idc import *
from ida_kernwin import *
from ida_segment import *
from PyQt5.Qt import QApplication

import os
import worm_base as base

Choose2 = Choose if IDA_SDK_VERSION >= 740 else Choose2
askfile_c = ask_file if IDA_SDK_VERSION >= 740 else askfile_c

def copyto_clipboard(data):
    print('[WormTools] Copied: ' + data)
    QApplication.clipboard().setText(data)

def jumpto_offset(*args):
    addr = ask_addr(BADADDR, 'RVA:')
    if addr:
        jumpto(get_imagebase() + addr)

try: SegName
except NameError: SegName = lambda ea: get_segm_name(getseg(ea))

# load a *.map file
def load_map(*args):
    path = askfile_c(False, '*.map', 'Open a *.map file')
    if path == None:
        return
    path = path if os.path.exists(path) else path.decode('utf-8').encode('gb2312')
    if not os.path.exists(path):
        print('File Not Exists')

    print('+------------- Load Map: Begin --------------')
    collision = {}
    for line in open(path, 'r'):
        t = line.split()
        if len(t) > 3 and '0' in t[0]:
            ea = int(t[2], 16)
            if ea != 0:
                # ignore the function in then .idata segment
                if SegName(ea) == '.idata':
                    name = get_name(ea)
                    if name:
                        print('| Ignore Import Name: ' + name)
                        continue
                function_name = t[1].replace('<','_').replace('>','_').replace('\\','_')
                if collision.has_key(function_name):
                    collision[function_name].append(ea)
                    set_name(ea, function_name + '_' + str(len(collision[function_name])), SN_NOWARN)
                else:
                    collision[function_name] = []
                    set_name(ea, function_name, SN_NOWARN)
    print('+------------- Load Map: End --------------')

def map_names(*args):
    import json
    rule_path = askfile_c(False, '*.json', 'Open Map Rules')
    if rule_path == None:
        return

    print('[NameMap] Loading json rules: ' + rule_path)
    rules = json.load(open(rule_path))
    print('[NameMap] Begin replace %d functions' % len(rules))

    for (k, v) in rules.items():
        k = str(k); v = str(v)
        addr = LocByName(k)
        if os.environ.get('LOG_DEBUG') == 'ON':
            print('[NameMap] %s: %s -> %s' % (hex(addr), k, v))
        if addr != BADADDR:
            set_name(addr, v)
            MakeComm(addr, k)

    print('[NameMap] DONE')

class SettingsForm(Form):
    def __init__(self):
        self.title = 'WormTools'
        Form.__init__(self, r'''STARTITEM 0
Settings for WormTools

{FormChangeCb}
<#Hint1#模块名:{edit_modulename}>
    ''', {
            'edit_modulename': Form.StringInput(tp = Form.FT_ASCII),
            'FormChangeCb': Form.FormChangeCb(self.OnFormChange),
        })

    def OnFormChange(self, fid):
        if fid == -1:           # Init
            self.SetControlValue(self.edit_modulename, base.idb_get('MODULE_NAME', os.path.basename(get_input_file_path())))
        elif fid == -2:         # OK
            module = self.GetControlValue(self.edit_modulename)
            base.idb_set('MODULE_NAME', module)
            print('[WormTools] Set MODULE_NAME ' + module)
            self.Close(0)

def show_settings(*args):
    form = SettingsForm()
    form.Compile(); form.Execute(); form.Free()

class WormForm(Form):
    def __init__(self):
        self.title = 'WormTools'
        Form.__init__(self, r"""STARTITEM 0
WormTools

复制：
    <##(~W~)单词:{btn_word}> <##(~A~)地址:{btn_address}> <##(~R~)偏移:{btn_offset}> <##(~F~)函数名 :{btn_funcname}> <##(~D~)模块+偏移:{btn_modoff}>
跳转：
    <##(~H~)函数头:{btn_funchead}> <##(~T~)函数尾:{btn_functail}>
其他：
    <## Load~M~ap :{btn_loadmap}> <## Map~N~ame :{btn_mapname}> <##(~1~)PatchNOP:{btn_patchnop}> <## ~S~ettings :{btn_settings}>
        """, {
            'btn_word': Form.ButtonInput(self.on_word),
            'btn_address': Form.ButtonInput(self.on_address),
            'btn_offset': Form.ButtonInput(self.on_offset),
            'btn_funcname': Form.ButtonInput(self.on_funcname),
            'btn_modoff': Form.ButtonInput(self.on_modoff),
            'btn_funchead': Form.ButtonInput(self.on_funchead),
            'btn_functail': Form.ButtonInput(self.on_functail),
            'btn_loadmap': Form.ButtonInput(self.on_loadmap),
            'btn_mapname': Form.ButtonInput(self.on_mapname),
            'btn_settings': Form.ButtonInput(self.on_settings),
            'btn_patchnop': Form.ButtonInput(self.on_patchnop),
        })

    def on_word(self, code = 0):
        try:
            copyto_clipboard(get_highlight(get_current_viewer())[0])
        except:
            pass
        self.Close(0)

    def on_address(self, code = 0):
        copyto_clipboard('0x%x' % here())
        self.Close(0)

    def on_offset(self, code = 0):
        copyto_clipboard('0x%x' % (here() - get_imagebase()))
        self.Close(0)

    def on_funcname(self, code = 0):
        copyto_clipboard(get_func_name(here()))
        self.Close(0)

    def on_modoff(self, code = 0):
        name = base.idb_get('MODULE_NAME', os.path.basename(get_input_file_path()))
        copyto_clipboard('%s+0x%x' % (name, here() - get_imagebase()))
        self.Close(0)

    def on_funchead(self, code = 0):
        addr = get_func_attr(here(), FUNCATTR_START)
        if addr == BADADDR:
            print('FUNCATTR_START NOT FOUND')
        else:
            jumpto(addr)
        self.Close(0)

    def on_functail(self, code = 0):
        addr = get_func_attr(here(), FUNCATTR_END)
        if addr == BADADDR:
            print('FUNCATTR_END NOT FOUND')
        else:
            jumpto(addr)
        self.Close(0)

    def on_loadmap(self, code = 0):
        self.Close(0)
        load_map()

    def on_mapname(self, code = 0):
        self.Close(0)
        map_names()

    def on_settings(self, code = 0):
        self.Close(0)
        show_settings()

    def on_patchnop(self, code = 0):
        from ida_bytes import patch_bytes
        patch_bytes(here(), '\x90' * get_item_size(here()))
        self.Close(0)

def show_wormtools(*args):
    form = WormForm()
    form.Compile(); form.Execute(); form.Free()

class JumpChooser(Choose2):
    def __init__(self, title, cols, items, flags=0, width=None, height=None, embedded=False, modal=False):
        Choose2.__init__(
            self,
            title,
            cols,
            flags = flags,
            width = width,
            height = height,
            embedded = embedded)
        self.items = items
        self.selcount = 0

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

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

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

def local_xref(ctx):
    func = get_func(here())
    start_ea, end_ea = func.start_ea, func.end_ea

    xrefs = []
    # refto = get_first_dref_from(here())
    refto = get_highlight(get_current_viewer())
    if not refto: return

    refto = LocByName(refto[0])
    if refto == BADADDR: return

    for ref in XrefsTo(refto):
        if ref.frm >= start_ea and ref.frm < end_ea:
            # print(hex(ref.frm))
            xrefs.append([str(ref.type), atoa(ref.frm), ref.frm])
    JumpChooser('local xref to ' + get_name(refto), [['Type', 10], ['Address', 30|Choose2.CHCOL_HEX]], xrefs).Show(True)

class WormTools(plugin_t):
    flags         = PLUGIN_KEEP | PLUGIN_FIX
    wanted_hotkey = ''
    wanted_name   = "WormTools"
    comment       = "https://github.com/metaworm/worm-ida"
    help          = "https://github.com/metaworm/worm-ida"

    def init(self):
        base.add_menu_item('Jump/', 'Jump to Offset', 'worm:jumptooffset',
                base.ActionHandler(jumpto_offset, enable = AST_ENABLE_FOR_IDB), 'Shift+G', 'Jump To Offset')
        base.add_menu_item('Edit/', 'WormTools', 'worm:wormtools',
                base.ActionHandler(show_wormtools, enable = lambda: True), 'W', 'WormTools')
        base.add_menu_item('Edit/', 'Name Mappings(~f~)', 'worm:namemap',
                base.ActionHandler(map_names), tooltip = 'Rename symbols by a table')

        base.add_menu_item('File/Load file/', '~M~ap file...', 'worm:loadmap',
                base.ActionHandler(load_map), tooltip = 'Load *.map file')
        base.add_menu_item('Jump/', 'Jump to local xrefs', 'worm:local_xref',
                base.ActionHandler(local_xref), 'Shift+X')

        msg("[WormTools] Plugin Initialized.\n")
        return PLUGIN_OK

    def term(self):
        msg('[WormTools] Plugin Terminated.\n')

    def run(self, arg):
        show_settings()

def PLUGIN_ENTRY():
    return WormTools()


repository: Cisco-Talos/DynDataResolver
file: ddr_installer/install_data/ddr_plugin.py
wanted_name: DDR
comment: Dynamic Data Resolver Plugin
created_at: 2018-12-10 22:28:40+00:00
pushed_at: 2022-12-08 10:02:47+00:00
url:
https://github.com/Cisco-Talos/DynDataResolver/blob/e630bd3d81062f790d932bbe7c594fde93da5f5f/ddr_installer/install_data/ddr_plugin.py

content:
#-------------------------------------------------------------------------------
#
#   IDA Pro Plug-in: Dynamic Data Resolver (DDR) Front End
#
#   Version 1.02beta 
#
#   Copyright (C) 2020 Cisco Talos
#   Author: Holger Unterbrink (hunterbr@cisco.com)
#
#   This program is free software; you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation; either version 2 of the License, or
#   (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License along
#   with this program; if not, write to the Free Software Foundation, Inc.,
#   51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
#
#   Requirements
#   ------------
#                           
#   Requests    (http://docs.python-requests.org)       
#   PEfile      (https://github.com/erocarrera/pefile)
#
#   e.g.:
#   python -m pip install --upgrade pip                     
#   pip install -U Requests   
#   pip install -U pefile
#
#   Hint: *Make sure you install these requirements for the same Python version IDA is using. 
# 
#  
#   Common setup:
#   -------------
#
#   Analyst PC (IDA, ddr_plugin.py) 
#              | 
#              | IP network (allow port DDR_CONFIG_SETTINGS.WEBSERVER_PORT (Default: 5000))
#              |
#              | IP addr = [DDR_CONFIG_SETTINGS.WEBSERVER/DDR_CONFIG_SETTINGS.WEBSERVER_PORT variable]
#              V
#   Malware PC (DynamoRio, ddr32.dll, ddr64.dll, ddr_server.py)
#           
#-----------------------------------------------------------------------------------

import sys            
import idaapi
import ida_kernwin
import ida_nalt
import ida_ua
import idautils
import ida_ida
import idc
import logging
import tempfile
import os
import glob
import errno
import json
import subprocess    
import zipfile
import hashlib
import time
import traceback
from datetime import datetime
from pprint import pprint
from collections import Counter, OrderedDict

# IDA version below 7.5 and Python equal or higher than 3.8 doesn't work
if idaapi.IDA_SDK_VERSION < 750 and sys.version_info >= (3, 8):
    print("[DDR][ERROR] Python version > 3.7 found. Problematic Python modules not loaded.")
    print("[DDR][ERROR] Python 3.8+ and DDR only works in IDA 7.5 and above.")
else:
    # these modules crash IDA if Python version > 3.7 is installed.
    from PyQt5 import QtCore, QtGui, QtWidgets
    import requests
    import pefile
        
# Setup Logging to file in %TEMP% --- Not used itm ---
# tmpdir = tempfile.gettempdir()
# ddr_logdir    = tmpdir + '\DDR'

# try:
#     os.makedirs(ddr_logdir, exist_ok=True)    # Python >3.2
# except TypeError:
#     try:
#         os.makedirs(ddr_logdir)               # Python >2.5
#     except OSError as e:
#         if e.errno == errno.EEXIST and os.path.isdir(ddr_logdir):
#             pass
#         else:
#             idaapi.warning("[DDR] Failed to make logging directory")

# ddr_logfilename = ddr_logdir + '\DDR-{:%Y-%m-%d-%H-%M-%S}.log'.format(datetime.now())
# logger = logging.getLogger("DDR")
# logger.setLevel(logging.DEBUG)
# fh = logging.FileHandler(ddr_logfilename)
# fh.setLevel(logging.DEBUG)
# formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
# fh.setFormatter(formatter)
# logger.addHandler(fh)

# Globals
# -------
DDR_PLUGIN_VERSION = "1.02 beta"
# Global DDR plugin configuration file
# DDR_CFG_FILE = idaapi.idadir("plugins\\ddr\\ddr_config.json") 
DDR_CFG_FILE = "C:\\Users\\Dex Dexter\\AppData\\Roaming\\Hex-Rays\\IDA Pro\\plugins\\ddr\\ddr_config.json"

# Colors used for highlighting the building blocks. If you want to change this, change the assignment to COLOR1-5 below.                                                     
GREEN                   = [0xa0ffa0, 0x70ff70, 0x00ff00, 0x00ef00, 0x00df00, 0x00cf00, 0x00bf00, 0x00af00]      
YELLOW                  = [0xd0ffff, 0xa0ffff, 0x70ffff, 0x00ffff, 0x00efff, 0x00dfff, 0x00cfff, 0x00bfff]      
RED                     = [0xC0C0ff, 0xB0B0ff, 0xA0A0ff, 0x8080ff, 0x6060ff, 0x4040ff, 0x2020ff, 0x0000ff]      
LIGHT_PURPLE            = 0xFF60FF                                                                              
PURPLE                  = 0xFF00FF                                                                                  
                                             # Number of times instruction executed                                                                                                                     
COLOR1                  = GREEN              # 0x0     - 0xF                             list [] of 8 colors
COLOR2                  = YELLOW             # 0x10    - 0xff                            list [] of 8 colors
COLOR3                  = RED                # 0x100   - 0xfff                           list [] of 8 colors
COLOR4                  = LIGHT_PURPLE       # 0x1000  - 0xffff                          color value
COLOR5                  = PURPLE             # 0x10000 - infinity                        color value
                 
EFLAGS                  = [ "CF", "PF", "AF", "ZF", "SF", "DF", "OF" ]

SAMPLE_FILENAME         = None
SAMPLE_SHA256           = None
SAMPLE_DIR              = None

JSONDATA                = None
JSONFILE_LOADED         = False
APIFILE_LOADED          = False
        
API_CALLS               = {}              
REG_LIST                = []

# Setup menu/action items: 
menu_items = OrderedDict([
("DDR_Action_Load_file"                   , { "menu_str":"Load DDR trace file"                              , "hotkey":"Ctrl+Shift+F11", "submenu":""              , "ah_id":"LoadFile"                   , "x64only":False, "hide_in_context":True }),
("DDR_Action_Load_api_file"               , { "menu_str":"Load DDR trace API file"                          , "hotkey":"Ctrl+Shift+F12", "submenu":""              , "ah_id":"LoadAPIFile"                , "x64only":False, "hide_in_context":True }),
("DDR_Action_Add_BB2List"                 , { "menu_str":"Add basic block to basic block list"              , "hotkey":"Ctrl+Shift+A"  , "submenu":"Select/"       , "ah_id":"Select_Add_BB2List"         , "x64only":False, "hide_in_context":False}),
("DDR_Action_Remove_BB2List"              , { "menu_str":"Remove basic block from basic block list"         , "hotkey":"Ctrl+Shift+D"  , "submenu":"Select/"       , "ah_id":"Select_Remove_BB2List"      , "x64only":False, "hide_in_context":False}),
("DDR_Action_Print_BB2List"               , { "menu_str":"Show basic block list"                            , "hotkey":None            , "submenu":"Select/"       , "ah_id":"Select_Print_BB2List"       , "x64only":False, "hide_in_context":False}),
("DDR_Action_Clear_BB2List"               , { "menu_str":"Clear basic block list"                           , "hotkey":None            , "submenu":"Select/"       , "ah_id":"Select_Clear_BB2List"       , "x64only":False, "hide_in_context":False}),
("DDR_Action_SetStartOffset"              , { "menu_str":"Set start offset"                                 , "hotkey":None            , "submenu":"Select/"       , "ah_id":"Select_SetStartOffset"      , "x64only":False, "hide_in_context":False}),
("DDR_Action_ClearStartOffset"            , { "menu_str":"Clear start offset"                               , "hotkey":None            , "submenu":"Select/"       , "ah_id":"Select_ClearStartOffset"    , "x64only":False, "hide_in_context":False}),
("DDR_Action_ShowStartOffset"             , { "menu_str":"Show start offset"                                , "hotkey":None            , "submenu":"Select/"       , "ah_id":"Select_ShowStartOffset"     , "x64only":False, "hide_in_context":False}),
("DDR_Action_SetBreakaddress"             , { "menu_str":"Set break address"                                , "hotkey":None            , "submenu":"Select/"       , "ah_id":"Select_SetBreakaddress"     , "x64only":False, "hide_in_context":False}),
("DDR_Action_ClearBreakaddress"           , { "menu_str":"Clear break address"                              , "hotkey":None            , "submenu":"Select/"       , "ah_id":"Select_ClearBreakaddress"   , "x64only":False, "hide_in_context":False}),
("DDR_Action_ShowBreakaddress"            , { "menu_str":"Show break address"                               , "hotkey":None            , "submenu":"Select/"       , "ah_id":"Select_ShowBreakaddress"    , "x64only":False, "hide_in_context":False}),
("DDR_Action_SetCmdArgs"                  , { "menu_str":"Set command line args"                            , "hotkey":None            , "submenu":"Select/"       , "ah_id":"Select_SetCmdArgs"          , "x64only":False, "hide_in_context":False}),
("DDR_Action_ClearCmdArgs"                , { "menu_str":"Clear command line args"                          , "hotkey":None            , "submenu":"Select/"       , "ah_id":"Select_ClearCmdArgs"        , "x64only":False, "hide_in_context":False}),
("DDR_Action_ShowCmdArgs"                 , { "menu_str":"Show command line args"                           , "hotkey":None            , "submenu":"Select/"       , "ah_id":"Select_ShowCmdArgs"         , "x64only":False, "hide_in_context":False}),
("DDR_Action_Trace_On_Range"              , { "menu_str":"Run full trace for marked address range"          , "hotkey":None            , "submenu":"Trace/"        , "ah_id":"Run_Trace_On_Range"         , "x64only":False, "hide_in_context":False}),
("DDR_Action_Run_Trace_On_BB"             , { "menu_str":"Run full trace for marked basic block"            , "hotkey":None            , "submenu":"Trace/"        , "ah_id":"Run_Trace_On_BB"            , "x64only":False, "hide_in_context":False}),
("DDR_Action_Run_Trace_On_BBList"         , { "menu_str":"Run full trace for basic block list"              , "hotkey":None            , "submenu":"Trace/"        , "ah_id":"Run_Trace_On_BB_list"       , "x64only":False, "hide_in_context":False}),
("DDR_Action_Run_Trace_On_Seg"            , { "menu_str":"Run full trace for segment"                       , "hotkey":None            , "submenu":"Trace/"        , "ah_id":"Run_Trace_On_Seg"           , "x64only":False, "hide_in_context":False}),
("DDR_Action_Run_Light_Trace_On_Seg"      , { "menu_str":"Run light trace for segment"                      , "hotkey":None            , "submenu":"Trace/"        , "ah_id":"Run_Light_Trace_On_Seg"     , "x64only":False, "hide_in_context":False}),
#("DDR_Action_Delete_Cached_Files"         , { "menu_str":"Delete cached traces"                             , "hotkey":None            , "submenu":"Trace/"        , "ah_id":"Delete_Cached_Files"        , "x64only":False, "hide_in_context":False}),
("DDR_Action_Dump_buffer_size"            , { "menu_str":"Use marked operand to get buffer size"            , "hotkey":None            , "submenu":"Dump/"         , "ah_id":"Dump_buffer_size_op"        , "x64only":False, "hide_in_context":False}),
("DDR_Action_Dump_buffer_addr"            , { "menu_str":"Use marked operand to get buffer address"         , "hotkey":None            , "submenu":"Dump/"         , "ah_id":"Dump_buffer_addr_op"        , "x64only":False, "hide_in_context":False}),
("DDR_Action_Dump_buffer_dump"            , { "menu_str":"Use marked address to dump buffer to file"        , "hotkey":None            , "submenu":"Dump/"         , "ah_id":"Dump_buffer_dump_addr"      , "x64only":False, "hide_in_context":False}),
("DDR_Action_Dump_buffer"                 , { "menu_str":"Execute sample and dump buffer"                   , "hotkey":None            , "submenu":"Dump/"         , "ah_id":"Dump_buffer"                , "x64only":False, "hide_in_context":False}),
("DDR_Action_Dump_buffer_nop"             , { "menu_str":"NOP out marked instruction at runtime"            , "hotkey":None            , "submenu":"Patch/"        , "ah_id":"Dump_buffer_nop_instr"      , "x64only":False, "hide_in_context":False}),
("DDR_Action_Dump_buffer_eflag"           , { "menu_str":"Toggle EFLAG at runtime"                          , "hotkey":None            , "submenu":"Patch/"        , "ah_id":"Dump_buffer_eflag_instr"    , "x64only":False, "hide_in_context":False}),
("DDR_Action_Dump_buffer_call"            , { "menu_str":"Skip function at marked address at runtime"       , "hotkey":None            , "submenu":"Patch/"        , "ah_id":"Dump_buffer_skip_call"      , "x64only":False, "hide_in_context":False}),
("DDR_Action_sample_run_only"             , { "menu_str":"Run patched sample"                               , "hotkey":None            , "submenu":"Patch/"        , "ah_id":"Sample_run_only"            , "x64only":False, "hide_in_context":False}),
("DDR_Action_sample_to_x64dbg"            , { "menu_str":"Create x64dbg script with bp at marked address"   , "hotkey":None            , "submenu":"Debugger/"     , "ah_id":"Sample_to_x64dbg"           , "x64only":False, "hide_in_context":False}),
("DDR_Action_sample_loop_addr"            , { "menu_str":"Create executable with loop at marked address"    , "hotkey":None            , "submenu":"Debugger/"     , "ah_id":"Sample_loop_addr"           , "x64only":False, "hide_in_context":False}),
("DDR_Action_Set_Num_Hits_for_Cmt"        , { "menu_str":"Set number of trace hits for IDA DISASM View"     , "hotkey":None            , "submenu":"Config/"       , "ah_id":"Get_Set_Num_Hits_Cmt"       , "x64only":False, "hide_in_context":False}),
("DDR_Action_Set_Num_Hits_for_IdaLog"     , { "menu_str":"Set number of trace hits for IDA log window"      , "hotkey":None            , "submenu":"Config/"       , "ah_id":"Get_Set_Num_Hits_IdaLog"    , "x64only":False, "hide_in_context":False}),
("DDR_Action_Set_Num_Max_Instr"           , { "menu_str":"Set number of instructions to log at runtime"     , "hotkey":None            , "submenu":"Config/"       , "ah_id":"Get_Set_Max_Instr"          , "x64only":False, "hide_in_context":False}),
("DDR_Action_Set_Num_API_timeout"         , { "menu_str":"Set number of seconds for API timeout"            , "hotkey":None            , "submenu":"Config/"       , "ah_id":"Get_Set_API_timeout"        , "x64only":False, "hide_in_context":False}),
("DDR_Action_Dump_buffer_edit_cfg"        , { "menu_str":"Edit DDR API request (Experts only)"              , "hotkey":None            , "submenu":"Config/"       , "ah_id":"Dump_buffer_edit_cfg"       , "x64only":False, "hide_in_context":False}),
("DDR_Action_Dump_buffer_clear_cfg"       , { "menu_str":"Clear configured DDR API request"                 , "hotkey":None            , "submenu":"Config/"       , "ah_id":"Dump_buffer_clear_cfg"      , "x64only":False, "hide_in_context":False}),
("DDR_Action_Get_Mem_Ptr_xax"             , { "menu_str":"Get memory for ptr in xax"                        , "hotkey":None            , "submenu":"Get Register/" , "ah_id":"Get_Mem_Ptr_xax"            , "x64only":False, "hide_in_context":False}),
("DDR_Action_Get_Mem_Ptr_xbx"             , { "menu_str":"Get memory for ptr in xbx"                        , "hotkey":None            , "submenu":"Get Register/" , "ah_id":"Get_Mem_Ptr_xbx"            , "x64only":False, "hide_in_context":False}),
("DDR_Action_Get_Mem_Ptr_xcx"             , { "menu_str":"Get memory for ptr in xcx"                        , "hotkey":None            , "submenu":"Get Register/" , "ah_id":"Get_Mem_Ptr_xcx"            , "x64only":False, "hide_in_context":False}),
("DDR_Action_Get_Mem_Ptr_xdx"             , { "menu_str":"Get memory for ptr in xdx"                        , "hotkey":None            , "submenu":"Get Register/" , "ah_id":"Get_Mem_Ptr_xdx"            , "x64only":False, "hide_in_context":False}),
("DDR_Action_Get_Mem_Ptr_xsp"             , { "menu_str":"Get memory for ptr in xsp"                        , "hotkey":None            , "submenu":"Get Register/" , "ah_id":"Get_Mem_Ptr_xsp"            , "x64only":False, "hide_in_context":False}),
("DDR_Action_Get_Mem_Ptr_xbp"             , { "menu_str":"Get memory for ptr in xbp"                        , "hotkey":None            , "submenu":"Get Register/" , "ah_id":"Get_Mem_Ptr_xbp"            , "x64only":False, "hide_in_context":False}),
("DDR_Action_Get_Mem_Ptr_xsi"             , { "menu_str":"Get memory for ptr in xsi"                        , "hotkey":None            , "submenu":"Get Register/" , "ah_id":"Get_Mem_Ptr_xsi"            , "x64only":False, "hide_in_context":False}),
("DDR_Action_Get_Mem_Ptr_xdi"             , { "menu_str":"Get memory for ptr in xdi"                        , "hotkey":None            , "submenu":"Get Register/" , "ah_id":"Get_Mem_Ptr_xdi"            , "x64only":False, "hide_in_context":False}),
("DDR_Action_Get_Mem_Ptr_r8"              , { "menu_str":"Get memory for ptr in r8"                         , "hotkey":None            , "submenu":"Get Register/" , "ah_id":"Get_Mem_Ptr_r8"             , "x64only":True , "hide_in_context":False}),
("DDR_Action_Get_Mem_Ptr_r9"              , { "menu_str":"Get memory for ptr in r9"                         , "hotkey":None            , "submenu":"Get Register/" , "ah_id":"Get_Mem_Ptr_r9"             , "x64only":True , "hide_in_context":False}),
("DDR_Action_Get_Mem_Ptr_r10"             , { "menu_str":"Get memory for ptr in r10"                        , "hotkey":None            , "submenu":"Get Register/" , "ah_id":"Get_Mem_Ptr_r10"            , "x64only":True , "hide_in_context":False}),
("DDR_Action_Get_Mem_Ptr_r11"             , { "menu_str":"Get memory for ptr in r11"                        , "hotkey":None            , "submenu":"Get Register/" , "ah_id":"Get_Mem_Ptr_r11"            , "x64only":True , "hide_in_context":False}),
("DDR_Action_Get_Mem_Ptr_r12"             , { "menu_str":"Get memory for ptr in r12"                        , "hotkey":None            , "submenu":"Get Register/" , "ah_id":"Get_Mem_Ptr_r12"            , "x64only":True , "hide_in_context":False}),
("DDR_Action_Get_Mem_Ptr_r13"             , { "menu_str":"Get memory for ptr in r13"                        , "hotkey":None            , "submenu":"Get Register/" , "ah_id":"Get_Mem_Ptr_r13"            , "x64only":True , "hide_in_context":False}),
("DDR_Action_Get_Mem_Ptr_r14"             , { "menu_str":"Get memory for ptr in r14"                        , "hotkey":None            , "submenu":"Get Register/" , "ah_id":"Get_Mem_Ptr_r14"            , "x64only":True , "hide_in_context":False}),
("DDR_Action_Get_Mem_Ptr_r15"             , { "menu_str":"Get memory for ptr in r15"                        , "hotkey":None            , "submenu":"Get Register/" , "ah_id":"Get_Mem_Ptr_r15"            , "x64only":True , "hide_in_context":False}),
("DDR_Action_GetSrcOpValue"               , { "menu_str":"Get values for source operand"                    , "hotkey":None            , "submenu":""              , "ah_id":"GetSrcOpValue"              , "x64only":False, "hide_in_context":False}),
("DDR_Action_GetDstOpValue"               , { "menu_str":"Get values for destination operand"               , "hotkey":None            , "submenu":""              , "ah_id":"GetDstOpValue"              , "x64only":False, "hide_in_context":False}),
("DDR_Action_GetSrcOpPtrValue"            , { "menu_str":"Get values for ptr in source operand"             , "hotkey":None            , "submenu":""              , "ah_id":"GetSrcOpPtrValue"           , "x64only":False, "hide_in_context":False}),
("DDR_Action_GetDstOpPtrValue"            , { "menu_str":"Get values for ptr in destination operand"        , "hotkey":None            , "submenu":""              , "ah_id":"GetDstOpPtrValue"           , "x64only":False, "hide_in_context":False}),
("DDR_Action_GetSrcOpPtrPtrValue"         , { "menu_str":"Get values for ptr ptr in source operand"         , "hotkey":None            , "submenu":""              , "ah_id":"GetSrcOpPtrPtrValue"        , "x64only":False, "hide_in_context":False}),
("DDR_Action_GetDstOpPtrPtrValue"         , { "menu_str":"Get values for ptr ptr in destination operand"    , "hotkey":None            , "submenu":""              , "ah_id":"GetDstOpPtrPtrValue"        , "x64only":False, "hide_in_context":False}),
("DDR_Action_GetSrcPtrMem"                , { "menu_str":"Get memory for ptr in source operand"             , "hotkey":None            , "submenu":""              , "ah_id":"GetSrcPtrMem"               , "x64only":False, "hide_in_context":False}),
("DDR_Action_GetDstPtrMem"                , { "menu_str":"Get memory for ptr in destination operand"        , "hotkey":None            , "submenu":""              , "ah_id":"GetDstPtrMem"               , "x64only":False, "hide_in_context":False}),
("DDR_Action_GetSrcPtrPtrMem"             , { "menu_str":"Get memory for ptr ptr in source operand"         , "hotkey":None            , "submenu":""              , "ah_id":"GetSrcPtrPtrMem"            , "x64only":False, "hide_in_context":False}),
("DDR_Action_GetDstPtrPtrMem"             , { "menu_str":"Get memory for ptr ptr in destination operand"    , "hotkey":None            , "submenu":""              , "ah_id":"GetDstPtrPtrMem"            , "x64only":False, "hide_in_context":False}),
("DDR_Action_Highlight_trace"             , { "menu_str":"Highlight traced instructions"                    , "hotkey":None            , "submenu":""              , "ah_id":"Highlight_trace"            , "x64only":False, "hide_in_context":False}),
("DDR_Action_Clear_Highlight_trace"       , { "menu_str":"Clear highlighted instructions"                   , "hotkey":None            , "submenu":""              , "ah_id":"Clear_highlighted_blocks"     , "x64only":False, "hide_in_context":False}),
("DDR_Action_Strings_View"                , { "menu_str":"Show strings trace"                               , "hotkey":None            , "submenu":""              , "ah_id":"Display_Strings_View"       , "x64only":False, "hide_in_context":False}),
("DDR_Action_Call_View"                   , { "menu_str":"Show API calls trace"                             , "hotkey":None            , "submenu":""              , "ah_id":"Display_Calls_View"         , "x64only":False, "hide_in_context":False}),
("DDR_Action_DeleteNonRepeatableComments" , { "menu_str":"Delete non-repeatable comments"                   , "hotkey":None            , "submenu":""              , "ah_id":"DeleteNonRepeatableComments", "x64only":False, "hide_in_context":False})
])

def PLUGIN_ENTRY():
    """ Set plugin entry class """
    return ddrPlugin()

def ddr_exception_handler1(msg, ex):
    """ DDR exception handler. Handle exceptions without killing the plugin """
    print("================================================= Exception ===========================================================")
    print(msg)
    print("-----------------------------------------------------------------------------------------------------------------------")
    ex_type, ex_value, ex_traceback = sys.exc_info()
    trace_back = traceback.extract_tb(ex_traceback)
    stack_trace = []
    for trace in trace_back:
        stack_trace.append("[DDR][ERROR] File: {}, Func.Name: {}, Line: {:d} \n[DDR][ERROR] Issue : {}".format(trace[0], trace[2], trace[1], trace[3]))
    print("[DDR][ERROR] Exception type : {} ".format(ex_type.__name__))
    print("[DDR][ERROR] Exception message : {}".format(ex_value))
    for msg in stack_trace:
        print("[DDR][ERROR] Stack trace : \n{}".format(msg))
    print("=======================================================================================================================")

class ddr_plugin_cfg():
    """ Class storing the global DDR configuration """  

    def __init__(self, ddr_cfg_file_json):
        self.reload(ddr_cfg_file_json)

    def reload(self, ddr_cfg_file_json):
        try:
            # IP address of host ddr_server.py is running on
            keyword = "WEBSERVER"
            self.WEBSERVER               = ddr_cfg_file_json[keyword]
            print("[DDR][INFO] Configured DDR server is      : {}".format(self.WEBSERVER))
            # Port DDRserver.py is using 
            keyword = "WEBSERVER_PORT"
            self.WEBSERVER_PORT          = ddr_cfg_file_json[keyword]
            print("[DDR][INFO] Configured DDR server port is : {}".format(self.WEBSERVER_PORT))
            # API key, check ddr_server.py start messages, generated by the ddr_server.py script at first startup
            keyword = "DDR_WEBAPI_KEY"
            self.DDR_WEBAPI_KEY          = ddr_cfg_file_json[keyword]
            # CA certificate
            keyword = "CA_CERT"
            self.CA_CERT                 = ddr_cfg_file_json[keyword]
            # Verify certificates
            keyword = "VERIFY_CERT"
            self.VERIFY_CERT             = ddr_cfg_file_json[keyword]
            # Config filename for dump buffer function. This file will be written into the sample directory
            keyword = "DUMP_CFG_FILE"
            self.DUMP_CFG_FILE           = ddr_cfg_file_json[keyword]
            # Max. timeout the ddr_server.py has for answering in seconds, before we kill the request.
            # Don't set this too low, it should be long enough to run the DynamoRio client analysis.
            keyword = "MAX_API_TIMEOUT"
            self.MAX_API_TIMEOUT         = ddr_cfg_file_json[keyword]
            # Debug Level, the higher the more is printed to the IDA log window. See DDR_print_mesg() for details.
            keyword = "DBG_LEVEL"
            self.DBG_LEVEL               = ddr_cfg_file_json[keyword]
            # Default max. number of instructions to execute before Dynamorio stops logging.  
            keyword = "MAX_INSTR_TO_EXECUTE"
            self.MAX_INSTR_TO_EXECUTE    = ddr_cfg_file_json[keyword]
            # Default max. number of found instruction values added to the IDA log window (max. 50, see below). 
            keyword = "MAX_LOG_ROUNDS"
            self.MAX_LOG_ROUNDS          = ddr_cfg_file_json[keyword]
            # Default max. number of found instruction values added to the IDA DISASM view as comments (max. 50, see below).
            keyword = "MAX_CMT_ROUNDS"
            self.MAX_CMT_ROUNDS          = ddr_cfg_file_json[keyword]
            # Max. number of instruction addresses wich will be added to the trace_instr_num_list list.  
            # This list holds all the occurences of instruction numbers which match the ea address.    
            # This is the max value for the MAX_LOG_ROUNDS and MAX_CMT_ROUNDS variables.       
            keyword = "MAX_INSTR_COUNT"
            self.MAX_INSTR_COUNT         = ddr_cfg_file_json[keyword]
            # Max. number of attempts to copy the sample file over to the DDR server
            keyword = "MAX_UPLOAD_ATTEMPTS"
            self.MAX_UPLOAD_ATTEMPTS     = ddr_cfg_file_json[keyword]
            # Sample architecture
            self.ARCH_BITS               = None
            # Config is successful initalized
            self.is_initalized = "PHASE1"
            print("[DDR][INFO] Configuration successfully parsed.")
        except:
            print("[DDR][ERROR] Failed to initalize {} settings.".format(keyword))
            self.is_initalized = False

# Load global DDR configuration
try:
    print("[DDR][INFO] Initalizing DDR plugin version {}.".format(DDR_PLUGIN_VERSION))
    with open(DDR_CFG_FILE) as json_file:
        ddr_cfg_file_json = json.load(json_file)
        print("[DDR][INFO] Configuration loaded from {}".format(DDR_CFG_FILE))
        DDR_CONFIG_SETTINGS = ddr_plugin_cfg(ddr_cfg_file_json)
except:
    if DDR_CFG_FILE: 
        print("[DDR][ERROR] Failed loading DDR configuration file: {}".format(DDR_CFG_FILE))
        raise
    else:
        print("[DDR][ERROR] Failed loading DDR configuration file.")
        raise


def DDR_print_mesg(msg, debuglevel=0, printout=False):
    """ Print msg to output window in IDA """
    func = "not implemented"

    if printout:
        if debuglevel == 0 and DDR_CONFIG_SETTINGS.DBG_LEVEL == 8:
            print("[DDR][INFO] {}".format(msg))
        if debuglevel == 1 and DDR_CONFIG_SETTINGS.DBG_LEVEL == 9:
            print("[DDR][WARNING] {}".format(msg))
        if debuglevel == 2 and DDR_CONFIG_SETTINGS.DBG_LEVEL == 10:
            print("[DDR][ERROR] {}".format(msg))
        return

    if debuglevel == 0 and DDR_CONFIG_SETTINGS.DBG_LEVEL >= 0:
        print("[DDR][INFO] {}".format(msg))
    if debuglevel == 1 and DDR_CONFIG_SETTINGS.DBG_LEVEL >= 1:
        print("[DDR][WARNING] {}".format(msg))
    if debuglevel == 2 and DDR_CONFIG_SETTINGS.DBG_LEVEL >= 2:
        print("[DDR][ERROR] {}".format(msg))
    if debuglevel == 3 and DDR_CONFIG_SETTINGS.DBG_LEVEL >= 3:                           
        print("[DDR][DEBUG_L1][{}] {}".format(func, msg))
    if debuglevel == 4 and DDR_CONFIG_SETTINGS.DBG_LEVEL >= 4:
        print("[DDR][DEBUG_L2][{}] {}".format(func, msg))
    if debuglevel == 5 and DDR_CONFIG_SETTINGS.DBG_LEVEL >= 5:
        print("[DDR][DEBUG_L3][{}] {}".format(func, msg))
    if debuglevel == 6 and DDR_CONFIG_SETTINGS.DBG_LEVEL >= 6:
        print("[DDR][DEBUG_L4][{}] {}".format(func, msg))  
    if debuglevel == 7 and DDR_CONFIG_SETTINGS.DBG_LEVEL >= 7:
        print("[DDR][DEBUG_L5][{}] {}".format(func, msg))

def DDR_upload_sample(filename):
    """ Upload sample to DDR server """
    url = "https://" + DDR_CONFIG_SETTINGS.WEBSERVER + ":" + DDR_CONFIG_SETTINGS.WEBSERVER_PORT + "/" +  "uploadsample"
    files = {'file': open(filename, 'rb')}
    postpara = { 'apikey' : DDR_CONFIG_SETTINGS.DDR_WEBAPI_KEY, 'id' : '0' }

    DDR_print_mesg("Uploading file: {} ...".format(filename))

    if DDR_CONFIG_SETTINGS.VERIFY_CERT:
        res = requests.post(url, verify=DDR_CONFIG_SETTINGS.CA_CERT, files=files, data=postpara, timeout=DDR_CONFIG_SETTINGS.MAX_API_TIMEOUT)
    else:
        res = requests.post(url, verify=False, files=files, data=postpara, timeout=DDR_CONFIG_SETTINGS.MAX_API_TIMEOUT) 
    
    DDR_print_mesg("File upload status: {}".format(res.json()["return_status"]))

    if res.status_code == 200:
        return True
    else:
        return False

def get_hash(filename):
    """ Calculate MD5,SHA1,SHA256 hash for given file """
    BUF_SIZE = 65536
    md5 = hashlib.md5()
    sha1 = hashlib.sha1()
    sha256 = hashlib.sha256()
    with open(filename, 'rb') as f:
        while True:
            data = f.read(BUF_SIZE)
            if not data:
                break
            md5.update(data)
            sha1.update(data)
            sha256.update(data)
    return {"md5_sum"   :md5.hexdigest().upper(),
            "sha1_sum"  :sha1.hexdigest().upper(), 
            "sha256_sum":sha256.hexdigest().upper()}


class basicblock_obj(object):
    """ Basic Block class """
    num         = 0
    start_addr  = 0
    end_addr    = 0

    def __init__(self, start_addr, end_addr):
        self.start      = start_addr
        self.end        = end_addr

    def __eq__(self, other):
        if self.start == other.start and self.end == other.end:
            return True
        else:
            return False

class basicblock_list(object):
    """ Global basic block list class """
    blocks      = []
    count       = 0

    def __init__(self):
        pass

    def __iter__(self):
        for bb in self.blocks:
            yield bb

    def add(self, start_addr, end_addr):
        new_bb = basicblock_obj(start_addr, end_addr)
        for bb in self.blocks:
            if bb.start == new_bb.start:
                DDR_print_mesg("Basicblock already in list. Not added.",1)
                return
        self.count += 1
        DDR_print_mesg("Basicblock added to list.")
        self.blocks.append(new_bb)

    def remove(self, start_addr, end_addr):
        new_bb = basicblock_obj(start_addr, end_addr)
        for bb in self.blocks:
           if bb.start == new_bb.start:
              self.blocks.remove(bb)
              self.count -= 1
              DDR_print_mesg("Basicblock removed from list.")
              return
        DDR_print_mesg("Basicblock not found in list.",1)

    def show(self):
        DDR_print_mesg("Basicblocks list({:d} entries):".format(self.count))
        for b in self.blocks:
            DDR_print_mesg("{:x} - {:x}".format(b.start, b.end))
        return True

    def clear(self, dont_ask=False):
        if dont_ask:
            self.blocks = []
            self.count = 0
            DDR_print_mesg("Basicblock list cleared.")
            return True
        else:
            ok = ida_kernwin.ask_yn(0,"Do you really want to clear the basicblocks list ?")
            if ok == 1:
                self.blocks = []
                self.count = 0
                DDR_print_mesg("Basicblock list cleared.")
                return True
            else:
                DDR_print_mesg("Clearing basicblock list canceled.")
                return False


    def get_bblist_from_json(self, jsoncfg):
        tmpcfg = json.loads(jsoncfg)
        for n in tmpcfg["trace_start"]:
            DDR_print_mesg("n={}  trace_start={}    trace_end={}".format(n, tmpcfg["trace_start"][n],tmpcfg["trace_end"][n]),8)
            self.add(int(tmpcfg["trace_start"][n],16), int(tmpcfg["trace_end"][n],16))

BB_LIST = basicblock_list()


class ddr_api_v1_cfg(object):
    """ JSON API configuration """
    def __init__(self, BB_LIST=None): 
        self.cfg = { "apikey"                : DDR_CONFIG_SETTINGS.DDR_WEBAPI_KEY,
                     "id"                    : None,
                     "arch_bits"             : None,
                     "sample_file"           : SAMPLE_FILENAME,
                     "sample_args"           : None,
                     "sample_sha256"         : SAMPLE_SHA256,
                     "buf_size_addr"         : {},
                     "buf_size_op"           : {},
                     "buf_addr_addr"         : {},
                     "buf_addr_op"           : {},
                     "buf_dump_addr"         : {},
                     "nop_start_addr"        : {},
                     "nop_end_addr"          : {},
                     "eflag_name"            : {},
                     "eflag_addr"            : {},
                     "call_addr"             : {},
                     "call_ret"              : {},
                     "trace_light"           : {},
                     "trace_start"           : {},
                     "trace_end"             : {},
                     "trace_max_instr"       : {},
                     "trace_breakaddress"    : {},
                     "trace_startoffset"     : {}, 
                     "filelist2del"          : [],
                     "dl_file"               : [],
                     "run_opt"               : None,
                     "other"                 : {}
                    }

        self.clear_cfg()

    def set_cfg_para(self, para, value1=None, value2=None, value3=None, value4=None, value5=None, value6=None):
        
        if para.upper().startswith("0X"):
            para = para[2:]
        try:
            if value1.upper().startswith("0X"):
                value1 = value1[2:]
        except:
            pass 

        try:
            if value2.upper().startswith("0X"):
                value2 = value2[2:]
        except:
            pass

        try:
            if value3.upper().startswith("0X"):
                value3 = value3[2:]
        except:
            pass

        try:
            if value4.upper().startswith("0X"):
                value4 = value4[2:]
        except:
            pass

        try:
            if value5.upper().startswith("0X"):
                value5 = value5[2:]
        except:
            pass

        try:
            if value6.upper().startswith("0X"):
                value6 = value6[2:]
        except:
            pass

        if para == "dumpbuffer":
            if value1:
                self.cfg["buf_size_addr"].update({"{}".format(self.counter_bufdump) : "{}".format(value1.upper())})
            if value2:
                self.cfg["buf_size_op"].update({"{}".format(self.counter_bufdump) : "{}".format(value2.upper())})
            if value3:
                self.cfg["buf_addr_addr"].update({"{}".format(self.counter_bufdump) : "{}".format(value3.upper())})
            if value4:
                self.cfg["buf_addr_op"].update({"{}".format(self.counter_bufdump) : "{}".format(value4.upper())})
            if value5:
                self.cfg["buf_dump_addr"].update({"{}".format(self.counter_bufdump) : "{}".format(value5.upper())})
            self.set_buffer_counter()

        elif para == "trace":
            if value1:
                self.cfg["trace_start"].update({"{}".format(self.counter_traces) : "{}".format(value1.upper())})
            if value2:
                self.cfg["trace_end"].update({"{}".format(self.counter_traces) : "{}".format(value2.upper())})
            if value3:
                self.cfg["trace_max_instr"].update({"{}".format(self.counter_traces) : "{}".format(value3.upper())})
            if value4:
                self.cfg["trace_breakaddress"].update({"{}".format(self.counter_traces) : "{}".format(value4.upper())})
            if value5:      
                self.cfg["trace_startoffset"].update({"{}".format(self.counter_traces) : "{}".format(value5.upper())})
            if value6:
                self.cfg["trace_light"].update({"{}".format(self.counter_traces) : "{}".format(value6.upper())})
            else:
                self.cfg["trace_light"].update({"{}".format(self.counter_traces) : "{}".format("FALSE")})
            self.set_traces_counter()

        elif para == "nops":
            self.cfg["nop_start_addr"].update({"{}".format(self.counter_nops) : "{}".format(value1.upper())})
            self.cfg["nop_end_addr"].update({"{}".format(self.counter_nops) : "{}".format(value2.upper())})
            self.counter_nops += 1

        elif para == "eflags":
            self.cfg["eflag_name"].update({"{}".format(self.counter_eflags) : "{}".format(value1)})
            self.cfg["eflag_addr"].update({"{}".format(self.counter_eflags) : "{}".format(value2.upper())})
            self.counter_eflags += 1

        elif para == "calls":
            self.cfg["call_addr"].update({"{}".format(self.counter_calls) : "{}".format(value1.upper())})
            self.cfg["call_ret"].update({"{}".format(self.counter_calls) : "{}".format(value2.upper())})
            self.counter_calls += 1

        else:
            DDR_print_mesg("Unknown error setting a config value.")
            return

    def set_buffer_counter(self):
        try:
            # all values set ?
            n = str(self.counter_bufdump)
            if not self.cfg["buf_size_addr"][n] or not \
                   self.cfg["buf_size_op"][n]   or not \
                   self.cfg["buf_addr_addr"][n] or not \
                   self.cfg["buf_addr_op"][n]   or not \
                   self.cfg["buf_dump_addr"][n]:
                return
        except:
            return

        self.counter_bufdump = len(self.cfg["buf_size_addr"]) 
        if self.cfg_all_buf_set():
            DDR_print_mesg("All mandatory buffer parameters set. Number of buffers to dump: {:d}".format(self.counter_bufdump))

    def set_traces_counter(self):
        try:
            # all mandatory values set ?
            n = str(self.counter_traces)
            if not self.cfg["trace_start"][n] or not \
                   self.cfg["trace_end"][n]:
                return
        except:
            return

        self.counter_traces = len(self.cfg["trace_start"]) 
        if self.cfg_all_trace_set():
            DDR_print_mesg("All mandatory trace parameters found. Number of traces: {:d}".format(self.counter_traces))

    def clear_cfg_para(self, para):
        if para == "dumpbuffer":
            self.cfg["buf_size_addr"] = {}
            self.cfg["buf_size_op"]   = {}
            self.cfg["buf_addr_addr"] = {}
            self.cfg["buf_addr_op"]   = {}
            self.cfg["buf_dump_addr"] = {}
            self.counter_bufdump = 0
        elif para == "trace":
            self.cfg["trace_start"]        = {}
            self.cfg["trace_end"]          = {}
            self.cfg["trace_max_instr"]    = {}
            self.cfg["trace_light"]        = {}
            self.counter_traces = 0
        elif para == "nops":
            self.cfg["nop_start_addr"] = {}
            self.cfg["nop_end_addr"]   = {}
            self.counter_nops = 0
        elif para == "eflags":
            self.cfg["eflag_name"] = {}
            self.cfg["eflag_addr"] = {}
            self.counter_eflags = 0
        elif para == "calls":
            self.cfg["call_addr"] = {}
            self.cfg["call_ret"]  = {}
            self.counter_calls = 0
        else:
            DDR_print_mesg("Something went wrong clearing the config setting. This should not happen",2)
            return

    def cfg_all_trace_set(self, printout=False):
        
        # check start and end address
        for n in range(0, self.counter_traces):
           
            # check if any madatory values are hex, if not exception
            try:
                int(self.cfg["trace_start"][str(n)], 16)  
            except:
                DDR_print_mesg("'trace_start' value is not a hex value.",10, printout=True)
            try:
                int(self.cfg["trace_end"][str(n)], 16) 
            except:
                DDR_print_mesg("'trace_end' value is not a hex value.",10, printout=True)
           
        # is max instr set ?
        if self.cfg["trace_max_instr"]:
            # check if max instr is a number
            for n in range(0, self.counter_traces):
                if not all(c in "0123456789" for c in self.cfg["trace_max_instr"][str(n)]):
                    DDR_print_mesg("'trace_max_instr' value error. Should be something similar to SP0 or DP0",10, printout=True)
                    return False

        # is break addr set ?
        if self.cfg["trace_breakaddress"]:
            for n in range(0, self.counter_traces+1):
                # check if values are hex, if not exception
                try:
                    int(self.cfg["trace_breakaddress"][str(n)], 16)
                except:
                    DDR_print_mesg("'trace_breakaddress' value is not a hex value.", 10, printout=True)

        # is light trace set ? 
        if self.cfg["trace_light"]:
            # if not "TRUE" or "FALSE" aka {} empty
            for n in range(0, self.counter_traces):
                if not self.cfg["trace_light"][str(n)] == "TRUE" and \
                   not self.cfg["trace_light"][str(n)] == "FALSE":
                    DDR_print_mesg("'trace_light' value error. Should be 'TRUE' or 'FALSE' string", 10, printout=True)
                    return False
    
        return True

    def cfg_all_buf_set(self, printout=False):
                    
        for n in range(0, self.counter_bufdump):
            # check if values are hex
            try:
                int(self.cfg["buf_size_addr"][str(n)], 16) 
            except:
                DDR_print_mesg("'buf_size_addr' value is not a hex value.",10, printout=True)
                return False

            if not all(c in "SDP012345" for c in self.cfg["buf_size_op"][str(n)]):
                DDR_print_mesg("'buf_size_op' value error. Should be something similar to SP0 or DP0", 10, printout=True)
                return False

            try:
                int(self.cfg["buf_addr_addr"][str(n)], 16) 
            except:
                DDR_print_mesg("'buf_addr_addr' value is not a hex value.",10, printout=True)
                return False

            if not all(c in "SDP012345" for c in self.cfg["buf_addr_op"][str(n)]):
                DDR_print_mesg("'buf_addr_op' value error. Should be something similar to SP0 or DP0", 10, printout=True)
                return False

            try:
                int(self.cfg["buf_dump_addr"][str(n)], 16) 
            except:
                DDR_print_mesg("'buf_dump_addr' value is not a hex value.", 10, printout=True)
                return False
        
        return True
       
    def cfg_all_nop_set(self, printout=False):
        if not self.cfg["nop_start_addr"] or not self.cfg["nop_end_addr"]:
            return False
        for key in self.cfg["nop_start_addr"]:
            if not self.cfg["nop_start_addr"][key]:
                return False
        for key in self.cfg["nop_end_addr"]:
            if not self.cfg["nop_end_addr"][key]:
                return False
        return True

    def cfg_all_eflag_set(self, printout=False):
        if not self.cfg["eflag_name"] or not self.cfg["eflag_addr"]:
            return False
        for key in self.cfg["eflag_name"]:
            if not self.cfg["eflag_addr"][key]:
                return False
        for key in self.cfg["eflag_name"]:
            if not self.cfg["eflag_addr"][key]:
                return False
        return True

    def cfg_all_call_set(self, printout=False):
        if not self.cfg["call_addr"] or not self.cfg["call_ret"]:
            return False
        for key in self.cfg["call_addr"]:
            if not self.cfg["call_ret"][key]:
                return False
        for key in self.cfg["call_addr"]:
            if not self.cfg["call_ret"][key]:
                return False
        return True

    def cfg_is_empty(self):
        if not self.cfg["buf_size_addr"]    and not \
           self.cfg["buf_size_op"]          and not \
           self.cfg["buf_addr_addr"]        and not \
           self.cfg["buf_addr_op"]          and not \
           self.cfg["buf_dump_addr"]        and not \
           self.cfg["trace_start"]          and not \
           self.cfg["trace_end"]            and not \
           self.cfg["trace_max_instr"]      and not \
           self.cfg["trace_breakaddress"]   and not \
           self.cfg["trace_startoffset"]    and not \
           self.cfg["trace_light"]          and not \
           self.cfg["nop_start_addr"]       and not \
           self.cfg["nop_end_addr"]         and not \
           self.cfg["eflag_name"]           and not \
           self.cfg["eflag_addr"]           and not \
           self.cfg["call_addr"]            and not \
           self.cfg["call_ret"]:
           return True
        return False

    def cfg_bufdump_not_set(self):
        if not self.cfg["buf_size_addr"]    and not \
           self.cfg["buf_size_op"]          and not \
           self.cfg["buf_addr_addr"]        and not \
           self.cfg["buf_addr_op"]          and not \
           self.cfg["buf_dump_addr"]:
           return True
        return False

    def cfg_trace_not_set(self):
        if not self.cfg["trace_end"]        and not \
           self.cfg["trace_max_instr"]      and not \
           self.cfg["trace_breakaddress"]:
           return True
        return False

    def cfg_verified(self, printout=False, empty_cfg_not_accepted=True):   
        ret = True
        # is cfg empty:
        if self.cfg_is_empty() and empty_cfg_not_accepted:
            DDR_print_mesg("Configuration is empty.", 2)
            return False

        # if any dump buffer value is set, check dump buffer config
        if self.cfg["buf_size_addr"] or \
           self.cfg["buf_size_op"]   or \
           self.cfg["buf_addr_addr"] or \
           self.cfg["buf_addr_op"]   or \
           self.cfg["buf_dump_addr"]:
           ret = self.cfg_all_buf_set(printout)
        # if any trace value is set, check trace config
        if self.cfg["trace_start"]          or \
           self.cfg["trace_end"]            or \
           self.cfg["trace_max_instr"]      or \
           self.cfg["trace_breakaddress"]   or \
           self.cfg["trace_light"]:
           ret = self.cfg_all_trace_set(printout)
        # if any NOP value is set, check NOP config
        if self.cfg["nop_start_addr"] or \
           self.cfg["nop_end_addr"]:
           ret = self.cfg_all_nop_set(printout)
        # if any EFLAG value is set, check EFLAG config
        if self.cfg["eflag_name"]     or \
           self.cfg["eflag_addr"]:
           ret = self.cfg_all_eflag_set(printout)
        # if any CALL value is set, check CALL config
        if self.cfg["call_addr"]      or \
           self.cfg["call_ret"]:
           ret = self.cfg_all_call_set(printout)

        return ret

    def get_missing_cfg_values(self, values):
        
        missing_values = []
        if "buf_dump" in values:
            missing_values = [k for k in ['buf_size_addr', 'buf_size_op', 'buf_addr_addr', 'buf_addr_op','buf_dump_addr'] if self.cfg[k] == {}]  
        if "trace" in values:
            missing_values += [k for k in ['trace_start', 'trace_end'] if self.cfg[k] == {}] 
        if "nop" in values:
            missing_values += [k for k in ['nop_start_addr', 'nop_end_addr'] if self.cfg[k] == {}] 
        if "eflag" in values:
            missing_values += [k for k in ['eflag_name', 'eflag_addr'] if self.cfg[k] == {}] 
        if "call" in values:
            missing_values += [k for k in ['call_addr', 'call_ret'] if self.cfg[k] == {}] 
          
        return missing_values 

    def get_missing_cfg_values_msg(self, values):
        msg = {'buf_size_addr' : "Buffer size PC address",
               'buf_size_op'   : "Buffer size operand",
               'buf_addr_addr' : "Buffer address PC address",
               'buf_addr_op'   : "Buffer address operand",
               'buf_dump_addr' : "Buffer dump PC address",
               'nop_start_addr': "NOP PC start address",
               'nop_end_addr'  : "NOP PC end address",
               'eflag_name'    : "EFLAG name",
               'eflag_addr'    : "EFLAG PC address",
               'call_addr'     : "Call PC address",
               'call_ret'      : "Call return value", 
               'trace_start'   : "Trace start",
               'trace_end'     : "Trace end",
               'trace_light'   : "Trace light switch" }

        missing_values = self.get_missing_cfg_values(values)

        missing_values_msg = []
        for x in missing_values:
            missing_values_msg.append("{}{}{}{}".format(msg[x],"(",x,")"))

        return missing_values_msg

    def get_cfg_json_str_formatted(self): 
        cfg_json = json.dumps(self.cfg,indent=4, sort_keys=True)
        return cfg_json

    def print_cfg(self):
        DDR_print_mesg("{}".format(self.get_cfg_json_str_formatted()))

    def edit_cfg(self):

        for block in BB_LIST:
            DDR_print_mesg("block [0x{:x} - 0x{:x}]".format(block.start, block.end))
            self.set_cfg_para("trace", hex(block.start), hex(block.end), str(DDR_CONFIG_SETTINGS.MAX_INSTR_TO_EXECUTE)) 

        old_cfg = self.get_cfg_json_str_formatted()

        DDR_print_mesg("Config length: {:d}".format(len(old_cfg)),7)

        new_cfg = ida_kernwin.ask_text(5000, old_cfg, "Edit configuration (Expert only - Use at own risk)")

        if not new_cfg:
            DDR_print_mesg("User canceled edit configuration dialog. Nothing changed",2)
            return

        try:
            self.cfg = json.loads(new_cfg)

            try:
                BB_LIST.clear(dont_ask=True)
                BB_LIST.get_bblist_from_json(new_cfg)
            except:
                raise

            try:
                self.counter_bufdump = len(self.cfg["buf_size_addr"]) 
                DDR_print_mesg("counter_bufdump: {:d}".format(self.counter_bufdump),8)
            except:
                self.counter_bufdump = 0
                DDR_print_mesg("Except: counter_bufdump: {:d}".format(self.counter_bufdump),8)

            try:
                self.counter_nops    = len(self.cfg["nop_start_addr"]) 
                DDR_print_mesg("counter_nops: {:d}".format(self.counter_nops))
            except:
                self.counter_nops    = 0
                DDR_print_mesg("Except: counter_nops: {:d}".format(self.counter_nops),8)

            try:
                self.counter_eflags  = len(self.cfg["eflag_name"]) 
                DDR_print_mesg("counter_eflags: {:d}".format(self.counter_eflags),8)
            except:
                self.counter_eflags  = 0
                DDR_print_mesg("Except:counter_eflags: {:d}".format(self.counter_eflags),8)

            try:
                self.counter_calls   = len(self.cfg["call_addr"]) 
                DDR_print_mesg("counter_calls: {:d}".format(self.counter_calls),8)
            except:
                self.counter_calls   = 0
                DDR_print_mesg("Except: counter_calls: {:d}".format(self.counter_calls),8)

            try:
                self.counter_traces  = len(self.cfg["trace_start"]) 
                DDR_print_mesg("counter_traces: {:d}".format(self.counter_traces),8)
            except:
                self.counter_traces  = 0
                DDR_print_mesg("Except: counter_traces: {:d}".format(self.counter_traces),8)

            if not self.cfg_verified(printout=True):
                self.cfg = json.loads(old_cfg)
                DDR_print_mesg("Config check failed. Old config restored.",2)
            else:
                DDR_print_mesg("Changed configuration:")
                self.print_cfg()

        except ValueError as e:
            DDR_print_mesg("Failed to parse JSON config",2)
            DDR_print_mesg("{}".format(e),2)

    def clear_cfg(self):
        self.cfg = { "apikey"             : DDR_CONFIG_SETTINGS.DDR_WEBAPI_KEY,
                     "id"                 : None,
                     "arch_bits"          : None,
                     "sample_file"        : SAMPLE_FILENAME,
                     "sample_args"        : None,
                     "sample_sha256"      : SAMPLE_SHA256,
                     "buf_size_addr"      : {},
                     "buf_size_op"        : {},
                     "buf_addr_addr"      : {},
                     "buf_addr_op"        : {},
                     "buf_dump_addr"      : {},
                     "nop_start_addr"     : {},
                     "nop_end_addr"       : {},
                     "eflag_name"         : {},
                     "eflag_addr"         : {},
                     "call_addr"          : {},
                     "call_ret"           : {},
                     "trace_light"        : {},
                     "trace_start"        : {},
                     "trace_end"          : {},
                     "trace_max_instr"    : {},
                     "trace_breakaddress" : {},
                     "trace_startoffset"  : {},
                     "run_opt"            : None,
                     "other"              : {}
                    }  

        self.counter_nops    = 0
        self.counter_eflags  = 0
        self.counter_calls   = 0
        self.counter_traces  = 0
        self.counter_bufdump = 0

        BB_LIST.clear(dont_ask=True)

        info = idaapi.get_inf_structure()
        if info.is_64bit():
            self.cfg["arch_bits"] = 64
        elif info.is_32bit():
            self.cfg["arch_bits"] = 32

        DDR_print_mesg("Configuration cleared.")        

class DDR_ida_action_handler(idaapi.action_handler_t):
    """
    DDR Action handler for all menu entries, gets created when user picks menu 
    """

    def __init__(self, usrcmd):
        idaapi.action_handler_t.__init__(self)
        self.cmd = usrcmd
        DDR_print_mesg("cmd = {}".format(self.cmd), 7)

    def update(self, ctx):
        """
        Ensure the context menu is always available in IDA.
        """
        DDR_print_mesg("[{}] updated.".format(self.cmd), 7)
        return idaapi.AST_ENABLE_ALWAYS

    def activate(self, ctx):
        """
        Execute the embedded action_function when this context menu is invoked.
        """
        DDR_print_mesg("Command selected by user: [{}]".format(self.cmd), 7)
        
        global JSONDATA
        global JSONFILE_LOADED
        global APIFILE_LOADED
        global BB_LIST
        global EFLAGS
      
        ea = idc.get_screen_ea()

        pe = pefile.PE(idaapi.get_input_file_path(), fast_load=True)
        eop = pe.OPTIONAL_HEADER.AddressOfEntryPoint
        eop = eop + pe.OPTIONAL_HEADER.ImageBase
        

        if eop != ida_ida.inf_get_start_ea():
            DDR_print_mesg("Org. Entrypoint of PE on disk and IDA Entrypoint are not the same.",2)
            DDR_print_mesg("OEP PEFILE: {}".format(hex(eop)),2)
            DDR_print_mesg("OEP IDA   : {}".format(hex(ida_ida.inf_get_start_ea())),2)
            DDR_print_mesg("DDR is using the absolute virtual addresses coming from IDA for its operations, e.g. 0x410010.",2)
            DDR_print_mesg("In case of relocation this will become e.g. 0x420010, but the ddr_server needs the org. virtual address from the PE on disk e.g. 0x410010.",2)
            DDR_print_mesg("This will not work in most cases as far as the DDR server will use the wrong address (0x420010) for its operations.",2)
            DDR_print_mesg("This usually happens when you manually rebase the segment or you have used the IDA debugger with a file using relocation",2)
            DDR_print_mesg("For security reasons we canceled the operation.",2)
            idaapi.warning("Warning: Operation canceled. IDA PE entrypoint and PE file on disk entrypoint are not the same. See log window for more infos.")
            return 1

        # Command handler
        if self.cmd == ("Select_Add_BB2List"):
            DDR_print_mesg("Adding basic block to basic blocks list",7)
            if self._Add_BB2BBlist(ea):
                DDR_print_mesg("Done. Basic block added to basic block list.")
                DDR_print_mesg("-------------------------------------------------------------------------------")
            else:
                DDR_print_mesg("Done. Failed to add basic block to basic block list.")
                DDR_print_mesg("-------------------------------------------------------------------------------")
            return 1

        if self.cmd == ("Select_Remove_BB2List"):
            DDR_print_mesg("Removing basic block to basic blocks list",7)
            if self._Remove_BB2BBlist(ea):
                DDR_print_mesg("Done. Basic block removed from basic block list.")
                DDR_print_mesg("-------------------------------------------------------------------------------")
            else:
                DDR_print_mesg("Done. Failed to remove basic block from basic block list.")
                DDR_print_mesg("-------------------------------------------------------------------------------")
            return 1

        if self.cmd == ("Select_Print_BB2List"):
            DDR_print_mesg("Show basic blocks list",7)
            if BB_LIST.show():
                DDR_print_mesg("Done. Basic block list printed.")
                DDR_print_mesg("-------------------------------------------------------------------------------")
            else:
                DDR_print_mesg("Done. Failed to print basic block list.")
                DDR_print_mesg("-------------------------------------------------------------------------------")
            return 1

        if self.cmd == ("Select_Clear_BB2List"):
            DDR_print_mesg("Show basic blocks list",7)
            if BB_LIST.clear():
                DDR_print_mesg("Done. Basic block list cleared.")
                DDR_print_mesg("-------------------------------------------------------------------------------")
            else:
                DDR_print_mesg("Done. Failed to clear basic block list.")
                DDR_print_mesg("-------------------------------------------------------------------------------")
            return 1

        if self.cmd == ("Select_SetStartOffset"):
            startoffstr = hex(ea)
            startoffstr = startoffstr[2:]
            # There will likely always be only one start offset for all trace basicblock list entries (multiple traces),
            # but we keep the dictionary format for future use of the field
            DDR_API_V1_CONFIG_SETTINGS.cfg["trace_startoffset"]["0"] = startoffstr.upper()
            DDR_print_mesg("Start offset set to 0x{}".format(DDR_API_V1_CONFIG_SETTINGS.cfg["trace_startoffset"]["0"]))
            return 1

        if self.cmd == ("Select_ClearStartOffset"):
            DDR_API_V1_CONFIG_SETTINGS.cfg["trace_startoffset"] = {}
            DDR_print_mesg("Start offset cleared")
            return 1

        if self.cmd == ("Select_ShowStartOffset"):
            try:
                DDR_print_mesg("Start offset: 0x{}".format(DDR_API_V1_CONFIG_SETTINGS.cfg["trace_startoffset"]["0"]))
            except:
                DDR_print_mesg("Start offset: not set")
            return 1

        if self.cmd == ("Select_SetBreakaddress"):
            startoffstr = hex(ea)
            startoffstr = startoffstr[2:]
            # There will likely always be only one breakaddress for all trace basicblock list entries (multiple traces),
            # but we keep the dictionary format for future use of the field
            DDR_API_V1_CONFIG_SETTINGS.cfg["trace_breakaddress"]["0"] = startoffstr.upper()
            DDR_print_mesg("Break address set to 0x{}".format(DDR_API_V1_CONFIG_SETTINGS.cfg["trace_breakaddress"]["0"]))
            return 1

        if self.cmd == ("Select_ClearBreakaddress"):
            DDR_API_V1_CONFIG_SETTINGS.cfg["trace_breakaddress"] = {}
            DDR_print_mesg("Break address cleared")
            return 1

        if self.cmd == ("Select_ShowBreakaddress"):
            try:
                DDR_print_mesg("Break address: 0x{}".format(DDR_API_V1_CONFIG_SETTINGS.cfg["trace_breakaddress"]["0"]))
            except:
                DDR_print_mesg("Break address: not set")
            return 1

        if self.cmd == ("Select_SetCmdArgs"):
            
            if DDR_API_V1_CONFIG_SETTINGS.cfg["sample_args"]:
                cmdarg_txt = DDR_API_V1_CONFIG_SETTINGS.cfg["sample_args"]
            else:
                cmdarg_txt = ""

            sample_args = ida_kernwin.ask_text(5000, cmdarg_txt, "Set command line arguments for sample e.g. -a 1 -b \"test string\"")
            DDR_API_V1_CONFIG_SETTINGS.cfg["sample_args"] = sample_args
            DDR_print_mesg("Command line args for sample set to: {}".format(sample_args))
            return 1

        if self.cmd == ("Select_ClearCmdArgs"):
            DDR_API_V1_CONFIG_SETTINGS.cfg["sample_args"] = None
            DDR_print_mesg("Command line args cleared")
            return 1

        if self.cmd == ("Select_ShowCmdArgs"):
            if DDR_API_V1_CONFIG_SETTINGS.cfg["sample_args"]:
                DDR_print_mesg("Command line args: {}".format(DDR_API_V1_CONFIG_SETTINGS.cfg["sample_args"]))
            else:
                DDR_print_mesg("Command line args: not set")
            return 1


        if self.cmd == ("Dump_buffer_size_op"):
            try:
                op_str = ida_kernwin.get_highlight(idaapi.get_current_viewer())[0]
            except:
                idaapi.warning("No operand highlighted. Please highlight the operand you want to use for finding the buffer size")
                DDR_print_mesg("No operand highlighted. Please highlight the operand you want to use for finding the buffer size",2)
                return 1

            dump_size_addr, dump_size_op = self._get_dump_para_addr_n_op(op_str, ea)
            if (dump_size_addr and dump_size_op):
                DDR_print_mesg("Using operand {} at PC address 0x{:x} to find size of buffer to dump".format(dump_size_op, dump_size_addr))
                DDR_API_V1_CONFIG_SETTINGS.set_cfg_para("dumpbuffer", value1=hex(dump_size_addr))
                DDR_API_V1_CONFIG_SETTINGS.set_cfg_para("dumpbuffer", value2=dump_size_op)
            return 1

        if self.cmd == ("Dump_buffer_addr_op"):
            try:
                op_str = ida_kernwin.get_highlight(idaapi.get_current_viewer())[0]
            except:
                idaapi.warning("No operand highlighted. Please highlight the operand you want to use for finding the buffer address")
                DDR_print_mesg("No operand highlighted. Please highlight the operand you want to use for finding the buffer address",2)
                return 1

            dump_addr_addr, dump_addr_op = self._get_dump_para_addr_n_op(op_str, ea)
            try:
                DDR_print_mesg("Using operand {} at PC address 0x{:x} to find memory address of buffer to dump".format(dump_addr_op, dump_addr_addr))
                DDR_API_V1_CONFIG_SETTINGS.set_cfg_para("dumpbuffer", value3=hex(dump_addr_addr))
                DDR_API_V1_CONFIG_SETTINGS.set_cfg_para("dumpbuffer", value4=dump_addr_op)
            except:
                DDR_print_mesg("No config parameters set.",2)
            return 1

        if self.cmd == ("Dump_buffer_dump_addr"):
            DDR_print_mesg("Using PC address 0x{:x} to finally dump the buffer to file".format(ea))
            DDR_API_V1_CONFIG_SETTINGS.set_cfg_para("dumpbuffer", value5=hex(ea))
            return 1

        if self.cmd == ("Dump_buffer_edit_cfg"):
            DDR_print_mesg("Manually edit config",7)
            try:
                DDR_API_V1_CONFIG_SETTINGS.edit_cfg();
                DDR_API_V1_CONFIG_SETTINGS.clear_cfg_para("trace") # clean the trace cfg, to load BB_LIST at execution
            except:
                DDR_print_mesg("Edit configuration failed.",2)
                raise
            return 1

        if self.cmd == ("Dump_buffer_clear_cfg"):
            DDR_print_mesg("Clear config",7)
            try:
                ok = ida_kernwin.ask_yn(0,"Do you really want to clear the configuration ?")
                if ok == 1:
                    DDR_API_V1_CONFIG_SETTINGS.clear_cfg();
                else:
                    DDR_print_mesg("Clearing configuration canceled.")
            except:
                DDR_print_mesg("Clear configuration failed.",2)
                raise
            return 1

        if self.cmd == ("Dump_buffer"):
            DDR_print_mesg("Executing sample and trying to dump the buffer",7)
            # --------------------
            if not DDR_API_V1_CONFIG_SETTINGS.cfg_verified():
                idaapi.warning("Dump config verification failed. Please first set the buffer size, buffer address and dump address via the 'Dump' menu.")
                DDR_print_mesg("Dump configuration verification failed.",2)
                DDR_print_mesg("Missing mandatory values: {}".format(DDR_API_V1_CONFIG_SETTINGS.get_missing_cfg_values_msg("buf_dump")),2)
                DDR_print_mesg("Please first set the buffer size, buffer address and dump address via the 'Dump' menu.",2)
            else:
                # We can either dump a buffer or trace the file itm.
                DDR_API_V1_CONFIG_SETTINGS.clear_cfg_para("trace") # clear trace config, we have it still saved in the BB list

                DDR_print_mesg("Sending following config to server:")
                DDR_API_V1_CONFIG_SETTINGS.print_cfg()

                if not self._call_api(cmd_id=5, dump_para_json = DDR_API_V1_CONFIG_SETTINGS.cfg):
                    DDR_print_mesg("Problem detected while dumping buffer. Check log above.",2)
                else:
                    DDR_print_mesg("Done. Buffer dumped to file.")
                    DDR_print_mesg("-------------------------------------------------------------------------------")
            return 1

        if self.cmd == ("Dump_buffer_nop_instr"):  
            try:
                start = idc.read_selection_start()
                end   = idc.prev_head(idc.read_selection_end())

                # User selected multiple lines
                if start != idc.BADADDR and end != idc.BADADDR:
                    DDR_print_mesg("NOP'ing out instructions from instruction 0x{:x} to 0x{:x}".format(start, end))
                    DDR_API_V1_CONFIG_SETTINGS.set_cfg_para("nops", hex(start), hex(end))
                # User selected single line
                else: 
                    DDR_print_mesg("NOP'ing out instruction at 0x{:x}".format(ea))
                    DDR_API_V1_CONFIG_SETTINGS.set_cfg_para("nops", hex(ea), hex(ea))
                return 1
            except:
                raise

        if self.cmd == ("Dump_buffer_eflag_instr"):  
            try:
                eflag_str = ida_kernwin.ask_text(2, "ZF", "Please enter EFLAG name:").upper()
            except:
                DDR_print_mesg("EFLAG dialog canceled.")
                return 1

            DDR_print_mesg("EFLAG: {} entered.".format(eflag_str))

            if not eflag_str in EFLAGS:
                DDR_print_mesg("EFLAG is not a valid EFLAG. Please use one from this list: {}".format(",".join(EFLAGS)))
                idaapi.warning("EFLAG is not a valid EFLAG. Please use one from this list: {}".format(",".join(EFLAGS)))
                return 1

            try:
                start = idc.read_selection_start()
                end   = idc.read_selection_end()

                # User selected multiple lines
                if start != idc.BADADDR and end != idc.BADADDR:
                    ea=start
                    while(ea < end):  
                         DDR_print_mesg("Toggle EFLAG {} at 0x{:x}".format(eflag_str, ea)) 
                         DDR_API_V1_CONFIG_SETTINGS.set_cfg_para("eflags", eflag_str, hex(ea))         
                         ea = idc.next_head(ea)
                # User selected single line
                else: 
                    DDR_print_mesg("Toggle EFLAG {} at line 0x{:x}".format(eflag_str, ea))
                    DDR_API_V1_CONFIG_SETTINGS.set_cfg_para("eflags", eflag_str, hex(ea)) 
                return 1
            except:
                raise

        if self.cmd == ("Dump_buffer_skip_call"):  

            if idaapi.get_inf_structure().is_64bit():
                max_size = 64 + 2
            else:
                max_size = 32 + 2

            try:
                start = idc.read_selection_start()
                end   = idc.read_selection_end()

                # User selected multiple lines
                if start != idc.BADADDR and end != idc.BADADDR:
                    idaapi.warning("Selecting a range for function skipping is not supported. Pls select a single line.")
                    DDR_print_mesg("Selecting a range for function skipping is not supported. Pls select a single line.",2)
                # User selected single line
                else: 
                    call_retval = ida_kernwin.ask_text(max_size, "D0D0", "Please enter return value (in hex) for skipped function:")
                    try:
                        if call_retval:
                            r = int(call_retval,16)
                        else:
                            DDR_print_mesg("Return value dialog canceled.")
                            return 1
                    except:
                        DDR_print_mesg("Return value is not a hex value. You need to enter a hex value. Skip function operation not executed.",1)
                        idaapi.warning("Return value is not a hex value. Skip function operation not executed.")
                        return 1

                    DDR_print_mesg("Skipping function at 0x{:x} at runtime. Return value set to: {}".format(ea, call_retval))
                    DDR_API_V1_CONFIG_SETTINGS.set_cfg_para("calls", hex(ea), call_retval)
                return 1
            except:
                raise

        if self.cmd == ("Sample_to_x64dbg"):

            DDR_print_mesg("Creating x64dbg script.")

            DDR_API_V1_CONFIG_SETTINGS.cfg["other"].update({"breakaddr" : hex(ea)})
            DDR_API_V1_CONFIG_SETTINGS.cfg["other"].update({"imagebase" : hex(idaapi.get_imagebase())})  

            try:
                # Fixing nop_end_addr + 1 
                org_nop_end_addr = DDR_API_V1_CONFIG_SETTINGS.cfg["nop_end_addr"]['0']
                new_nop_end_addr = idc.next_head(int(org_nop_end_addr,16))
                DDR_API_V1_CONFIG_SETTINGS.cfg["nop_end_addr"]['0'] = hex(new_nop_end_addr)
            except:
                pass

            DDR_API_V1_CONFIG_SETTINGS.print_cfg()

            if not self._call_api(cmd_id=12, dump_para_json = DDR_API_V1_CONFIG_SETTINGS.cfg):
                    DDR_print_mesg("Problem detected while patching sample. Check log above.",2)
            else:
                DDR_print_mesg("Done. x64dbg Script created.")
                DDR_print_mesg("-------------------------------------------------------------------------------")

            try:
                DDR_API_V1_CONFIG_SETTINGS.cfg["nop_end_addr"]['0'] = org_nop_end_addr
            except:
                pass

            return 1


        if self.cmd == ("Sample_loop_addr"):
            offset_on_disk = idaapi.get_fileregion_offset(ea)
            DDR_print_mesg("Create binary with and endless loop at: 0x{:x}".format(ea))
            DDR_print_mesg("Address offset on disk                : 0x{:x}".format(offset_on_disk))

            DDR_API_V1_CONFIG_SETTINGS.cfg["other"].update({"offset_disk" : offset_on_disk})
            DDR_API_V1_CONFIG_SETTINGS.cfg["other"].update({"loop_addr"   : ea})

            if not self._call_api(cmd_id=11, dump_para_json = DDR_API_V1_CONFIG_SETTINGS.cfg):
                    DDR_print_mesg("Problem detected while patching sample with endless loop. Check log above.",2)
            else:
                DDR_print_mesg("Done. Binary with endless loop created.")
                DDR_print_mesg("-------------------------------------------------------------------------------")

            return 1

        if self.cmd == ("Sample_run_only"):

            if DDR_API_V1_CONFIG_SETTINGS.cfg_verified(empty_cfg_not_accepted=False):
                DDR_print_mesg("Executing sample on server in DynamoRio incl. all patches.")
                old_cfg = DDR_API_V1_CONFIG_SETTINGS.get_cfg_json_str_formatted()
                DDR_API_V1_CONFIG_SETTINGS.clear_cfg_para("dumpbuffer") 
                DDR_API_V1_CONFIG_SETTINGS.clear_cfg_para("trace")
                DDR_API_V1_CONFIG_SETTINGS.cfg["run_opt"] = "RUN_ONLY"

                DDR_print_mesg("Sending following config to server:")
                DDR_API_V1_CONFIG_SETTINGS.print_cfg()

                if not self._call_api(cmd_id=10, dump_para_json = DDR_API_V1_CONFIG_SETTINGS.cfg):
                    DDR_print_mesg("Problem detected while running sample on server. Check log above.",2)
                else:
                    DDR_print_mesg("Done. Sample executed.")
                    DDR_print_mesg("-------------------------------------------------------------------------------")

                DDR_API_V1_CONFIG_SETTINGS.clear_cfg_para("trace")
                DDR_API_V1_CONFIG_SETTINGS.cfg = json.loads(old_cfg)
                DDR_API_V1_CONFIG_SETTINGS.cfg["run_opt"] = None
            else:
                idaapi.warning("Config verification failed.")
                DDR_print_mesg("Config verification failed.",2)

            return 1

        if self.cmd == ("Delete_Cached_Files"):
            jsonfile_name =  JSONFILE_DIR + "\\" + "DDR_log_" + '.'.join(idaapi.get_root_filename().split('.')[:-1]) + "*"
            for filename_to_delete in glob.glob(jsonfile_name):
                DDR_print_mesg("Deleting file(s): {}".format(filename_to_delete))
                os.remove(filename_to_delete)
            return 1

        if self.cmd == ("Run_Light_Trace_On_Seg"):
            b_start        = idc.get_segm_start(ea)
            b_end          = idc.get_segm_end(ea)

            if b_start != idc.BADADDR and b_end != idc.BADADDR:
                old_cfg = DDR_API_V1_CONFIG_SETTINGS.get_cfg_json_str_formatted()
                DDR_API_V1_CONFIG_SETTINGS.clear_cfg_para("dumpbuffer") 
                DDR_API_V1_CONFIG_SETTINGS.clear_cfg_para("trace")

                DDR_print_mesg("Running trace on segment: [0x{:x} - 0x{:x}]".format(b_start, b_end))
                DDR_API_V1_CONFIG_SETTINGS.set_cfg_para("trace", hex(b_start), hex(b_end), str(DDR_CONFIG_SETTINGS.MAX_INSTR_TO_EXECUTE), value6="TRUE")


                DDR_print_mesg("Sending following config to server:")
                DDR_API_V1_CONFIG_SETTINGS.print_cfg()

                if not self._call_api(cmd_id=7, dump_para_json = DDR_API_V1_CONFIG_SETTINGS.cfg):
                    DDR_print_mesg("Failed to run light trace for segment",2)
                    idaapi.warning("Failed to run light trace for segment.")
                else:
                    DDR_print_mesg("Done. Light trace on segment executed.")
                    DDR_print_mesg("-------------------------------------------------------------------------------")

                DDR_API_V1_CONFIG_SETTINGS.clear_cfg_para("trace")
                DDR_API_V1_CONFIG_SETTINGS.cfg = json.loads(old_cfg)

            else:
                DDR_print_mesg("Failed finding segment.")
           
            # old: DDR_print_mesg("Running instruction address trace only for selected segment range: {:x} - {:x}.".format(start, end))
            # old: self._exec_dynRio(start_addr=start, end_addr=end, instr_count=DDR_CONFIG_SETTINGS.MAX_INSTR_TO_EXECUTE, options='light_trace_only')
            return 1

        if self.cmd == ("Run_Trace_On_Seg"):
            b_start        = idc.get_segm_start(ea)
            b_end          = idc.get_segm_end(ea)
            
            if b_start != idc.BADADDR and b_end != idc.BADADDR:
                old_cfg = DDR_API_V1_CONFIG_SETTINGS.get_cfg_json_str_formatted()
                DDR_API_V1_CONFIG_SETTINGS.clear_cfg_para("dumpbuffer") 
                DDR_API_V1_CONFIG_SETTINGS.clear_cfg_para("trace")

                DDR_print_mesg("Running trace on segment: [0x{:x} - 0x{:x}]".format(b_start, b_end))
                DDR_API_V1_CONFIG_SETTINGS.set_cfg_para("trace", hex(b_start), hex(b_end), str(DDR_CONFIG_SETTINGS.MAX_INSTR_TO_EXECUTE))

                DDR_print_mesg("Sending following config to server:")
                DDR_API_V1_CONFIG_SETTINGS.print_cfg()

                if not self._call_api(cmd_id=7, dump_para_json = DDR_API_V1_CONFIG_SETTINGS.cfg):
                    DDR_print_mesg("Failed to run full trace for segment",2)
                    idaapi.warning("Failed to run full trace for segment.")
                else:
                    DDR_print_mesg("Done. Full trace on segment executed.")
                    DDR_print_mesg("-------------------------------------------------------------------------------")

                DDR_API_V1_CONFIG_SETTINGS.clear_cfg_para("trace")
                DDR_API_V1_CONFIG_SETTINGS.cfg = json.loads(old_cfg)

            else:
                DDR_print_mesg("Failed finding segment.")

            return 1

        if self.cmd == ("Run_Trace_On_Range"):
            b_start = idc.read_selection_start()
            b_end   = idc.prev_head(idc.read_selection_end())

            if b_start != idc.BADADDR and b_end != idc.BADADDR:

                old_cfg = DDR_API_V1_CONFIG_SETTINGS.get_cfg_json_str_formatted()
                DDR_API_V1_CONFIG_SETTINGS.clear_cfg_para("dumpbuffer") 
                DDR_API_V1_CONFIG_SETTINGS.clear_cfg_para("trace")

                DDR_print_mesg("Running trace on address range: [0x{:x} - 0x{:x}]".format(b_start, b_end))
                DDR_API_V1_CONFIG_SETTINGS.set_cfg_para("trace", hex(b_start), hex(b_end), str(DDR_CONFIG_SETTINGS.MAX_INSTR_TO_EXECUTE))

                DDR_print_mesg("Sending following config to server:")
                DDR_API_V1_CONFIG_SETTINGS.print_cfg()

                if not self._call_api(cmd_id=7, dump_para_json = DDR_API_V1_CONFIG_SETTINGS.cfg):
                    DDR_print_mesg("Failed to run full trace for selected range.",2)
                    idaapi.warning("Failed to run full trace for selected range.")
                DDR_print_mesg("Done. Full trace on address range executed.")
                DDR_print_mesg("-------------------------------------------------------------------------------")

                DDR_API_V1_CONFIG_SETTINGS.clear_cfg_para("trace")
                DDR_API_V1_CONFIG_SETTINGS.cfg = json.loads(old_cfg)

            else:
                DDR_print_mesg("Failed finding address range.")
            return 1

        if self.cmd == ("Run_Trace_On_BB"):
            DDR_print_mesg("Running trace for selected basic block.")
    
            f = idaapi.get_func(ea)

            if not f:
                DDR_print_mesg("No function found at 0x{:x}. This only works inside of functions.".format(ea))
                idaapi.warning("No function found at 0x{:x}. This only works inside of functions.".format(ea))
                return

            fc = idaapi.FlowChart(f)
            for block in fc:
                if block.start_ea <= ea:
                    if block.end_ea > ea:
                        b_start = block.start_ea
                        b_end   = idc.prev_head(block.end_ea) 

            if b_start and b_end and b_start != idc.BADADDR and b_end != idc.BADADDR:
                old_cfg = DDR_API_V1_CONFIG_SETTINGS.get_cfg_json_str_formatted()
                DDR_API_V1_CONFIG_SETTINGS.clear_cfg_para("dumpbuffer") 
                DDR_API_V1_CONFIG_SETTINGS.clear_cfg_para("trace")

                DDR_print_mesg("Running trace on basic block [0x{:x} - 0x{:x}]".format(b_start, b_end))
                DDR_API_V1_CONFIG_SETTINGS.set_cfg_para("trace", hex(b_start), hex(b_end), str(DDR_CONFIG_SETTINGS.MAX_INSTR_TO_EXECUTE)) 

                DDR_print_mesg("Sending following config to server:")
                DDR_API_V1_CONFIG_SETTINGS.print_cfg()  

                if not self._call_api(cmd_id=7, dump_para_json = DDR_API_V1_CONFIG_SETTINGS.cfg):
                    DDR_print_mesg("Failed to run full trace for selected basic block.",2)
                    idaapi.warning("Failed to run full trace for selected basic block.")
                else:
                    DDR_print_mesg("Done. Full trace on basic block executed.")
                    DDR_print_mesg("-------------------------------------------------------------------------------")

                DDR_API_V1_CONFIG_SETTINGS.clear_cfg_para("trace")
                DDR_API_V1_CONFIG_SETTINGS.cfg = json.loads(old_cfg)
            else:
               DDR_print_mesg("Basic block start and end address not found.",2) 

            return 1

        if self.cmd == ("Run_Trace_On_BB_list"):
            DDR_print_mesg("Running trace for selected basic block list:")

            # We can either dump a buffer or trace the file itm. 
            old_cfg = DDR_API_V1_CONFIG_SETTINGS.get_cfg_json_str_formatted()

            for block in BB_LIST:
                DDR_print_mesg("Block [0x{:x} - 0x{:x}]".format(block.start, block.end))
                DDR_API_V1_CONFIG_SETTINGS.set_cfg_para("trace", hex(block.start), hex(block.end), str(DDR_CONFIG_SETTINGS.MAX_INSTR_TO_EXECUTE)) 

            if not DDR_API_V1_CONFIG_SETTINGS.cfg_verified():
                idaapi.warning("Trace config verification failed. Please first select basic blocks via 'Select' menu.")
                DDR_print_mesg("Missing mandatory values: {}".format(DDR_API_V1_CONFIG_SETTINGS.get_missing_cfg_values_msg("trace")),2)
                DDR_print_mesg("Trace configuration verification failed. Please first select basic blocks via 'Select' menu.",2)        
            else:
                DDR_API_V1_CONFIG_SETTINGS.cfg["id"] = 7
                # We can either dump a buffer or trace the file itm. 
                DDR_API_V1_CONFIG_SETTINGS.clear_cfg_para("dumpbuffer") 

                DDR_print_mesg("Sending following config to server:")
                DDR_API_V1_CONFIG_SETTINGS.print_cfg()
                if not self._call_api(cmd_id=7, dump_para_json = DDR_API_V1_CONFIG_SETTINGS.cfg):
                    DDR_print_mesg("Failed to run full trace for basic block list.",2)
                    idaapi.warning("Failed to run full trace for basic block list.")
                else:
                    DDR_print_mesg("Done. Full trace on basic block list executed.")
                    DDR_print_mesg("-------------------------------------------------------------------------------")

                DDR_API_V1_CONFIG_SETTINGS.cfg = json.loads(old_cfg)

            DDR_API_V1_CONFIG_SETTINGS.clear_cfg_para("trace") # clear trace config, we have it still saved in the BB list

            return 1

        if self.cmd == ("Highlight_trace"): 
            if(JSONFILE_LOADED == False):
                self._trace_file_not_loaded()
                return 1
            else:
                self._highlight_trace()
                DDR_print_mesg("Done. Instructions highlighted.")
                DDR_print_mesg("-------------------------------------------------------------------------------")
                return 1

        if self.cmd == ("Clear_highlighted_blocks"):
                self._clear_highlight_trace(ea)
                DDR_print_mesg("Done. Highlighted instructions cleared.")
                DDR_print_mesg("-------------------------------------------------------------------------------")
                return 1

        if self.cmd == ("Get_Set_Num_Hits_Cmt"):
            cmt_rnds = ida_kernwin.ask_long(DDR_CONFIG_SETTINGS.MAX_CMT_ROUNDS, "Please enter max. number of hits to display in IDA comments (max. {:d}):".format(DDR_CONFIG_SETTINGS.MAX_INSTR_COUNT))
            try:
                if cmt_rnds > 0 and cmt_rnds <= DDR_CONFIG_SETTINGS.MAX_INSTR_COUNT:
                    DDR_CONFIG_SETTINGS.MAX_CMT_ROUNDS = cmt_rnds
                    DDR_print_mesg("Number of hits for repeatable comments {:d} set".format(DDR_CONFIG_SETTINGS.MAX_CMT_ROUNDS)) 
                return 1
            except:
                DDR_print_mesg("Setting value canceled. Number of hits for repeatble comments {:d}".format(DDR_CONFIG_SETTINGS.MAX_CMT_ROUNDS), 1) 
                return 1

        if self.cmd == ("Get_Set_Num_Hits_IdaLog"):
            try:
                cmt_rnds = ida_kernwin.ask_long(DDR_CONFIG_SETTINGS.MAX_LOG_ROUNDS, "Please enter max. number of hits to display in IDA comments (max. {:d}):".format(DDR_CONFIG_SETTINGS.MAX_INSTR_COUNT))
                if cmt_rnds > 0 and cmt_rnds <= DDR_CONFIG_SETTINGS.MAX_INSTR_COUNT:
                    DDR_CONFIG_SETTINGS.MAX_LOG_ROUNDS = cmt_rnds
                    DDR_print_mesg("Number of hits for IDA log window {:d} set".format(DDR_CONFIG_SETTINGS.MAX_LOG_ROUNDS)) 
                return 1
            except:
                DDR_print_mesg("Setting value canceled. Number of hits for IDA log window {:d}".format(DDR_CONFIG_SETTINGS.MAX_LOG_ROUNDS), 1) 
                return 1


        if self.cmd == ("Get_Set_Max_Instr"):
            try:
                cmt_rnds = ida_kernwin.ask_long(DDR_CONFIG_SETTINGS.MAX_INSTR_TO_EXECUTE, "Please enter max. number of instructions to process in DynamoRio analysis (Default: 20000):")
                if cmt_rnds > 0:
                    if cmt_rnds > 40000:
                        idaapi.warning("Setting max. instructions to greater than 40000 might be slow. Only use it for 'Run Light Trace on Segment' feature " + 
                            "or make sure you do not run into API timeouts (DDR_CONFIG_SETTINGS.MAX_API_TIMEOUT) because the analysis takes too long on the server side.")
                    DDR_CONFIG_SETTINGS.MAX_INSTR_TO_EXECUTE = cmt_rnds
                    DDR_print_mesg("Max. Number of instructions to process in DynamoRio analysis set to: {:d}".format(DDR_CONFIG_SETTINGS.MAX_INSTR_TO_EXECUTE))
                return 1
            except:
                DDR_print_mesg("Setting value canceled. Max. Number of instructions to process in DynamoRio analysis: {:d}".format(DDR_CONFIG_SETTINGS.MAX_INSTR_TO_EXECUTE), 1)
                return 1

        if self.cmd == ("Get_Set_API_timeout"):
            try:
                cmt_rnds = ida_kernwin.ask_text(20, str(DDR_CONFIG_SETTINGS.MAX_API_TIMEOUT), "Please enter max. number of seconds to wait for the DDR server API to answer (Default: 30):")
                DDR_CONFIG_SETTINGS.MAX_API_TIMEOUT = float(cmt_rnds)
                DDR_print_mesg("Max. number of seconds to wait for a response from the DDR server set to: {:f} seconds".format(DDR_CONFIG_SETTINGS.MAX_API_TIMEOUT))
                return 1
            except:
                DDR_print_mesg("Setting value canceled. Max. Number of to wait for a response from the DDR server: {:f}".format(DDR_CONFIG_SETTINGS.MAX_API_TIMEOUT), 1)
                return 1

        if self.cmd == ("Display_Strings_View"):
            if(JSONFILE_LOADED == False):
                self._trace_file_not_loaded()
                return 1    
            else:
                items = self._get_chooser_row_list_strings()
                DDR_print_mesg("Len items: {:d}".format(len(items)),7)
                columns = [ ["Instruction #", 6], ["Address", 12], ["Type", 15], ["Data Hex", 25], ["Data Address", 12], ["String", 12], ["Disasm", 35] ]
                c = DDR_MyChoose("Strings Trace", items=items, modal=False, columns=columns)
                if not c.show():
                    DDR_print_mesg("Failed to start Strings Trace View", 7)
                else:
                    DDR_print_mesg("Strings Trace View created", 7)
                return 1 

        if self.cmd == ("Display_Calls_View"):
            if APIFILE_LOADED == False:
                DDR_print_mesg("API trace file not loaded. Pls use 'Trace' menu to run a trace first. To generate a trace and API trace file.", 1)
                DDR_print_mesg("Alternativly, you can use 'File/Load File/Load DDR trace ... file' menu to load custom files.", 1)
                idaapi.warning("API trace file not loaded. Pls use DDRs 'Trace' menu to run a trace first OR load custom trace files via 'File/Load File/Load DDR trace ... file'.") 
                return 1

            if JSONFILE_LOADED == False:
                self._trace_file_not_loaded()
                return 1
            else:
                items = self._get_chooser_row_list_calls()
                DDR_print_mesg("Len items: {:d}".format(len(items)),7)
                columns = [ ["Instruction #", 6], ["Address", 12], ["Type", 15], ["Data Address", 12], ["API call", 12], ["API module", 12], ["Disasm", 35] ]
                c = DDR_MyChoose("API Call Trace", items=items, modal=False, columns=columns)
                if not c.show():
                    DDR_print_mesg("Failed to start Call Trace View", 7)
                else:
                    DDR_print_mesg("Call Trace View created", 7)
                return 1 

        # Load DynamoRIO JSON File 
        if self.cmd == "LoadFile" :
            self._load_JSON_file()
            return 1

        # Load DynamoRIO JSON API File 
        if self.cmd == "LoadAPIFile" :
            api_filename=self._get_JSON_APIfile()
            if api_filename == None:
                DDR_print_mesg("File dialog canceled. API file not loaded.",1)
                return 1
            if self._load_APICalls_file(api_filename):
                DDR_print_mesg("API trace file imported: {}".format(api_filename),7)
            return 1

        # Delete comments
        if self.cmd == "DeleteNonRepeatableComments":
            try:
                start = idc.read_selection_start()
                end   = idc.read_selection_end()

                DDR_print_mesg("Selected processing range: {:x} - {:x} ea={:x}".format(start, end, ea),7)
                # User selected multiple lines
                if start != idc.BADADDR and end != idc.BADADDR:
                    DDR_print_mesg("Selected processing range: {:x} - {:x}".format(start, end),7)
                    ea=start
                    while(ea < end):  
                        DDR_print_mesg("Delete non-repeatable comment at {:x}".format(ea)) 
                        idc.set_cmt(ea, "",0)         
                        ea = idc.next_head(ea)
                # User selected single line
                else: 
                    DDR_print_mesg("Selected address: {:x}".format(ea),7) 
                    DDR_print_mesg("Delete non-repeatable comment at line {:x}".format(ea))
                    idc.set_cmt(ea, "",0)

                DDR_print_mesg("Done. repeatable Comments deleted.")
                DDR_print_mesg("-------------------------------------------------------------------------------")

                return 1
            except:
                raise


        # Parse all the other menu items (parse_ea) which can be applied to a single ea or a range (e.g. Get src op value etc)
        if(JSONFILE_LOADED == True):
            start = idc.read_selection_start()
            end   = idc.read_selection_end()
            # User selected multiple lines
            if start != idc.BADADDR and end != idc.BADADDR:
                DDR_print_mesg("Selected processing range: {:x} - {:x}".format(start, end),7)
                ea=start
                while(ea < end):
                    self._parse_ea(ea)
                    ea = idc.next_head(ea)
            # User selected single line
            else:
                DDR_print_mesg("Selected address: {:x}".format(ea),7)
                self._parse_ea(ea)

            DDR_print_mesg("Done.")
            DDR_print_mesg("-------------------------------------------------------------------------------")
        else:
            self._trace_file_not_loaded()
            
        return 1


    def op_count(self, ea):
        """Return the number of operands of given instruction"""
        for i in range(10):
            op = idc.print_operand(ea,i)
            if op:
                continue
            break
        return i    


    def _get_dump_para_addr_n_op(self, op_str, ea):
        """ Get dump parameters """
        DDR_print_mesg("op_str = {}".format(op_str),7)
        num_operands = self.op_count(ea)
        DDR_print_mesg("Number of operands for instruction at {:x}: {:d}".format(ea, num_operands),7)

        if num_operands < 2:
            op_src = idc.print_operand(ea,0)
            DDR_print_mesg("op_src = {}".format(op_src),7)
            if op_str in op_src:
                DDR_print_mesg("src op match",7)
                ret = "SP0"
            else:
                DDR_print_mesg("Highlighted operand not found pls set it manually",1)
                ret = ida_kernwin.ask_text(3, "SP0", "Highlighted operand not found, pls enter value manually.")
                if ret == "":
                    ret = None
                    ea  = None

        elif num_operands < 3:
            op_dst = idc.print_operand(ea,0)
            op_src = idc.print_operand(ea,1)
            DDR_print_mesg("op_dst = {}".format(op_dst),7)
            DDR_print_mesg("op_src = {}".format(op_src),7)
            if op_str in op_src:
                DDR_print_mesg("src op match",7)
                ret = "SP0"
            elif op_str in op_dst:
                DDR_print_mesg("dst op match",7)
                ret = "DP0"
            else:
                DDR_print_mesg("Highlighted operand not found pls set it manually",1)
                ret = ida_kernwin.ask_text(3, "SP0", "Highlighted operand not found, pls enter value manually.")
                if ret == "":
                    ret = None
                    ea  = None

        else:
            idaapi.warning("Sorry, instructions with more than 2 operands are not supported in the moment.")
            DDR_print_mesg("Sorry, instructions with more than 2 operands are not supported in the moment.",2)
            return (None,None)

        return (ea,ret)

    def _parse_ea(self, ea):
        """
        Main parser method for parsing and analysing data at address (ea)
        """
        
        op0 = idc.print_operand(ea,0)
        op1 = idc.print_operand(ea,1)
        op0_type = idc.get_operand_type(ea,0)
        op1_type = idc.get_operand_type(ea,1)

        DDR_print_mesg("Address: 0x{:x}:".format(ea))
        DDR_print_mesg("  DISASM   : {}".format(idc.GetDisasm(ea)))
        DDR_print_mesg("  CMD      : {}".format(self.cmd), 7)
        DDR_print_mesg("  OP0      : {}".format(op0),7)
        DDR_print_mesg("  OP1      : {}".format(op1),7)
        DDR_print_mesg("  OP0 type : {:x}".format(op0_type),7)
        DDR_print_mesg("  OP1 type : {:x}".format(op1_type),7)
        DDR_print_mesg("  OP0 value: {:x}".format(idc.get_operand_value(ea,0)),7)
        DDR_print_mesg("  OP1 value: {:x}".format(idc.get_operand_value(ea,1)),7)

        # Handle selected menu entries
        if self.cmd == "GetSrcOpValue":
            trace_instr_num_list = self._get_trace_instr_list(ea) 
            self._set_ida_src_op_name(ea)
            self._handle_trace_list(ea, self.src_op, "inst_mem_addr_src0", trace_instr_num_list, "SrcVal")

        if self.cmd == "GetDstOpValue":
            trace_instr_num_list = self._get_trace_instr_list(ea) 
            self._set_ida_dst_op_name(ea)
            self._handle_trace_list(ea, self.dst_op, "inst_mem_addr_dst0", trace_instr_num_list, "DstVal")
          
        if self.cmd == "GetSrcOpPtrValue":
            DDR_print_mesg("Menu GetSrcOpPtrValue selected")
            trace_instr_num_list = self._get_trace_instr_list(ea) 
            self._set_ida_src_op_name(ea)
            self._handle_trace_list_ptr_value(ea, self.src_op, "inst_mem_addr_src0_data", trace_instr_num_list, "SrcValPtr")

        if self.cmd == "GetDstOpPtrValue":
            DDR_print_mesg("Menu GetDstOpPtrValue selected")
            trace_instr_num_list = self._get_trace_instr_list(ea) 
            self._set_ida_dst_op_name(ea)
            self._handle_trace_list_ptr_value(ea, self.dst_op, "inst_mem_addr_dst0_data", trace_instr_num_list, "DstValPtr")

        if self.cmd == "GetSrcOpPtrPtrValue":
            DDR_print_mesg("Menu GetSrcOpPtrPtrValue selected")
            trace_instr_num_list = self._get_trace_instr_list(ea) 
            self._set_ida_src_op_name(ea)
            self._handle_trace_list_ptr_value(ea, self.src_op, "inst_mem_addr_src0_data_ptr_data", trace_instr_num_list, "SrcValPtrPtr")

        if self.cmd == "GetDstOpPtrPtrValue":
            DDR_print_mesg("Menu GetDstOpPtrPtrValue selected")
            trace_instr_num_list = self._get_trace_instr_list(ea) 
            self._set_ida_dst_op_name(ea)
            self._handle_trace_list_ptr_value(ea, self.dst_op, "inst_mem_addr_src0_data_ptr_data", trace_instr_num_list, "DstValPtrPtr")

        if self.cmd=="GetSrcPtrMem":
            trace_instr_num_list = self._get_trace_instr_list(ea)
            self._set_ida_src_op_name(ea)
            self._handle_trace_list_ptr(ea, self.src_op, "inst_mem_addr_src0_data", trace_instr_num_list, "SrcPtrMem")  

        if self.cmd=="GetDstPtrMem":
            trace_instr_num_list = self._get_trace_instr_list(ea)
            self._set_ida_dst_op_name(ea)
            self._handle_trace_list_ptr(ea, self.dst_op, "inst_mem_addr_dst0_data", trace_instr_num_list, "DstPtrMem")  

        if self.cmd=="GetSrcPtrPtrMem":
            trace_instr_num_list = self._get_trace_instr_list(ea)
            self._set_ida_src_op_name(ea)
            self._handle_trace_list_ptr(ea, self.src_op, "inst_mem_addr_src0_data_ptr_data", trace_instr_num_list, "SrcPtrPtrMem")

        if self.cmd=="GetDstPtrPtrMem":
            trace_instr_num_list = self._get_trace_instr_list(ea)
            self._set_ida_dst_op_name(ea)
            self._handle_trace_list_ptr(ea, self.dst_op, "inst_mem_addr_dst0_data_ptr_data", trace_instr_num_list, "DstPtrPtrMem")

        # --- Register menu selected ---
        # REG_LIST(64) = ["xax","xbx","xcx","xdx","xsp","xbp","xsi","xdi","r8","r9","r10","r11","r12","r13","r14","r15"]
        # REG_LIST(32) = ["xax","xbx","xcx","xdx","xsp","xbp","xsi","xdi"]
        for reg in REG_LIST:
            DDR_print_mesg("reg:{}".format(reg), 7) 
            if self.cmd=="Get_Mem_Ptr_" + reg:
                trace_instr_num_list = self._get_trace_instr_list(ea)
                self._set_ida_comment_regmenu_mem_ptr(ea, reg, trace_instr_num_list)

        return True

    def _trace_file_not_loaded(self):
        DDR_print_mesg("Trace file not loaded. Pls use DDRs 'Trace' menu to run a trace first.", 1)
        DDR_print_mesg("Alternativly, you can use 'File/Load File/Load DDR trace file' to load a custom trace file.", 1)
        idaapi.warning("Trace file not loaded. Pls use DDRs 'Trace' menu to run a trace first OR load a custom trace file via 'File/Load File/Load DDR trace file'.") 

    def _load_APICalls_file(self, filename = None):
        global APIFILE_LOADED
        global API_CALLS
        API_CALLS = {}      # reset API calls dict

        try:
            with open(filename) as apicall_file: 
                apicall_data = json.load(apicall_file)
            DDR_print_mesg("API calls file loaded: {}".format(filename))
            APIFILE_LOADED = True
        except IOError as e:
            DDR_print_mesg("Failed to open API calls file. Quitting operation.")
            APIFILE_LOADED = False
            return False
        except ValueError as e:
            DDR_print_mesg("Failed to parse API calls file: {}. Quitting operation.".format(filename),2)
            DDR_print_mesg("Please check file format is JSON and try again.",2)
            e_args = ''.join(map(str,e.args))
            DDR_print_mesg("{}".format(e_args))
            APIFILE_LOADED = False
            idaapi.warning("[DDR] Failed to open API calls file. Check Logging window for details.")
            return False

        # Build call dict with address as key
        for call in apicall_data["apicalls"]:
            if call["address"] in API_CALLS:
                old_name = API_CALLS[call["address"]]
                API_CALLS[call["address"]] = call["name"] + "," + old_name  
            else:
                API_CALLS[call["address"]] = call["name"] + ";" + call["module"]

        return True


    def _load_JSON_file(self, filename = None):
        """ Loads DynamoRIO generated JSON file into memory """
        global JSONDATA
        global JSONFILE_LOADED

        if filename == None:
            filename = self._get_JSON_file()
            if filename == None:
                # Dialog canceled
                DDR_print_mesg("File dialog canceled. Trace file not loaded.",1)
                return False

        try:
            with open(filename) as json_file:  
                JSONDATA = json.load(json_file)
                JSONFILE_LOADED = True
                DDR_print_mesg("JSON file was generated for executable: {}".format((JSONDATA["samplename"])))
                if DDR_CONFIG_SETTINGS.ARCH_BITS == 64 and JSONDATA["architecture"] != "x64":
                    idaapi.warning("JSON File has wrong architecture: {}".format(JSONDATA["architecture"]))
                    DDR_print_mesg("JSON File has wrong architecture: {}".format(JSONDATA["architecture"]))
                    JSONFILE_LOADED = False
                    return False
                elif DDR_CONFIG_SETTINGS.ARCH_BITS == 32 and JSONDATA["architecture"] != "x32":
                    idaapi.warning("JSON File has wrong architecture: {}".format(JSONDATA["architecture"]))
                    DDR_print_mesg("JSON File has wrong architecture: {}".format(JSONDATA["architecture"]))
                    JSONFILE_LOADED = False
                    return False
                else:
                    DDR_print_mesg("Instructions traced in JSON file: {:d}".format(len(JSONDATA["instruction"])))
                    return True
        except IOError as e:
            DDR_print_mesg("Failed to open JSON file. Quitting operation.")
            JSONFILE_LOADED = False
            return False
        except ValueError as e:
            DDR_print_mesg("Failed to parse JSON file: {}. Quitting operation.".format(filename),2)
            DDR_print_mesg("Please check file format is JSON and try again.",2)
            e_args = ''.join(map(str,e.args))
            DDR_print_mesg("{}".format(e_args))
            JSONFILE_LOADED = False
            idaapi.warning("[DDR] Failed to open JSON file. Check Logging window for details.")
            return False
        else:
            DDR_print_mesg("Failed to open JSON file: {} with unknown error. Quitting operation.".format(filename),2)
            JSONFILE_LOADED = False
            idaapi.warning("[DDR] Failed to open JSON file. Check Logging window for details.")
            raise


    def _get_JSON_file(self):
        """
        Prompt a file selection dialog, returning file selections.
        """
        my_dir = JSONFILE_DIR

        # create & configure a Qt File Dialog for immediate use
        file_dialog = QtWidgets.QFileDialog(
            None,
            'Open DDR generated instructions trace JSON file', # IDA does not show this capition itm.
            #idautils.GetIdbDir(),
            my_dir,
            'All Files (*.*)'
        )
        file_dialog.setFileMode(QtWidgets.QFileDialog.ExistingFiles)

        # prompt the user with the file dialog, and await filename(s)
        filename, _ = file_dialog.getOpenFileName()

        if filename == '':
            return None

        # log the selected filenames from the dialog
        DDR_print_mesg("DDR instructions trace JSON file loaded: {}".format(filename))

        # return the captured filenames
        return filename


    def _get_JSON_APIfile(self):
        """
        Prompt a file selection dialog, returning file selections.
        """
        my_dir = JSONFILE_DIR

        # create & configure a Qt File Dialog for immediate use
        file_dialog = QtWidgets.QFileDialog(
            None,
            'Open DDR generated API trace JSON file',  # IDA does not show this capition itm.  
            #idautils.GetIdbDir(),
            my_dir,
            'All Files (*.*)'
        )
        file_dialog.setFileMode(QtWidgets.QFileDialog.ExistingFiles)

        # prompt the user with the file dialog, and await filename(s)
        filename, _ = file_dialog.getOpenFileName()

        if filename == '':
            return None

        # log the selected filenames from the dialog
        DDR_print_mesg("DDR API trace JSON file loaded: {}".format(filename),7)

        # return the captured filenames
        return filename


    def _highlight_trace(self):
        """ Highlight basics blocks in IDA graph disasm view"""
        
        #idaapi.get_item_color(ea)   # Get actual color at ea

        # COLOR                                                                                             # Number of instr. executed
        # -----------------------------------------------------------------------------------------------------------------------------
        # green         = [0xa0ffa0, 0x70ff70, 0x00ff00, 0x00ef00, 0x00df00, 0x00cf00, 0x00bf00, 0x00af00]  # 0x0   - 0xf
        # yellow        = [0xd0ffff, 0xa0ffff, 0x70ffff, 0x00ffff, 0x00efff, 0x00dfff, 0x00cfff, 0x00bfff]  # 0x10  - 0xff
        # red           = [0xC0C0ff, 0xB0B0ff, 0xA0A0ff, 0x8080ff, 0x6060ff, 0x4040ff, 0x2020ff, 0x0000ff]  # 0x100 -  0xfff
        # light_purple  = 0xFF60FF                                                                          # 0x1000  - 0xffff
        # purple        = 0xFF00FF                                                                          # 0x10000 - infinity

        inst_addr_list = []
        for instr in JSONDATA["instruction"]:
            inst_addr_list.append(int(instr["address"], 16))

        inst_addr_counter = Counter(inst_addr_list)

        # paint lines depending on the number of occurences in the trace 
        # Default: (green 0-0xf, yellow 0x10 - 0xff, red 0x100 - 0xfff, light_purple 0x1000 - 0xffff, purple > 0xffff)
        for ea in inst_addr_counter:
            if inst_addr_counter[ea] > 0xffff:  
                idaapi.set_item_color(ea, COLOR5)
                continue

            if inst_addr_counter[ea] > 0xfff:  
                idaapi.set_item_color(ea, COLOR4)
                continue

            if inst_addr_counter[ea] > 0xff: 
                if inst_addr_counter[ea] <= 7:
                    color = COLOR3[inst_addr_counter[ea]]
                else:
                    color = COLOR3[7]
                idaapi.set_item_color(ea, color)
                continue

            if inst_addr_counter[ea] > 0xf: 
                if inst_addr_counter[ea] <= 7:
                    color = COLOR2[inst_addr_counter[ea]]
                else:
                    color = COLOR2[7]
                idaapi.set_item_color(ea, color)
                continue

            if inst_addr_counter[ea] <= 0xf:
                if inst_addr_counter[ea] <= 7:
                    color = COLOR1[inst_addr_counter[ea]]
                else:
                    color = COLOR1[7]
                idaapi.set_item_color(ea, color)
     
    def _clear_highlight_trace(self, ea):
        # Loop through all instructions
        for ea in idautils.Segments():
            seg_start = idc.get_segm_start(ea)
            seg_end = idc.get_segm_end(ea)

            addr = seg_start
            while (addr < seg_end) and (addr != idaapi.BADADDR):
                # Set color to default white background
                idaapi.set_item_color(addr, 0xffffffff)
                addr = idc.next_head(addr)

    def _find_apicall_for_addr(self, addr):
        """ Check if addr is a symbol / known api address"""
        if addr in API_CALLS:
            apiname_with_module = API_CALLS[addr]
            name = apiname_with_module.split(";")[0]
            module = apiname_with_module.split(";")[1]
            return name, module
        return "-", "-"

    def _get_chooser_row_list_strings(self) :
        """ Fill the lines list used for the Trace Window item list"""

        lines = []
        for instr in JSONDATA["instruction"]:                               
            for s in [ "inst_mem_addr_src0_data", "inst_mem_addr_src0_data_ptr_data"]:
                try:
                    line  = []
                    line.append("{:12d}".format(int(instr["instr_num"],10)))    # instruction number in trace
                    line.append(str(instr["address"]))                          # ea address of instruction
                    if s == "inst_mem_addr_src0_data":
                        datatype = "Pointer"                                    
                    else:
                        datatype = "Pointer-Pointer"                            
                    line.append(datatype)                                       # JSON field (type) of string
                    mystring = str(instr[s])
                    stringlist = mystring.split("    ")
                    line.append(stringlist[0])                                  # String in Hex
                    data_addr = self._get_addr_from_data(mystring)
                    line.append(data_addr)                                      # Data address extracted from mystring
                    line.append("".join(stringlist[1:]))                        # String
                    line.append(str(instr["disasm"]))                           # Disasm of instruction
                    lines.append(line)                                          # add generated line to lines list
                except KeyError as e:
                    break                                                       # key doesn't exists in JSON file, try next instruction

        return lines

    def _get_chooser_row_list_calls(self) :
        """ Fill the lines list used for the API calls Window item list"""
        lines = []
        for instr in JSONDATA["instruction"]:    
            try:
                disasm = instr["disasm"]
            except:
                disasm = "na"                          
            
            try:
                val = instr["inst_mem_addr_src0"]
            except:
                val = "0x0"

            try:
                ptr = instr["inst_mem_addr_src0_data_ptr"]
            except:
                ptr = "0x0"

            if disasm == "na":
                return lines

            else:
                # correlate api function addresses with inst_mem_addr_src0 or inst_mem_addr_src0_data_ptr
                apicall_val, apimodule_val = self._find_apicall_for_addr(val)
                apicall_ptr, apimodule_ptr = self._find_apicall_for_addr(ptr) 

                if apicall_val == '-' and apicall_ptr == '-':
                    continue

                
                if apicall_val != '-':
                    datatype  = "Value"
                    data_addr = val
                    apicall   = apicall_val
                    apimodule = apimodule_val
                else:
                    datatype  = "Pointer"
                    data_addr = ptr
                    apicall   = apicall_ptr
                    apimodule = apimodule_ptr

                line  = []
                line.append("{:12d}".format(int(instr["instr_num"],10)))    # instruction number in trace
                line.append(str(instr["address"]))                          # ea address of instruction
                line.append(datatype)                                       # JSON field (type) of string
                line.append(data_addr)                                      # address called
                line.append(apicall)                                        # API call name
                line.append(apimodule)                                      # API module name
                line.append(disasm)                                         # Disasm of instruction
                lines.append(line)                                          # add generated line to lines list

        return lines

    def _set_ida_src_op_name(self, ea):
        """ Find source operand of instruction and set src_op name"""
        op0 = idc.print_operand(ea,0)
        op1 = idc.print_operand(ea,1)
        op0_type = idc.get_operand_type(ea,0)
        op1_type = idc.get_operand_type(ea,1)

        if op1_type == idaapi.o_void:
            self.src_op = op0
        else:
            self.src_op = op1

        if op0_type == idaapi.o_void:
            self.src_op = "src_op not set"

    def _set_ida_dst_op_name(self, ea):
        """ Find destination operand of instruction and set dst_op name"""
        op0 = idc.print_operand(ea,0)
        op1_type = idc.get_operand_type(ea,1)

        if op1_type == idaapi.o_void:
            self.dst_op = "dst_op not set"
        else:
            self.dst_op = op0

    def _handle_trace_list_ptr_value(self, ea, op, json_instr_field, trace_instr_num_list, usrcmd):
        comment = ""
        if not trace_instr_num_list:
            comment = "Instruction at {:x} not in trace".format(ea)

        #TBD: sub optimal algorithm, fix this
        for c,i in enumerate(trace_instr_num_list):
            try:
                if c < DDR_CONFIG_SETTINGS.MAX_LOG_ROUNDS:
                    data = JSONDATA["instruction"][i][json_instr_field]
                    addr = self._get_addr_from_data(data)
                    apicall, apimodule = self._find_apicall_for_addr(addr)
                    DDR_print_mesg("    {}:{}  APIname:{}  Module:{}".format(usrcmd,addr,apicall,apimodule))
                if c < DDR_CONFIG_SETTINGS.MAX_CMT_ROUNDS:
                    data = JSONDATA["instruction"][i][json_instr_field]
                    addr = self._get_addr_from_data(data)
                    apicall, apimodule = self._find_apicall_for_addr(addr)
                    if apicall == '-':
                        comment += "{}:0x{:x}; ".format(usrcmd, int(addr,16))  # delete leading zeros
                    else:
                        comment += "{}:{}; ".format(usrcmd, apicall)  
            except:
                if c < DDR_CONFIG_SETTINGS.MAX_LOG_ROUNDS:
                    DDR_print_mesg("    {}:InstrAddr=0x{:x} instr number:{:07d} = No pointer found".format(usrcmd, ea, i))
                if c < DDR_CONFIG_SETTINGS.MAX_CMT_ROUNDS:
                    comment += "{}:PtrNotFound; ".format(usrcmd)


        # Get existing commend and add our dynamic data to it.
        org_cmt = idc.GetCommentEx(ea, False)
        if org_cmt != None:
            new_comment = org_cmt + "\n" + comment
        else:
            new_comment = comment

        idc.set_cmt(ea, str(new_comment),0)   # IDA doesn't like unicode from JSON file 


    def _handle_trace_list(self, ea, op, json_instr_field, trace_instr_num_list, usrcmd):
        """ Get content from 'json_instr_field' key from JSON file and generate corrosponding comment for value"""
        comment = ""
        if not trace_instr_num_list:
            comment = "Instruction at 0x{:x} not in trace".format(ea)

        for c,i in enumerate(trace_instr_num_list):
            try:
                addr = JSONDATA["instruction"][i][json_instr_field]
                if c < DDR_CONFIG_SETTINGS.MAX_LOG_ROUNDS:
                    apicall, apimodule = self._find_apicall_for_addr(addr)
                    DDR_print_mesg("    {}({})@0x{:x}@instr number:{:07d} = 0x{:x} (APIname:{}  Module:{})".format(usrcmd, op, ea, i, int(addr,16), apicall,apimodule))

                if c < DDR_CONFIG_SETTINGS.MAX_CMT_ROUNDS:
                    apicall, apimodule = self._find_apicall_for_addr(addr)
                    if apicall == '-':
                        comment += "{}:0x{:x}; ".format(usrcmd, int(addr,16))  # delete leading zeros
                    else:
                        comment += "{}:{}; ".format(usrcmd, apicall)  
            except:
                if c < DDR_CONFIG_SETTINGS.MAX_LOG_ROUNDS:
                    DDR_print_mesg("    {}({})@0x{:x}@instr number:{:07d} = No memory data found in JSON file".format(usrcmd, op, ea, i))
                if c < DDR_CONFIG_SETTINGS.MAX_CMT_ROUNDS:
                    comment += "{}({}):NotFound; ".format(usrcmd,op) 

        # Get existing commend and add our dynamic data to it.
        org_cmt = idc.GetCommentEx(ea, False)
        if org_cmt != None:
            new_comment = org_cmt + "\n" + comment
        else:
            new_comment = comment

        idc.set_cmt(ea, str(new_comment),0)    

    def _handle_trace_list_ptr(self, ea, op, json_instr_field, trace_instr_num_list, usrcmd):
        """ Get content from 'json_instr_field' key from JSON file and generate corrosponding comment for pointer or pointer-pointer"""
        comment = ""
        if not trace_instr_num_list:
            comment = "Instruction at 0x{:x} not in trace".format(ea)

        for c,i in enumerate(trace_instr_num_list):
            try:
                if c < DDR_CONFIG_SETTINGS.MAX_LOG_ROUNDS:
                    DDR_print_mesg("    {}({})@0x{:x}@instr number:{:07d} = {}".format(usrcmd, op, ea, i, str(JSONDATA["instruction"][i][json_instr_field])))
                if c < DDR_CONFIG_SETTINGS.MAX_CMT_ROUNDS:
                    comment += "{}({}):{}; ".format(usrcmd, op, str(JSONDATA["instruction"][i][json_instr_field])) 
            except:
                if c < DDR_CONFIG_SETTINGS.MAX_LOG_ROUNDS:
                    DDR_print_mesg("    {}({})@0x{:x}@instr number:{:07d} = No memory data found in JSON file".format(usrcmd, op, ea, i))
                if c < DDR_CONFIG_SETTINGS.MAX_CMT_ROUNDS:
                    comment += "{}({}):NotFound; ".format(usrcmd,op) 

        # Get existing commend and add our dynamic data to it.
        org_cmt = idc.GetCommentEx(ea, False)
        if org_cmt != None:
            new_comment = org_cmt + "\n" + comment
        else:
            new_comment = comment

        idc.set_cmt(ea, new_comment,0)        
 
    def _set_ida_comment_regmenu_mem_ptr(self, ea, reg, trace_instr_num_list):
        """ Prints out the content stored in the register and the memory data the register value is pointing to, then adds a comment to the disasm view"""
        comment = ""

        if not trace_instr_num_list:
            comment = "Instruction at 0x{:x} not in trace".format(ea)

        for c,i in enumerate(trace_instr_num_list):
            if c < DDR_CONFIG_SETTINGS.MAX_LOG_ROUNDS:
                try:
                    DDR_print_mesg("    {}(={})@instr number:{:07d} = {}".format(reg, JSONDATA["instruction"][i][reg], i, JSONDATA["instruction"][i][reg + "_ptr_data"]))
                except:
                    DDR_print_mesg("    {}@instr number:{:07d} = not found in data. Did you run a full trace ?".format(reg,i))
            if c < DDR_CONFIG_SETTINGS.MAX_CMT_ROUNDS:
                s_pdata  = "NO_DATA"
                regval   = 0

                try:
                    s_pdata  = str(JSONDATA["instruction"][i][reg + "_ptr_data"])
                    regval   = int(JSONDATA["instruction"][i][reg],16)
                except:
                    pass

                s_cmt = str("{}(0x{:x})={}; ".format(reg, regval, s_pdata)) #IDA doesn't like unicode type

                comment += s_cmt

        # Get existing commend and add our dynamic data to it.
        org_cmt = idc.GetCommentEx(ea, False)
        if org_cmt != None:
            new_comment = org_cmt + "\n" + comment
        else:
            new_comment = comment

        idc.set_cmt(ea, new_comment,0)

    def _get_addr_from_data(self, data):
        """ Extracts an address from the data bytes stream e.g. 04 03 02 01 00 00 00 ... => 01020304. Returns the addrs as string with a leading 0x... """
        if DDR_CONFIG_SETTINGS.ARCH_BITS == 64:
            s_addr = data[21:23] + data[18:20] + data[15:17] + data[12:14] + data[9:11] + data[6:8] + data[3:5] + data[:2] 
        else:
            s_addr = data[9:11] + data[6:8] + data[3:5] + data[:2] 

        return "0x" + s_addr

    def _get_trace_instr_list(self, ea):
        """ Get instruction number list from JSON trace file for ea address, if operand is a memory address"""

        trace_instr_num_list = []
        trace_instr_num_list = self._find_instr_addr_in_json(ea, trace_instr_num_list)

        DDR_print_mesg("Build list for ea=0x{:x} length_trace_instr_num_list={:d}".format(ea, len(trace_instr_num_list)),7)

        if not trace_instr_num_list:
            DDR_print_mesg("Instruction address 0x{:x} not found in trace file.".format(ea))

        return trace_instr_num_list

    def _find_instr_addr_in_json(self, ea, trace_instr_num_list):
        """ Return a list of all instructions from the JSON file which match the ea address"""
        i=0
        c=0
        for instr in JSONDATA["instruction"]:                               # TBD change to eval..
            jsonaddr = int(JSONDATA["instruction"][i]["address"], 16)
            if jsonaddr == ea:
                DDR_print_mesg("    {:d}. Found instr address    : {:x}".format(i, ea),7)
                trace_instr_num_list.append(i)
                c += 1
                if c > DDR_CONFIG_SETTINGS.MAX_INSTR_COUNT-1: break
            i += 1

        DDR_print_mesg("    Number of times instruction at address (0x{:x}) executed in trace file: {:d}".format(ea, len(trace_instr_num_list)))

        return trace_instr_num_list

    def _Add_BB2BBlist(self, ea):
        """ runs DynamoRio analysis against the marked basic block (only works for functions)"""
        f = idaapi.get_func(ea)

        if not f:
            DDR_print_mesg("No function found at 0x{:x}. This only works inside of functions.".format(ea))
            idaapi.warning("No function found at 0x{:x}. This only works inside of functions.".format(ea))
            return False

        fc = idaapi.FlowChart(f)
        for block in fc:
            if block.start_ea <= ea:
                if block.end_ea > ea:
                    DDR_print_mesg("0x{:x} is part of block [0x{:x} - 0x{:x}]".format(ea, block.start_ea, idc.prev_head(block.end_ea)))
                    BB_LIST.add(block.start_ea, idc.prev_head(block.end_ea))
        return True

    def _Remove_BB2BBlist(self, ea):
        """ runs DynamoRio analysis against the marked basic block (only works for functions)"""
        f = idaapi.get_func(ea)

        if not f:
            DDR_print_mesg("No function found at 0x{:x}. This only works inside of functions.".format(ea))
            idaapi.warning("No function found at 0x{:x}. This only works inside of functions.".format(ea))
            return False

        fc = idaapi.FlowChart(f)
        for block in fc:
            if block.start_ea <= ea:
                if block.end_ea > ea:
                    DDR_print_mesg("0x{:x} is part of block [0x{:x} - 0x{:x}]".format(ea, block.start_ea, idc.prev_head(block.end_ea)))
                    BB_LIST.remove(block.start_ea, block.end_ea)
        return True

    def _exec_dynRIO_against_BB(self, ea):
        """ runs DynamoRio analysis against the marked basic block (only works for functions)"""
        f = idaapi.get_func(ea)

        if not f:
            DDR_print_mesg("No function found at 0x{:x}. This only works inside of functions.".format(ea))
            idaapi.warning("No function found at 0x{:x}. This only works inside of functions.".format(ea))
            return

        fc = idaapi.FlowChart(f)
        for block in fc:
            if block.start_ea <= ea:
                if block.end_ea > ea:
                    DDR_print_mesg("0x{:x} is part of block [0x{:x} - 0x{:x}]".format(ea, block.start_ea, idc.prev_head(block.end_ea)))
                    self._exec_dynRio(start_addr=block.start_ea, end_addr=idc.prev_head(block.end_ea), instr_count=DDR_CONFIG_SETTINGS.MAX_INSTR_TO_EXECUTE)
             
    #TBD: do some input checks
    def _exec_dynRio(self, start_addr=None, end_addr=None, break_addr=None, instr_count=None, options=None):
        """ Wrapper to prepare calling the DDRsever.py API _call_api method to send cmds"""
        if start_addr == None or end_addr == None or instr_count == None:
            DDR_print_mesg("start_addr, end_addr or instr_count not set")
            return False

        DDR_print_mesg("Block 0x{:x} - 0x{:x}".format(start_addr, end_addr),7)

        dynrio_sample = SAMPLE_FILENAME

        if options == "light_trace_only":
            cmd_id = 3
        else:
            cmd_id = 1

        self._call_api(cmd_id       = cmd_id,                                   # run dynamoRIO client against sample 
                      dynrio_sample = dynrio_sample,                            # sample file to analyse
                      start_addr    = start_addr,                               # start address to log
                      end_addr      = end_addr,                                 # end address to log
                      instr_count   = DDR_CONFIG_SETTINGS.MAX_INSTR_TO_EXECUTE, # max number of instructions to execute in the trace
                      arch_bits     = DDR_CONFIG_SETTINGS.ARCH_BITS,            # 32 or 64 bit sample architecture
                      options       = options)                                  # option e.g. light_trace_only        
        
        return True

    def _test_zip_is_valid(self, zipfilename):

        if not os.path.isfile(zipfilename):
             DDR_print_mesg("Zip file not found.", 2)
             return False

        try:
            zip_file = zipfile.ZipFile(zipfilename)
            res = zip_file.testzip()

            if res is not None:
                DDR_print_mesg("Bad file in zip archive: {}".format(res), 2)
                return False
            else:
                return True
        except:
            DDR_print_mesg("Zip file: {} is corrupt. File content:\n".format(zipfilename), 2)

            with open(zipfilename, 'r') as f:
                print(f.read())

            os.remove(zipfilename)
            DDR_print_mesg("Corrupt zip file: {} deleted.".format(zipfilename), 2)
            return False

    def get_hash_from_str(self,s):
        md5 = hashlib.md5()
        sha1 = hashlib.sha1()
        sha256 = hashlib.sha256() 
        md5.update(s.encode())
        sha1.update(s.encode())
        sha256.update(s.encode())
        return {"md5_sum"   :md5.hexdigest().upper(),
                "sha1_sum"  :sha1.hexdigest().upper(), 
                "sha256_sum":sha256.hexdigest().upper()}

    def http_ping(self, url):
        """ Quick test if server is responding """
        ret = False

        if DDR_CONFIG_SETTINGS.VERIFY_CERT:
            req_verify =  DDR_CONFIG_SETTINGS.CA_CERT
        else:
            req_verify =  False
                
        try :
            with requests.get(url, verify=req_verify, timeout=3) as res:
                if res.status_code == 200:
                    DDR_print_mesg("HTTP ping successful")
                    ret = True

        except requests.exceptions.SSLError as e:
            DDR_print_mesg("[API SSL ERROR]: {}".format(e),2)
            DDR_print_mesg("Did you configure a valid serial number for the certificate ?",2)
            DDR_print_mesg("For troubleshooting install the certificate from the <IDA plugin dir>/ddr into your browser.",2)
            DDR_print_mesg("Access the DDR server test webpage with your browser and check the error message.",2)
            DDR_print_mesg("Reminder: Test webpage is {}".format(url),2)

        except requests.exceptions.ConnectionError as e:
            DDR_print_mesg("[CONNECTION ERROR]: {}".format(e),2)
            DDR_print_mesg("Did you start and configure the DDR server script or is there any firewall inbetween ?",2)

        except requests.exceptions.HTTPError as e:
            DDR_print_mesg("[HTTP ERROR]: {}".format(e),2)

        except requests.exceptions.ReadTimeout as e:
            DDR_print_mesg("[TIMEOUT ERROR]: {}".format(e),2)

        except:
            DDR_print_mesg("[UNKNOWN ERROR] Unkown error happend. HTTP ping failed.",2)

        return ret


    def _call_api(self, cmd_id=0, dynrio_sample=None, start_addr=None, end_addr=None, instr_count=None, arch_bits=None, options=None, dump_para_json={}):
        """ Main method to call the DDRserver.py API"""

        ret = False

        # Let's test servers reachabillity first
        if not self.http_ping("https://" + DDR_CONFIG_SETTINGS.WEBSERVER + ":" + DDR_CONFIG_SETTINGS.WEBSERVER_PORT + "/"):
            DDR_print_mesg("HTTP ping failed.")
            return False
        
        # ----- New JSON API -------
        if cmd_id >= 5 and cmd_id <= 12:
            url = "https://" + DDR_CONFIG_SETTINGS.WEBSERVER + ":" + DDR_CONFIG_SETTINGS.WEBSERVER_PORT + "/api/v1/json"
            DDR_print_mesg("Using new JSON API...")
        # Error
        else:
            DDR_print_mesg("Unknown command id: {:d}\n".format(cmd_id), 2)
            return False

        if DDR_CONFIG_SETTINGS.VERIFY_CERT:
            req_verify =  DDR_CONFIG_SETTINGS.CA_CERT
        else:
            req_verify =  False

        try:            
            # ---- Dump buffer ----
            if cmd_id == 5:
                ret = False
                dump_para_json['id'] = cmd_id

                tmp_file_list      = []
                dump_filename_list = []
            
                for upload_counter in range(DDR_CONFIG_SETTINGS.MAX_UPLOAD_ATTEMPTS, -1, -1):
                    if upload_counter < 1:
                        DDR_print_mesg("Failed to upload sample file to server. Max number of upload attempts reached",2)
                        return False

                    with requests.post(url, verify=req_verify, json=dump_para_json, timeout=DDR_CONFIG_SETTINGS.MAX_API_TIMEOUT, stream=False) as res:
                        if res.status_code == 200:
                            try:
                                dump_files = res.json()["files"]
                                break
                            except:
                                DDR_print_mesg("No dumpfiles found in server response.",2)
                                dump_files = None
                                ret = False
                                break
                        else:
                            # API return status not 200
                            if "Sample file not found" in res.json()["return_status"]:
                                DDR_print_mesg("Sample file not found on DDR server, trying to upload it...")
                                if DDR_upload_sample(idaapi.get_input_file_path()):
                                    DDR_print_mesg("File {} successfully uploaded.".format(idaapi.get_input_file_path()))
                                continue

                            # other error than sample file not found
                            json_res = res.json()
                            idaapi.warning("Failed dumping buffer on server. See server log for details.")
                            DDR_print_mesg("Failed dumping buffer on server. API return_status = {}".format(json_res["return_status"]),2)
                            return False

                if dump_files:
                    DDR_print_mesg("Successfully dumped buffer on server.")
                    t = datetime.now()
                    timestamp = t.strftime("%d%b%y_%H%M%S")
                    # Prepare filename lists
                    for num, d_file in enumerate(dump_files):
                        df_addr = str(res.json()["fileaddrs"][num])
                        df_size = str(res.json()["filesizes"][num])
                        dump_filename = "dump_{}_{:d}_{}_{}_{}.bin".format(os.path.splitext(SAMPLE_FILENAME)[0], num, df_addr, df_size, timestamp)
                        tmp_file_list.append(d_file)
                        dump_filename_list.append(dump_filename)

                    # Download the dumped file(s)
                    for num, d_file in enumerate(tmp_file_list):
                        dump_para_json['id'] = 9
                        dump_para_json['dl_file'] = [d_file]

                        DDR_print_mesg("Start downloading file {}...".format(d_file))
                        DDR_print_mesg("Sending following config to server:\n{}".format(json.dumps(dump_para_json,indent=4, sort_keys=True)))

                        with requests.post(url, verify=req_verify, json=dump_para_json, timeout=DDR_CONFIG_SETTINGS.MAX_API_TIMEOUT, stream=True) as res:
                            dump_para_json['dl_file'] = []
                            #Success. Response code 200
                            if res.status_code == 200:
                                dump_filename = dump_filename_list[num]
                                dump_filename = ida_kernwin.ask_file(1,dump_filename,"Save dumped buffer to:")
                                if dump_filename:
                                    # Download file
                                    with open(dump_filename, 'wb') as f:
                                        for chunk in res.iter_content(chunk_size=1024): 
                                            if chunk: 
                                                f.write(chunk)
                                else:
                                    DDR_print_mesg("File '{}' not saved. Dialog was canceled.".format(d_file),1)
                                    
                            # Download failed. Response code is not 200
                            else:
                                DDR_print_mesg("Failed downloading {} from server".format(d_file))
                                return False
                    ret = True
                
                # try to delete tmp. buffer files on server
                if tmp_file_list:
                    dump_para_json['id'] = 6
                    dump_para_json['filelist2del'] = tmp_file_list

                    with requests.post(url, verify=req_verify, json=dump_para_json, timeout=DDR_CONFIG_SETTINGS.MAX_API_TIMEOUT, stream=False) as res:
                        if res.status_code == 200:
                            DDR_print_mesg("Temp. buffer file(s) on server successfully deleted")
                            dump_para_json['filelist2del'] = []
                        else:
                            DDR_print_mesg("Failed to delete temp. buffer file(s) on server",2)
                            dump_para_json['filelist2del'] = []
                            ret = False

                return ret

            # ----- Run trace -----
            if cmd_id == 7:
                ret = False
                dump_para_json['id'] = cmd_id
               
                for upload_counter in range(DDR_CONFIG_SETTINGS.MAX_UPLOAD_ATTEMPTS, -1, -1):

                    if upload_counter < 1:
                        DDR_print_mesg("Failed to upload sample file to server. Max number of upload attempts reached",2)
                        return False
                    
                    hash_list = []
                    hash_list.append(json.dumps(dump_para_json['trace_start']))
                    hash_list.append(json.dumps(dump_para_json['trace_end']))
                    hash_list.append(json.dumps(dump_para_json['trace_light']))
                    hash_str = "".join(hash_list)
      
                    trace_filename     = "trace_tmp_" + self.get_hash_from_str(hash_str)["md5_sum"] + ".json"
                    trace_filename_api = os.path.splitext(trace_filename)[0] + "_apicalls.json"
                    zipfilename        = os.path.splitext(trace_filename)[0] + ".zip"
                    
                    with requests.post(url, verify=req_verify, json=dump_para_json, timeout=DDR_CONFIG_SETTINGS.MAX_API_TIMEOUT, stream=True) as res:
                        if res.status_code == 200:
                            # Download file
                            with open(zipfilename, 'wb') as f:
                                for chunk in res.iter_content(chunk_size=1024): 
                                    if chunk: 
                                        f.write(chunk)
                            DDR_print_mesg("Analysis done. Downloaded zip file from server: {}".format(zipfilename))
                            break

                        # Error on the server side. Response is not 200        
                        else:
                            # Sample not found on server
                            if "Sample file not found" in res.json()["return_status"]:
                                DDR_print_mesg("Sample file not found on DDR server, trying to upload it...")
                                
                                if DDR_upload_sample(idaapi.get_input_file_path()):
                                    DDR_print_mesg("File {} successfully uploaded.".format(idaapi.get_input_file_path()))
                            # Other error on server      
                            else:
                                DDR_print_mesg("Failed to generate trace on server, please check server side for details.",2)
                                return False
                
                # Verify zip file is valid
                if not self._test_zip_is_valid(zipfilename):
                    DDR_print_mesg("Zip file verification failed for full trace", 2)
                    return False
                    
                # Extract files from zip archive
                with zipfile.ZipFile(zipfilename, "r") as ziparchive:
                    DDR_print_mesg("Extracting files to {}".format(JSONFILE_DIR))
                    ziparchive.extractall(JSONFILE_DIR)

                # Delete zip archive
                DDR_print_mesg("Delete archive: {}".format(zipfilename))
                try:
                    os.remove(zipfilename)
                except:
                    DDR_print_mesg("Failed to delete zip file: {}".format(zipfilename), 1)
                    pass

                # load and parse downloaded JSON trace file and the API file 
                if self._load_JSON_file(filename=trace_filename):
                    DDR_print_mesg("JSON file succesfully parsed: {}".format(trace_filename))
                    if self._load_APICalls_file(filename=trace_filename_api):
                        # Operation successful. We got valid files from the server.
                        ret = True
                    else:
                        ret = False
                else:
                    ret = False
                  
                # Delete tmp files on server
                dump_para_json['id'] = 8
                
                with requests.post(url, verify=req_verify, json=dump_para_json, timeout=DDR_CONFIG_SETTINGS.MAX_API_TIMEOUT, stream=False) as res:
                    DDR_print_mesg("Told server to delete local temp. files. Respondscode={:d} API status:{}".format(res.status_code, res.json()["return_status"]))
                    
                return ret

            # ----- Run sample only -----
            if cmd_id == 10:
                dump_para_json['id'] = cmd_id

                for upload_counter in range(DDR_CONFIG_SETTINGS.MAX_UPLOAD_ATTEMPTS, -1, -1):
                    with requests.post(url, verify=req_verify, json=dump_para_json, timeout=DDR_CONFIG_SETTINGS.MAX_API_TIMEOUT, stream=False) as res:
                        if res.status_code == 200:
                            DDR_print_mesg("Successfully executed sample on server",2)
                            ret = True
                            break
                        else:
                            if "Sample file not found" in res.json()["return_status"]:
                                DDR_print_mesg("Sample file not found on DDR server, trying to upload it...")
                                if DDR_upload_sample(idaapi.get_input_file_path()):
                                    DDR_print_mesg("File {} successfully uploaded.".format(idaapi.get_input_file_path()))
                            # all other errors
                            DDR_print_mesg("Failed to execute sample on server",2)
                            ret = False

                return ret

            # ----- Create patched binary with endless loop at ea (without DynRio) -----
            if cmd_id == 11:
                dump_para_json['id'] = cmd_id

                for upload_counter in range(DDR_CONFIG_SETTINGS.MAX_UPLOAD_ATTEMPTS, -1, -1):
                    with requests.post(url, verify=req_verify, json=dump_para_json, timeout=DDR_CONFIG_SETTINGS.MAX_API_TIMEOUT, stream=False) as res:
                        if res.status_code == 200:
                            DDR_print_mesg("Response: {}".format(res.json()["return_status"]))
                            ret = True
                            break
                        else:
                            if "Sample file not found" in res.json()["return_status"]:
                                DDR_print_mesg("Sample file not found on DDR server, trying to upload it...")
                                if DDR_upload_sample(idaapi.get_input_file_path()):
                                    DDR_print_mesg("File {} successfully uploaded.".format(idaapi.get_input_file_path()))
                            # all other errors
                            DDR_print_mesg("Failed to execute sample on server",2)
                            ret = False

                return ret

            # run sample in x64dbg (create x64dbg script)
            if cmd_id == 12:
                dump_para_json['id'] = cmd_id

                for upload_counter in range(DDR_CONFIG_SETTINGS.MAX_UPLOAD_ATTEMPTS, -1, -1):
                    with requests.post(url, verify=req_verify, json=dump_para_json, timeout=DDR_CONFIG_SETTINGS.MAX_API_TIMEOUT, stream=False) as res:
                        if res.status_code == 200:
                            DDR_print_mesg("Response: {}".format(res.json()["return_status"]))
                            ret = True
                            break
                        else:
                            if "Sample file not found" in res.json()["return_status"]:
                                DDR_print_mesg("Sample file not found on DDR server, trying to upload it...")
                                if DDR_upload_sample(idaapi.get_input_file_path()):
                                    DDR_print_mesg("File {} successfully uploaded.".format(idaapi.get_input_file_path()))
                            # all other errors
                            DDR_print_mesg("Failed to execute sample on server",2)
                            ret = False

                return ret


        except requests.exceptions.SSLError as e:
            DDR_print_mesg("[API SSL ERROR]: {}".format(e),2)
            DDR_print_mesg("Did you configure a valid serial number for the certificate ?",2)
            DDR_print_mesg("For troubleshooting install the certificate from the <IDA plugin dir>/ddr into your browser.",2)
            DDR_print_mesg("Access the DDR server test webpage with your browser and check the error message.",2)
            url = "https://" + DDR_CONFIG_SETTINGS.WEBSERVER + ":" + DDR_CONFIG_SETTINGS.WEBSERVER_PORT + "/"
            DDR_print_mesg("Reminder: Test webpage is {}".format(url),2)
            ret = False

        except requests.exceptions.ConnectionError as e:
            DDR_print_mesg("[API CONNECTION ERROR]: {}".format(e),2)
            DDR_print_mesg("Did you start and configure the DDR server script or is there any firewall inbetween ?",2)
            ret = False

        except requests.exceptions.HTTPError as e:
            DDR_print_mesg("[API HTTP ERROR]: {}".format(e),2)
            DDR_print_mesg("\nCMD execution failed. Check server side for details",2)
            ret = False

        except requests.exceptions.ReadTimeout as e:
            DDR_print_mesg("[API TIMEOUT ERROR]: {}".format(e),2)
            ret = False

        except ValueError as e:
            DDR_print_mesg("[API JSON ERROR] Failed to decode the returned JSON data from server: {}".format(e), 2)
            ret = False

        except:
            DDR_print_mesg("[API UNKNOWN ERROR] Unkown error happend. REST API request failed.",2)
            raise

        return ret

class ddrPlugin(idaapi.plugin_t):
    """
    The IDA plugin stub for DDR.
    """

    #
    # Plugin flags:
    # - PLUGIN_MOD: DDR may modify the database
    # - PLUGIN_PROC: Load/unload DDR when an IDB opens / closes
    # - PLUGIN_HIDE: Hide DDR from the IDA plugin menu
    #
    flags = idaapi.PLUGIN_PROC | idaapi.PLUGIN_MOD 
    comment = "Dynamic Data Resolver Plugin"
    help = ""
    wanted_name = "DDR"
    wanted_hotkey = ""

    def activate(self,ctx):
        if sys.version_info > (3, 7):
            return 0
        return 1

    def init(self):
        """
        This is called by IDA when it is loading the plugin.
        """     
        global REG_LIST
        global JSONFILE_DIR
        global DDR_API_V1_CONFIG_SETTINGS
        global SAMPLE_FILENAME
        global SAMPLE_SHA256
        global SAMPLE_DIR
            
        if not sys.version_info >= (3, 0):
            print("[DDR][WARNING] This script only supports Python 3. Install Python 3 64bit and use idapyswitch.exe to enable Python 3 for IDA.")
            print("[DDR][WARNING] DDR only support Python 3.7 and below in the moment due to http://www.hexblog.com/?cat=3.")
        
        if idaapi.IDA_SDK_VERSION < 750:
            if sys.version_info >= (3, 8):
                print("[DDR][ERROR] ------------------------------------------------------------------------------------------------")
                print("[DDR][ERROR] In IDA version below 7.5, the DDR plugin only supports Python 3.7 64 bit due to")
                print("[DDR][ERROR] http://www.hexblog.com/?cat=3. It is fine to use 3.8 on the DDR server side.")
                print("[DDR][ERROR] Unloading DDR plugin.")
                print("[DDR][ERROR] ------------------------------------------------------------------------------------------------")
                idaapi.warning("[DDR] Failed to load DDR plugin. Wrong Python version detected. Please see log for more details.")
                return idaapi.PLUGIN_SKIP   

        if DDR_CONFIG_SETTINGS.is_initalized == "PHASE1":
            print("[DDR][INFO] Configuration found.")
        else:
            print("[DDR][ERROR] Failed to initialze DDR configuration. Unloading Plugin.")
            return idaapi.PLUGIN_SKIP

        DDR_print_mesg("Activating plugin...")

        info = idaapi.get_inf_structure()
        orgfile = ida_nalt.get_input_file_path()

        try:
            SAMPLE_FILENAME         = idaapi.get_root_filename()
            #SAMPLE_SHA256           = ida_nalt.retrieve_input_file_sha256().decode()  # buggy itm, cuts the last byte of the hash
            SAMPLE_SHA256           = get_hash(orgfile)['sha256_sum']

            if orgfile:
                SAMPLE_DIR              = os.path.dirname(os.path.abspath(orgfile)) 
                DDR_CONFIG_SETTINGS.DUMP_CFG_FILE_WITHPATH  = SAMPLE_DIR + DDR_CONFIG_SETTINGS.DUMP_CFG_FILE

                if info.is_64bit():
                    DDR_CONFIG_SETTINGS.ARCH_BITS = 64
                    DDR_print_mesg("We are in a 64bit world...",0)
                    REG_LIST = ["xax","xbx","xcx","xdx","xsp","xbp","xsi","xdi","r8","r9","r10","r11","r12","r13","r14","r15"]
                elif info.is_32bit():
                    DDR_CONFIG_SETTINGS.ARCH_BITS = 32
                    DDR_print_mesg("We are in a 32bit world...",0)
                    REG_LIST = ["xax","xbx","xcx","xdx","xsp","xbp","xsi","xdi"]
                else:
                    #ERROR
                    DDR_CONFIG_SETTINGS.ARCH_BITS = 16
                    DDR_print_mesg("This plugin only supports 64/32bit enviroments and PE files. Unloading plugin.",2)
                    return idaapi.PLUGIN_SKIP

                DDR_API_V1_CONFIG_SETTINGS  = ddr_api_v1_cfg(BB_LIST)

            else:
                DDR_print_mesg("No input file loaded yet. Pls load a sample file into IDA")
                return idaapi.PLUGIN_SKIP
        except:
            DDR_print_mesg("Initalizing plugin failed.",2)
            raise
            return idaapi.PLUGIN_SKIP

        try:
            JSONFILE_DIR  = os.path.dirname(os.path.abspath(idaapi.get_input_file_path()))
        except:
            DDR_print_mesg("No input file loaded. Pls load a sample file into IDA")

        if not orgfile:
            DDR_print_mesg("No PE file loaded. Unloading plugin.",1)
            return idaapi.PLUGIN_SKIP

        try:
            pe = pefile.PE(orgfile)
        except:
            DDR_print_mesg("This plugin only supports PE files. Unloading plugin.",1)
            return idaapi.PLUGIN_SKIP

        if pe.is_dll():
            DDR_print_mesg("This plugin does not support DLLs in the moment.",1)
            return idaapi.PLUGIN_SKIP

        if not pe.is_exe():
            DDR_print_mesg("This plugin only supports PE files. Unloading plugin.",1)
            return idaapi.PLUGIN_SKIP
        
        # Register Menu Actions
        for m in menu_items:
            if DDR_CONFIG_SETTINGS.ARCH_BITS == 32:
                if menu_items[m]["x64only"]:
                    continue

            self._add_action(m, menu_items[m]["menu_str"], menu_items[m]["hotkey"], "DDR Tool", DDR_ida_action_handler(menu_items[m]["ah_id"]))

        # attach action to 'File/Load file/' menu
        result = idaapi.attach_action_to_menu(
            "File/Load file/",      # Relative path of where to add the action
            "DDR_Action_Load_file", # The action ID (see above)
            idaapi.SETMENU_APP      # We want to append the action after ^
        )
        if not result:
            idaapi.warning("[DDR] Failed to attach action (DDR_Action_Load_file) to main menu")
        else:
            DDR_print_mesg("Attached [DDR_Action_Load_file] to main menu", 7)

        # attach action(s) to the toolbar
        if idaapi.attach_action_to_toolbar("AnalysisToolBar", "DDR_Action_Load_file"):
            DDR_print_mesg("Action [DDR_Action_Load_file] attached to toolbar.", 7)
        else:
            idaapi.warning("[DDR] Failed to attach action [DDR_Action_Load_file] to toolbar.")


        # attach action to 'File/Load file/' menu
        result = idaapi.attach_action_to_menu(
            "File/Load file/",          # Relative path of where to add the action
            "DDR_Action_Load_api_file", # The action ID (see above)
            idaapi.SETMENU_APP          # We want to append the action after ^
        )
        if not result:
            idaapi.warning("[DDR] Failed to attach action (DDR_Action_Load_api_file) to main menu")
        else:
            DDR_print_mesg("Attached [DDR_Action_Load_api_file] to main menu", 7)

        # attach action(s) to the toolbar
        if idaapi.attach_action_to_toolbar("AnalysisToolBar", "DDR_Action_Load_api_file"):
            DDR_print_mesg("Action [DDR_Action_Load_api_file] attached to toolbar.", 7)
        else:
            idaapi.warning("[DDR] Failed to attach action [DDR_Action_Load_api_file] to toolbar.")

        # attach action(s) to context menu
        class Hooks(idaapi.UI_Hooks):
            #def finish_populating_tform_popup(self, form, popup):
            def finish_populating_widget_popup(self, form, popup):
                # register all menu items from global menu_item dict
                for m in menu_items:

                    if menu_items[m]["hide_in_context"]:
                        continue

                    if DDR_CONFIG_SETTINGS.ARCH_BITS != 64 and menu_items[m]["x64only"]:
                        continue

                    self._add_ida_popup(m, form, popup, "DDR/" + menu_items[m]["submenu"])

            def _add_ida_popup(self, act_name, form, popup, menuname):
                """
                Helper function to add a menu to the context menu in IDA if we are in the Disasm view
                """
                
                if ida_kernwin.get_widget_type(form) == idaapi.BWN_DISASM:
                    DDR_print_mesg("we are in DISASM view.",7)
                    if (idaapi.attach_action_to_popup(form, popup, act_name, menuname)):
                        DDR_print_mesg("added [{}] to context menu.".format(act_name),7)
                        pass
                    else:
                        idaapi.warning("[DDR] Failed to add [{}] to context menu.".format(act_name))

                else:
                   DDR_print_mesg("we are NOT in DISASM view.",7)
                   return False

                return True

        self.hooks = Hooks()
        if (self.hooks.hook()):
            DDR_print_mesg("All actions attached to context menu.",7)
        else:
            idaapi.warning("[DDR] Hooking context menu failed.")

        try:
            DDR_print_mesg("Done. DDR plugin started and initalized.")
            DDR_CONFIG_SETTINGS.is_initalized = "FULL"
            #logger.info("DDR plugin started")
            #DDR_print_mesg("Using Logfile: " + ddr_logfilename)
  
        except Exception as e:
            idaapi.warning("[DDR] init failed.")
            #logger.exception("Exception details:")
            return idaapi.PLUGIN_SKIP

        if not os.path.isfile(DDR_CONFIG_SETTINGS.CA_CERT):
            DDR_print_mesg("CA certificate file not found. Pls check the DDR_CONFIG_SETTINGS.CA_CERT variable in the ddr_plugin.py script. Fix it and restart IDA.",3)
            idaapi.warning("[DDR] CA file not found. Pls check the DDR_CONFIG_SETTINGS.CA_CERT variable in the ddr_plugin.py script. Fix it and restart IDA.")
            return idaapi.PLUGIN_SKIP

        # tell IDA to keep the plugin loaded (everything is okay)
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        """
        This is called by IDA when this file is loaded as a script.
        """
        idaapi.warning("[DDR] This plugin cannot be run as a script in IDA.")

    def update(self, ctx):
        return idaapi.AST_ENABLE_FOR_FORM if ctx.form_type == idaapi.BWN_DISASM else idaapi.AST_DISABLE_FOR_FORM

    def term(self):
        """
        This is called by IDA when it is unloading the plugin.
        """
        try:
            # unregister all menues. ToDo: check if action is registered
            if DDR_CONFIG_SETTINGS.is_initalized == "FULL":
                for m in menu_items:            

                    if DDR_CONFIG_SETTINGS.ARCH_BITS == 16:    # actions are only registered for 32/64 bit enviroments at init 
                        break

                    if DDR_CONFIG_SETTINGS.ARCH_BITS != 64 and menu_items[m]["x64only"]:
                        continue

                    self._unregister_ida_action(m)

            DDR_print_mesg("plugin terminated.",7)
        except Exception as e:
            idaapi.warning("[DDR] Failed to cleanly unload DDR from IDA.")
            #logger.exception("Failed to cleanly unload DDR from IDA.")

    def _unregister_ida_action(self, act_name):
        """
        Helper function to unregister an IDA action
        """

        if idaapi.unregister_action(act_name):
            DDR_print_mesg("[{}] action unregistered from IDA.".format(act_name),7)
        else:
            idaapi.warning("Failed to unregister ({action_name}) action in IDA.".format(action_name=act_name)) 

    def _add_action(self, act_name, act_text, act_shortcut, act_tooltip, act_handler):
        """
        Helper function to add an action description to IDA
        """
        action_desc = idaapi.action_desc_t(
            act_name,                  # The action name
            act_text,                  # The action text
            act_handler,               # The action handler
            act_shortcut,              # Optional: action shortcut
            act_tooltip                # Optional: tooltip
                                       # Optional: the action icon
        )

        # register the action with IDA
        result = idaapi.register_action(action_desc)
        if not result: 
            idaapi.warning("[DDR] Failed to register action({action_name}) with IDA".format(action_name=act_name))

class DDR_MyChoose(idaapi.Choose):
    """ Chooser2 class to display Strings and API calls views"""
    def __init__(self, title, items, flags=0, width=None, height=None, embedded=False, modal=False, columns=False):
        
        self.lines_n = []

        idaapi.Choose.__init__(
            self,
            title,
            columns,
            flags = flags,
            width = width,
            height = height,
            embedded = embedded)
        self.n = 0
        self.items = items
        self.icon = -1
        #self.selcount = 0
        self.modal = modal
        #self.popup_names = ["Inzert", "Del leet", "Ehdeet", "Ree frech"]

    def OnClose(self):
        # Window closed
        DDR_print_mesg("Trace View window closed" , 7)
        pass

    def OnEditLine(self, n):
        # Click on "Edit.. menu"
        DDR_print_mesg("Not implemented.")

    def OnInsertLine(self):
        # Click on "Insert.. menu"
        DDR_print_mesg("Not implemented.")

    def OnDeleteLine(self, n):
        # Click on "Delete menu"
        del self.items[n]
        DDR_print_mesg("Line {:d} deleted.".format(n) , 7)
        return n

    def OnRefresh(self, n):
        # Click on "Refresh menu"
        DDR_print_mesg("Trace View window refreshed {:d}".format(n), 7)
        return n

    def OnSelectLine(self, n):
        # double click on line
        self.lines_n.append(n)
        DDR_print_mesg("Line: {:d}  Address: {}  Name:{}".format(n, str(self.items[n][0]), str(self.items[n][1])), 7)
        ea = int(self.items[n][1],16) 
        ida_kernwin.jumpto(ea)

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

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

    def OnGetIcon(self, n):
        r = self.items[n]
        t = self.icon + r[1].count("*")
        return t

    def OnGetLineAttr(self, n):
        # Highlight line 1 in blue:
        if n in self.lines_n:
            self.highlight_line = False
            return [0xFF0000, 0]
        
        return

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


# --- For future use Tree View Class  ---
#
#  tree_window = CBaseTreeViewer()
#  tree_window.Show("TreeView Tab Name")
#
class DDR_CBaseTreeViewer(idaapi.PluginForm):
    def populate_tree(self):
        # Clear previous items
        self.tree.clear()

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

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

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

        TreeBranch1 = QtWidgets.QTreeWidgetItem(self.tree)
        TreeBranch1.setText(0, "TreeBranch1")
        TreeBranch1.ea = idc.BADADDR
        leaf1 = QtWidgets.QTreeWidgetItem(TreeBranch1)
        leaf1.setText(0, "{} {} {}".format("Leaf1.0","Leaf1.1","Leaf1.2"))
        leaf1.ea = idc.BADADDR
        leaf2 = QtWidgets.QTreeWidgetItem(TreeBranch1)
        leaf2.setText(0, "{} {} {}".format("Leaf2.0","Leaf2.1","Leaf2.2"))
        leaf2.ea = idc.BADADDR
        TreeBranch2 = QtWidgets.QTreeWidgetItem(self.tree)
        TreeBranch2.setText(0, "TreeBranch2")
        leaf3 = QtWidgets.QTreeWidgetItem(TreeBranch2)
        leaf3.setText(0, "{} {} {}".format("Leaf3.0","Leaf3.1","Leaf3.2"))
        leaf3.ea = idc.BADADDR

        # Populate PluginForm
        self.parent.setLayout(layout)

    def Show(self, title):
        return idaapi.PluginForm.Show(self, title, options = idaapi.PluginForm.FORM_PERSIST)


