# plugins

repository: shemesh999/oregami
file: offreg_plugin.py
description: IDA plugins and scripts for analyzing register usage frame
wanted_name: OffReg
comment: OffReg
created_at: 2020-03-16 11:25:54+00:00
pushed_at: 2023-06-08 14:27:51+00:00
url:
https://github.com/shemesh999/oregami/blob/eee466005af08c816275e3269a83704c653a5312/offreg_plugin.py

content:
from oregami.reg_utils import *
from oregami.reg_type import rf_settype, get_type_from_user
from oregami.reg_frame import RegFrame
import ida_offset

class OffRegPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_PROC
    comment = "OffReg"
    help = "Set offset for regs in their usage frame - only when " \
           "used as a specific variable"
    wanted_name = "OffReg"
    wanted_hotkey = "Shift+R"

    @staticmethod
    def init():
        return idaapi.PLUGIN_OK

    @staticmethod
    def term():
        pass

    @staticmethod
    def run(arg):
        start, _ = sark.get_selection()

        offreg_plugin_starter(start)


def PLUGIN_ENTRY():
    return OffRegPlugin()


def rf_setoff(rf, off_ea):
    for insn in rf.get_noinit_instructions():
        done_offset = False
        need_offset = False
        for opnd in insn.operands:
            if opnd.uf_is_read and (not opnd.uf_is_write) and \
                    (not opnd.uf_is_implicit):
                need_offset = True
                print('Setting offset {:x} for {:x} operand #{}'.format(off_ea, insn.ea, opnd.n))
                if opnd.type.name == 'Memory_Displacement':
                    ida_offset.op_offset(insn.ea, opnd.n, idc.REFINFO_NOBASE | idc.REF_OFF32, idc.BADADDR, off_ea)
                    done_offset = True
        if need_offset and (not done_offset):
            # probably another operand is an immediate value which needs this to be applied to it. May have false positives
            for opnd in insn.operands:
                if opnd.type.name == 'Immediate_Value':
                    ida_offset.op_offset(insn.ea, opnd.n, idc.REFINFO_NOBASE | idc.REF_OFF32, idc.BADADDR, off_ea)
                    break
                    


def offreg_plugin_starter(orig_ea):
    canon_list = conf.proc.get_reg_list()
    # print canon_list
    reg = get_reg_from_cursor(orig_ea, canon_list)

    if reg is None:
        # Ask for user input - may be used to look for a reg influencing
        #   the line - even if it doesn't exist on the line
        reg_idx = RegChoose(orig_ea, canon_list).Show(True)
        if reg_idx >= 0:
            reg = canon_list[reg_idx]
        else:
            return

    reg = RegName(orig_ea, canon_list).canon(reg)

    if reg is None:
        return

    # Get type name
    off_ea = ida_kernwin.ask_addr(0, 'Choose offset')
    if off_ea is None:
        return
        
    # global conf
    rf = RegFrame(orig_ea, reg, force=(not conf.cache_bool))
    rf_setoff(rf, off_ea)


repository: benwoo1110/PMA-Notes
file: Plugins/noping.py
wanted_name: Easy Nop
comment: Easy nopping tool for you nopping needs
created_at: 2022-09-27 01:21:48+00:00
pushed_at: 2022-12-10 06:36:38+00:00
url:
https://github.com/benwoo1110/PMA-Notes/blob/9d322ef0cbb99709f9227e322846b58dd1b1c9ab/Plugins/noping.py

content:
"""
Adapted from https://www.unknowncheats.me/forum/general-programming-and-reversing/343177-easy-nop-ida-plugin.html
"""

import idaapi
 
popup_action_names = []
 
class Hooks(idaapi.UI_Hooks):
    def __init__(self):
        idaapi.UI_Hooks.__init__(self)
 
    def finish_populating_widget_popup(self, form, popup):
        global popup_action_names
        widget_type = idaapi.get_widget_type(form)
        if widget_type == idaapi.BWN_DISASM:
            for action_name in popup_action_names:
                idaapi.attach_action_to_popup(form, popup, action_name, None)
 
class EasyNopPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_HIDE | idaapi.PLUGIN_PROC
    comment = "Easy nopping tool for you nopping needs"
    help = "Use Shift+N to nop out current line or selection"
    wanted_name = "Easy Nop"
    wanted_hotkey = ""
 
    def init(self):
        idaapi.msg("Easy Nop :: init\n")
        AddToPopup('nopaction:add_action', 'Nop Selection', NopAction(), 'Shift+N', None)
 
        self.hooks = Hooks()
        self.hooks.hook()
        return idaapi.PLUGIN_KEEP
 
    def run(self):
        idaapi.msg("Easy Nop :: run\n")
 
    def term(self):
        idaapi.msg("Easy Nop :: term\n")
        if self.hooks:
            self.hooks.unhook()
 
        idaapi.unregister_action('nopaction:add_action')
 
class EA_Action(idaapi.action_handler_t):
    def __init__(self):
        idaapi.action_handler_t.__init__(self)
 
    def update(self, ctx):
        return idaapi.AST_ENABLE_FOR_WIDGET if ctx.form_type == idaapi.BWN_DISASM else idaapi.AST_DISABLE_FOR_WIDGET
 
class NopAction(EA_Action):
    def activate(self, ctx):
        info = idaapi.get_inf_structure()
        result, start, end = idaapi.read_range_selection(None)
 
        if result == False:
            start = idaapi.get_screen_ea()
 
            if start == idaapi.BADADDR:
                print('Easy Nop :: Screen EA == idaapi.BADADDR')
                return 0
 
            end = start + idaapi.get_item_size(start)
        else:
            end += idaapi.get_item_size(end)
 
        if start == idaapi.BADADDR:
            print('Easy Nop :: Selection EA == idaapi.BADADDR')
            return 0
 
        if start == end:
            print('Easy Nop :: Nothing to nop')
            return 0
 
        for x in range(start, end):
            # Maybe theres a smarter way to get the nop value for different archs e.g. Assemble('nop') -> 0x90
            idaapi.patch_byte(x, 0x90)
 
        for x in range(start + 1, end):
            idaapi.hide_item(x)
 
        # Must do this else it bugs out on 2x 1 byte instructions being nopped
        idaapi.hide_item(start)
        idaapi.unhide_item(start)
 
        # Search for hidden nops and add to count
        while idaapi.get_byte(end) == 0x90 and idaapi.is_hidden_item(end) == True:
            end += 1
 
        count = end - start
 
        if count > 1:
            idaapi.set_cmt(start, "truncated nops (%d)" % (count), False)
 
        print(end)
        print(start)
 
        return 1
 
def AddToPopup(action_name, display, handler, shortcut, tooltip, icon=None):
    global popup_action_names
 
    if tooltip == None:
        tooltip = action_name
 
    if idaapi.register_action(idaapi.action_desc_t(
                action_name,
                display,
                handler,
                shortcut,
                tooltip
            )):
 
        popup_action_names.append(action_name)
    else:
        print('Easy Nop :: Error registering action %s' % (action_name))
 
def PLUGIN_ENTRY(*args, **kwargs):
    return EasyNopPlugin()


repository: williballenthin/idawilli
file: plugins/hint_calls/hint_calls.py
description: IDA Pro resources, scripts, and configurations
wanted_name: CallsHintsPlugin
comment: Display the calls and strings referenced by a function as hints.
created_at: 2016-12-14 02:23:46+00:00
pushed_at: 2025-05-12 10:45:42+00:00
url:
https://github.com/williballenthin/idawilli/blob/063417207eebd061fbd58f6441e6b5e9ffc43f8e/plugins/hint_calls/hint_calls.py

content:
'''
IDA plugin to display the calls and strings referenced by a function as hints.

Installation: put this file in your %IDADIR%/plugins/ directory.
Author: Willi Ballenthin <william.ballenthin@fireeye.com>
Licence: Apache 2.0
'''
import logging

import idc
import idaapi
import idautils
import ida_xref
import ida_nalt
import ida_bytes
import ida_lines
import ida_kernwin
import re


logger = logging.getLogger(__name__)
DEFAULT_IMPORTANT_LINES_NUM = 5


###############################################################################
#### begin: idapython_lex_curline
# via: https://gist.github.com/williballenthin/466eb28679d30e212ffac57e4a9ceaa5
# note: inline here for simplicity


# inverse mapping of color value to name.
# ref: https://www.hex-rays.com/products/ida/support/sdkdoc/group___s_c_o_l_o_r__.html#ga6052470f86411b8b5ffdf4af4bbee225
INV_COLORS = {
    0x1: 'COLOR_DEFAULT',  #= 0x01,         // Default
    0x2: 'COLOR_REGCMT',   #= 0x02,         // Regular comment
    0x3: 'COLOR_RPTCMT',   #= 0x03,         // Repeatable comment (comment defined somewhere else)
    0x4: 'COLOR_AUTOCMT',  #= 0x04,         // Automatic comment
    0x5: 'COLOR_INSN',     #= 0x05,         // Instruction
    0x6: 'COLOR_DATNAME',  #= 0x06,         // Dummy Data Name
    0x7: 'COLOR_DNAME',    #= 0x07,         // Regular Data Name
    0x8: 'COLOR_DEMNAME',  #= 0x08,         // Demangled Name
    0x9: 'COLOR_SYMBOL',   #= 0x09,         // Punctuation
    0xa: 'COLOR_CHAR',     #= 0x0A,         // Char constant in instruction
    0xb: 'COLOR_STRING',   #= 0x0B,         // String constant in instruction
    0xc: 'COLOR_NUMBER',   #= 0x0C,         // Numeric constant in instruction
    0xd: 'COLOR_VOIDOP',   #= 0x0D,         // Void operand
    0xe: 'COLOR_CREF',     #= 0x0E,         // Code reference
    0xf: 'COLOR_DREF',     #= 0x0F,         // Data reference
    0x10: 'COLOR_CREFTAIL', #= 0x10,         // Code reference to tail byte
    0x11: 'COLOR_DREFTAIL', #= 0x11,         // Data reference to tail byte
    0x12: 'COLOR_ERROR',    #= 0x12,         // Error or problem
    0x13: 'COLOR_PREFIX',   #= 0x13,         // Line prefix
    0x14: 'COLOR_BINPREF',  #= 0x14,         // Binary line prefix bytes
    0x15: 'COLOR_EXTRA',    #= 0x15,         // Extra line
    0x16: 'COLOR_ALTOP',    #= 0x16,         // Alternative operand
    0x17: 'COLOR_HIDNAME',  #= 0x17,         // Hidden name
    0x18: 'COLOR_LIBNAME',  #= 0x18,         // Library function name
    0x19: 'COLOR_LOCNAME',  #= 0x19,         // Local variable name
    0x1A: 'COLOR_CODNAME',  #= 0x1A,         // Dummy code name
    0x1B: 'COLOR_ASMDIR',   #= 0x1B,         // Assembler directive
    0x1C: 'COLOR_MACRO',    #= 0x1C,         // Macro
    0x1D: 'COLOR_DSTR',     #= 0x1D,         // String constant in data directive
    0x1E: 'COLOR_DCHAR',    #= 0x1E,         // Char constant in data directive
    0x1F: 'COLOR_DNUM',     #= 0x1F,         // Numeric constant in data directive
    0x20: 'COLOR_KEYWORD',  #= 0x20,         // Keywords
    0x21: 'COLOR_REG',      #= 0x21,         // Register name
    0x22: 'COLOR_IMPNAME',  #= 0x22,         // Imported name
    0x23: 'COLOR_SEGNAME',  #= 0x23,         // Segment name
    0x24: 'COLOR_UNKNAME',  #= 0x24,         // Dummy unknown name
    0x25: 'COLOR_CNAME',    #= 0x25,         // Regular code name
    0x26: 'COLOR_UNAME',    #= 0x26,         // Regular unknown name
    0x27: 'COLOR_COLLAPSED',#= 0x27,         // Collapsed line

    #  // Fictive colors
    0x28: 'COLOR_ADDR',     #= 0x28, // hidden address marks
                            #        // The address is represented as 8digit
                            #        // hex number: 01234567
                            #        // It doesn't have COLOR_OFF pair
                            #        // NB: for 64-bit IDA, the address is 16digit

    0x29: 'COLOR_OPND1',    #= COLOR_ADDR+1, // Instruction operand 1
    0x2A: 'COLOR_OPND2',    #= COLOR_ADDR+2, // Instruction operand 2
    0x2B: 'COLOR_OPND3',    #= COLOR_ADDR+3, // Instruction operand 3
    0x2C: 'COLOR_OPND4',    #= COLOR_ADDR+4, // Instruction operand 4
    0x2D: 'COLOR_OPND5',    #= COLOR_ADDR+5, // Instruction operand 5
    0x2E: 'COLOR_OPND6',    #= COLOR_ADDR+6, // Instruction operand 6

    0x32: 'COLOR_UTF8',     #= COLOR_ADDR+10;// Following text is UTF-8 encoded
}


class Symbol(object):
    def __init__(self, type):
        super(Symbol, self).__init__()
        self.type = type

    def __str__(self):
        raise NotImplementedError()


class StringSymbol(Symbol):
    def __init__(self, string):
        super(StringSymbol, self).__init__('string')
        self.string = string

    def __str__(self):
        return 'STRING=' + self.string


class ColorOnSymbol(Symbol):
    def __init__(self, color):
        super(ColorOnSymbol, self).__init__('coloron')
        self.color = ord(color)

    def __str__(self):
        return 'COLORON=' + INV_COLORS[self.color]


class ColorOffSymbol(Symbol):
    def __init__(self, color):
        super(ColorOffSymbol, self).__init__('coloroff')
        self.color = ord(color)

    def __str__(self):
        return 'COLOROFF=' + INV_COLORS[self.color]


class ColorInvSymbol(Symbol):
    def __init__(self):
        super(ColorInvSymbol, self).__init__('colorinv')

    def __str__(self):
        return 'COLORINV'



def lex(curline):
    '''
    split the line returned by `get_custom_viewer_curline` into symbols.
    it pulls out the strings, color directives, and escaped characters.
    
    Args:
      curline (str): a line returned by `ida_kernwin.get_custom_viewer_curline`
    
    Returns:
      generator: generator of Symbol subclass instances
    '''

    offset = 0
    cur_word = []
    while offset < len(curline):

        c = curline[offset]

        if c == ida_lines.COLOR_ON:
            if cur_word:
                yield StringSymbol(''.join(cur_word))
                cur_word = []

            offset += 1
            color = curline[offset]

            yield ColorOnSymbol(color)
            offset += 1

        elif c == ida_lines.COLOR_OFF:
            if cur_word:
                yield StringSymbol(''.join(cur_word))
                cur_word = []

            offset += 1
            color = curline[offset]

            yield ColorOffSymbol(color)
            offset += 1

        elif c == ida_lines.COLOR_ESC:
            if cur_word:
                yield StringSymbol(''.join(cur_word))
                cur_word = []

            offset += 1
            c = curline[offset]

            cur_word.append(c)
            offset += 1

        elif c == ida_lines.COLOR_INV:
            if cur_word:
                yield StringSymbol(''.join(cur_word))
                cur_word = []

            yield ColorInvSymbol()
            offset += 1

        else:
            cur_word.append(c)
            offset += 1


def get_color_at_char(curline, index):
    curlen = 0
    curcolor = 0
    for sym in lex(curline):
        if sym.type == 'string':
            curlen += len(sym.string)
            if curlen >= index:
                return curcolor
        elif sym.type == 'coloron':
            curcolor = sym.color
        elif sym.type == 'coloroff':
            curcolor = 0
        else:
            curcolor = 0

    return curcolor


def get_token_at_char(curline, index):
    curlen = 0
    curcolor = 0
    for sym in lex(curline):
        if sym.type == 'string':
            curlen += len(sym.string)
            if curlen >= index:
                return sym.string
        else:
            continue

    return ''


#### end: idapython_lex_curline
###############################################################################


def enum_function_addrs(fva):
    '''
    yield the effective addresses of each instruction in the given function.
    these addresses are not guaranteed to be in any order.
    
    Args:
      fva (int): the starting address of a function
    
    Returns:
      sequence[int]: the addresses of each instruction
    '''
    f = idaapi.get_func(fva)
    if not f:
        raise ValueError('not a function')

    for block in idaapi.FlowChart(f):
        ea = block.start_ea
        while ea <= block.end_ea:
            yield ea
            ea = ida_bytes.next_head(ea, idc.BADADDR)


def enum_calls_in_function(fva):
    '''
    yield the call instructions in the given function.
    
    Args:
      fva (int): the starting address of a function
    
    Returns:
      sequence[tuple[int, str]]: the address of a call instruction, and the disassembly line at that address
    '''
    for ea in enum_function_addrs(fva):
        if idaapi.is_call_insn(ea):
            disasm = ida_lines.generate_disassembly(ea, 16, True, False)[1][0]
            # replace consequent whitespaces by a single whitespaces
            disasm = re.sub("\s\s+", " ", disasm)
            yield ea, disasm


def enum_string_refs_in_function(fva):
    '''
    yield the string references in the given function.
    
    Args:
      fva (int): the starting address of a function
    
    Returns:
      sequence[tuple[int, int, str]]: tuples of metadata, including:
       - the address of the instruction referencing a string
       - the address of the string
       - the string
    '''
    for ea in enum_function_addrs(fva):
        for ref in idautils.DataRefsFrom(ea):
            stype = ida_nalt.get_str_type(ref)
            if stype < 0 or stype > 7:
                continue

            CALC_MAX_LEN = -1
            try:
                s = ida_bytes.get_strlit_contents(ref, CALC_MAX_LEN, stype).decode("utf-8")
            except UnicodeDecodeError:
                s = str(ida_bytes.get_strlit_contents(ref, CALC_MAX_LEN, stype))

            yield ea, ref, s


def uniq(s):
    '''
    yield the unique items in the given sequence.
    only provide the first copy encountered, skipping subsequent copies.
    
    Args:
      s (sequence[T]): a sequence of things

    Returns:
      sequence[T]: a sequence of the unique copies
    '''
    seen = set([])

    for item in s:
        if item in seen:
            continue

        yield item

        seen.add(item)


def render_function_hint(fva):
    '''
    create a textual report for the call and string references in the given function.
    
    eg.

        3 calls, 4 strings

        calls:
          - call     new
          - call     gethostbyname
          - call     delete
    
        strings:
          - www.google.com
          - www.bing.com
          - www.yahoo.com
          - www.cnn.com
    
    Args:
      fva (int): the starting address of a function

    Returns:
      str: the report
    '''
    ret = []

    # use `uniq` here (vs using a set) cause we want to maintain *some* semblance of order.
    calls = list(uniq(d for f, d in enum_calls_in_function(fva)))
    strings = list(uniq(s for o, r, s in enum_string_refs_in_function(fva)))
    xrefs = [xref.frm for xref in idautils.XrefsTo(fva, ida_xref.XREF_ALL) 
             if ida_bytes.is_code(ida_bytes.get_flags(xref.frm))]

    # this would be a good place to use Jinja2 templating,
    #  but lets not require that external dependency
    title = ida_lines.COLSTR('%d calls, ' % len(calls), ida_lines.SCOLOR_CODNAME)
    title += ida_lines.COLSTR('%s strings, ' % len(strings), ida_lines.SCOLOR_DSTR)
    title += ida_lines.COLSTR('%d xrefs ' % len(xrefs), ida_lines.SCOLOR_DEFAULT)
    ret.append(title)

    ret.append('')
    if calls:
        ret.append(ida_lines.COLSTR('calls:', ida_lines.SCOLOR_DEFAULT))
        for call in calls:
            ret.append('  - ' + call)
        ret.append('')
    
    if strings:
        ret.append(ida_lines.COLSTR('strings:', ida_lines.SCOLOR_DEFAULT))
        for s in strings:
            ret.append('  - "' + ida_lines.COLSTR(s, ida_lines.SCOLOR_DSTR) + '"')

    return '\n'.join(ret)


class CallsHintsHook(ida_kernwin.UI_Hooks):
    def get_custom_viewer_hint(self, view, place):
        try:
            widget = ida_kernwin.get_current_widget()
            if ida_kernwin.get_widget_type(widget) != ida_kernwin.BWN_DISASM:
                return None

            curline = ida_kernwin.get_custom_viewer_curline(view, True)
            
            # sometimes get_custom_viewer_place() returns [x, y] and sometimes [place_t, x, y].
            # we want the place_t.
            viewer_place = ida_kernwin.get_custom_viewer_place(view, True)
            if len(viewer_place) != 3:
                return None

            _, x, y = viewer_place
            ea = place.toea()

            # "color" is a bit of misnomer: its the type of the symbol currently hinted
            color = get_color_at_char(curline, x)
            if color != ida_lines.COLOR_ADDR:
                return None

            # grab the FAR references to code (not necessarilty a branch/call/jump by itself)
            far_code_references = [xref.to for xref in idautils.XrefsFrom(ea, ida_xref.XREF_FAR) 
                                   if ida_bytes.is_code(ida_bytes.get_flags(xref.to))]
            if len(far_code_references) != 1:
                return None

            fva = far_code_references[0]

            # ensure its actually a function
            if not idaapi.get_func(fva):
                return None

            # this magic constant is the number of "important lines" to display by default.
            # the remaining lines get shown if you scroll down while the hint is displayed, revealing more lines.
            return render_function_hint(fva), DEFAULT_IMPORTANT_LINES_NUM
        except Exception as e:
            logger.warning('unexpected exception: %s. Get in touch with @williballenthin.', e, exc_info=True)
            return None


class CallsHintsPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_KEEP
    comment = 'Display the calls and strings referenced by a function as hints.'
    help = 'Display the calls and strings referenced by a function as hints.'
    wanted_name = "CallsHintsPlugin"
    wanted_hotkey = "Ctrl-'"

    def init(self):
        self.hooks = CallsHintsHook()
        if self.hooks.hook():
            return idaapi.PLUGIN_KEEP
        else:
            logger.warning('error setting hooks.')
            return idaapi.PLUGIN_SKIP

    def run(self, arg):
        pass
        

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


def PLUGIN_ENTRY():
    return CallsHintsPlugin()


repository: securisec/mixto-scripts
file: mixto-ida/mixto_ida_py3.py
description: Repo holds scripts to integrate with other tools
wanted_name: Mixto
created_at: 2020-12-15 01:13:08+00:00
pushed_at: 2025-05-12 01:29:09+00:00
url:
https://github.com/securisec/mixto-scripts/blob/4457843c8fc51cccd1b584b455996086c0a8e707/mixto-ida/mixto_ida_py3.py

content:
# Mixto lite lib for python3

from typing import List
from urllib.request import Request, urlopen
from urllib.parse import urlencode, urljoin
from urllib.error import HTTPError
from pathlib import Path
from os import getenv
import json

MIXTO_ENTRY_ID = getenv("MIXTO_ENTRY_ID")
MIXTO_HOST = getenv("MIXTO_HOST")
MIXTO_API_KEY = getenv("MIXTO_API_KEY")


class MissingRequired(Exception):
    """Missing params"""

    pass


class BadResponse(Exception):
    """Bad response from Mixto API"""

    pass


class MixtoLite:
    def __init__(self, host: str = None, api_key: str = None) -> None:
        super().__init__()
        self.host = host
        self.api_key = api_key
        self.workspace = None
        self.status = 0
        self.commit_type = "tool"
        self.entry_id = None

        # if envars are set, always use those values
        if MIXTO_HOST is not None and self.host is None:
            self.host = MIXTO_HOST
        if MIXTO_API_KEY is not None and self.api_key is None:
            self.api_key = MIXTO_API_KEY

        # if host or apikey is not available, read config file
        if self.host == None or self.api_key == None:
            try:
                conf_path = str(Path().home() / ".mixto.json")
                with open(conf_path) as f:
                    j = json.loads(f.read())
                    self.host = j["host"]
                    self.api_key = j["api_key"]
                    self.workspace = j["workspace"]
            except:
                print("Cannot read mixto config file")
                raise

    def MakeRequest(
        self,
        method: str,
        uri: str,
        body: dict = {},
        query: dict = {},
        isJSON: bool = True,
    ):
        """Generic method helpful in extending this lib for other Mixto
        API calls. Refer to Mixto docs for all available API endpoints.

        Args:
            method (str): Request method
            uri (str): Mixto URI.
            body (dict, optional): Body. Defaults to {}.
            query (dict, optional): Query params. Defaults to {}.
            isJSON (bool, optional): If the response is of type JSON. Defaults to True.

        Raises:
            BadResponse: [description]
            BadResponse: [description]

        Returns:
            [type]: [description]
        """
        url = urljoin(str(self.host), uri)
        q = ""
        if query:
            q = "?" + urlencode(query)
        req = Request(
            method=method.upper(),
            url=url + q,
            data=json.dumps(body).encode(),
            headers={
                "x-api-key": self.api_key,
                "user-agent": "mixto-lite-py",
            },
        )
        if body:
            req.add_header("Content-Type", "application/json")
        try:
            res = urlopen(req)
            body = res.read().decode()
            self.status = res.getcode()
            if self.status > 300:
                raise BadResponse(self.status, res)
            if isJSON:
                return json.loads(str(body))
            else:
                return body
        except HTTPError as e:
            raise BadResponse(e.code, e.read())

    def AddCommit(self, data: str, entry_id: str = None, title: str = ""):
        """Add/commit data to an entry. This is the primary functionality of
        an integration

        Args:
            data (str): Data to add
            entry_id (str, optional): Entry ID. Will use MIXTO_ENTRY_ID as primary. Defaults to None.
            title (str, optional): Title for commit. Defaults to "Untitled".

        Raises:
            MissingRequired: If entry id is missing

        Returns:
            dict: Commit added response
        """
        if MIXTO_ENTRY_ID is None and entry_id is None:
            raise MissingRequired("Entry id is missing")

        e_id = MIXTO_ENTRY_ID if MIXTO_ENTRY_ID else entry_id
        r = self.MakeRequest(
            "POST",
            "/api/entry/{}/{}/commit".format(self.workspace, e_id),
            {"data": data, "type": self.commit_type, "title": title},
        )
        return r

    def GetWorkspaces(self) -> List[dict]:
        """Get all workspaces, entries and commits in a compact format.
        Helpful when trying to populate entry ID and commit ID's or
        filter by workspace

        Returns:
            List[dict]: Array of workspace items
        """
        return self.MakeRequest(
            "GET",
            "/api/misc/workspaces/{}".format(self.workspace),
            None,
            {"all": "true"},
            True,
        )

    def GetEntryIDs(self) -> List[str]:
        """Get all entry ids filtered by the current workspace

        Returns:
            List[str]: List of entry ids
        """
        # get all workspaces
        workspaces = self.GetWorkspaces()
        # filter workspaces by current workspace
        return [w["entry_id"] for w in workspaces if w["workspace"] == self.workspace]


"""
Mixto IDA Plugin
"""
import idc
import idaapi
import idautils
import ida_hexrays

# Retrieving imports from IDAPython is a little weird
# So need this global variable
AllImports = ""


def imports_cb(ea, name, ord):
    global AllImports
    if not name:
        AllImports += "0x{:x}: ord#{}\n".format(ea, ord)
    else:
        AllImports += "0x{:x}: {} ord#{}\n".format(ea, name, ord)
    # True -> Continue enumeration
    # False -> Stop enumeration
    return True


def start_ea_of(o):
    return getattr(o, "start_ea" if idaapi.IDA_SDK_VERSION >= 700 else "startEA")


def end_ea_of(o):
    return getattr(o, "end_ea" if idaapi.IDA_SDK_VERSION >= 700 else "endEA")


def get_flags_at(ea):
    return getattr(
        idaapi, "get_flags" if idaapi.IDA_SDK_VERSION >= 700 else "getFlags"
    )(ea)


def is_data(flags):
    return getattr(idaapi, "is_data" if idaapi.IDA_SDK_VERSION >= 700 else "isData")(
        flags
    )


MenuAllFunc = "Mixto:AllFunc"
MenuImports = "Mixto:Imports"
MenuExports = "Mixto:Exports"
MenuAllComments = "Mixto:AllComments"
MenuDecFunc = "Mixto:DecompileFunc"
MenuAskEntryId = "Mixto:AskEntryId"


class mixto_t(idaapi.plugin_t):

    flags = idaapi.PLUGIN_UNL
    comment = ""
    help = ""
    wanted_name = "Mixto"
    wanted_hotkey = ""
    mixto = MixtoLite()

    def init(self):
        self.add_menu_items()
        return idaapi.PLUGIN_OK

    @classmethod
    def add_menu_item_helper(self, name, text, tooltip, icon, shortcut):
        description = idaapi.action_desc_t(
            name, text, self.StartHandler(self, name), shortcut, tooltip, icon
        )
        idaapi.register_action(description)
        idaapi.attach_action_to_menu(
            self.wanted_name + "/" + text, name, idaapi.SETMENU_APP
        )

    @classmethod
    def add_menu_items(self):
        idaapi.create_menu(self.wanted_name, self.wanted_name)
        self.add_menu_item_helper(
            MenuAskEntryId,
            "Set Mixto Entry ID",
            "Specify entry id to send subsequent data",
            -1,
            "",
        )
        self.add_menu_item_helper(
            MenuAllFunc,
            "Send all functions",
            "Send all function names and addresses to mixto",
            -1,
            "",
        )
        self.add_menu_item_helper(
            MenuDecFunc,
            "Send decompilation of selected function",
            "Send decompilation of selected function to mixto",
            -1,
            "",
        )
        self.add_menu_item_helper(
            MenuImports, "Send all imports", "Send all imports to mixto", -1, ""
        )
        self.add_menu_item_helper(
            MenuExports, "Send all exports", "Send all exports to mixto", -1, ""
        )
        self.add_menu_item_helper(
            MenuAllComments,
            "Send all comments",
            "Send all comments and respective addresses to mixto",
            -1,
            "",
        )

    class StartHandler(idaapi.action_handler_t):
        def __init__(self, outer_self, menu_title):
            idaapi.action_handler_t.__init__(self)
            self.menu_title = menu_title
            self.outer_self = outer_self

        def activate(self, ctx):
            if self.menu_title == MenuAskEntryId:
                self.outer_self.mixto.entry_id = idaapi.ask_str(
                    "", 1000, "Mixto Entry Id"
                )

            else:
                if self.outer_self.mixto.entry_id is None:
                    self.outer_self.mixto.entry_id = idaapi.ask_str(
                        "", 1000, "Mixto Entry Id"
                    )

            if self.menu_title == MenuAllFunc:
                all_func = ""
                # Get count of all functions in the binary
                count = idaapi.get_func_qty()

                for i in range(count):
                    fn = idaapi.getn_func(i)

                    # Function should not have dummy name such as sub_*
                    # and should not start with underscore (possible library functions)
                    if not idaapi.has_dummy_name(
                        get_flags_at(start_ea_of(fn))
                    ) and not idaapi.get_func_name(start_ea_of(fn)).startswith("_"):
                        all_func += "{} @ 0x{:x}\n".format(
                            idaapi.get_func_name(start_ea_of(fn)), start_ea_of(fn)
                        )
                self.outer_self.mixto.AddCommit(
                    all_func, self.outer_self.mixto.entry_id, "(IDA) All Functions"
                )

            elif self.menu_title == MenuImports:
                global AllImports
                AllImports = ""
                # Get count of all import modules in the binary
                count = idaapi.get_import_module_qty()

                for i in range(count):
                    module_name = idaapi.get_import_module_name(i)

                    AllImports += "{}:\n".format(module_name)
                    idaapi.enum_import_names(i, imports_cb)
                self.outer_self.mixto.AddCommit(
                    AllImports, self.outer_self.mixto.entry_id, "(IDA) All Imports"
                )

            elif self.menu_title == MenuDecFunc:
                addr_current = idc.get_screen_ea()
                addr_func = idaapi.get_func(addr_current)

                if not addr_func:
                    idaapi.msg("Place cursor inside a function!")
                    return False
                else:
                    err = None
                    out = ida_hexrays.decompile_func(addr_func, err)
                    # print(out)
                    self.outer_self.mixto.AddCommit(
                        str(out),
                        self.outer_self.mixto.entry_id,
                        "(IDA) Function Decompilation",
                    )

            elif self.menu_title == MenuExports:
                all_exports = ""
                for entry in idautils.Entries():
                    _, ord, ea, name = entry
                    if not name:
                        all_exports += "0x{:x}: ord#{}\n".format(ea, ord)
                    else:
                        all_exports += "0x{:x}: {} ord#{}\n".format(ea, name, ord)
                self.outer_self.mixto.AddCommit(
                    all_exports, self.outer_self.mixto.entry_id, "(IDA) All Exports"
                )

            elif self.menu_title == MenuAllComments:
                raise NotImplementedError("Comments not yet implemented TODO")

            return True

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

    def run(self, arg):
        pass

    def term(self):
        pass


def PLUGIN_ENTRY():
    return mixto_t()


repository: vmallet/ida-describekey
file: describekey.py
description: IDA plugin to quickly learn what a shortcut does
created_at: 2022-04-08 23:00:45+00:00
pushed_at: 2022-04-08 23:04:18+00:00
url:
https://github.com/vmallet/ida-describekey/blob/0b781202d3c835001c063659a088d79f68714e82/describekey.py

content:
# SPDX-FileCopyrightText: 2022 Vincent Mallet <vmallet@gmail.com>
# SPDX-License-Identifier: MIT

"""
DescribeKey: a simple action to display actions associated with a
shortcut, live.
"""

from collections import defaultdict
from typing import Optional

from PyQt5 import QtWidgets, QtCore
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QKeyEvent, QKeySequence, QResizeEvent
from PyQt5.QtWidgets import QDialog, QTableWidget, QTableWidgetSelectionRange, \
    QVBoxLayout, QHeaderView, QLabel, QHBoxLayout

import ida_idaapi
import ida_kernwin
import idaapi

__author__ = "https://github.com/vmallet"

# TODO: use some sort of built-in window/dialog and have iDA remember last position
# TODO: seems to be having issues with ". ; : / - = <arrow-keys>" keys..˚

VERSION = 0.1

DESCRIBE_KEY_ACTION = "DescribeKey"
DESCRIBE_KEY_TEXT = "Describe Key"
DESCRIBE_KEY_SHORTCUT = "Alt-Shift-K"
DESCRIBE_KEY_TOOLTIP = "List all actions associated with a shortcut"

DEBUG = True

# Column indices for action table
COL_LABEL = 0
COL_ACTION = 1
COL_SHORTCUT = 2
COL_TOOLTIP = 3
COL_STATE = 4
COL_VISIBILITY = 5
COL_CHECKABLE = 6
COL_CHECKED = 7
COLUMN_COUNT = 8

DIRECT_KEYS = r",./<>?;':\"[]{}()`~!@#$%^&*-_=+"
"""Keys that shouldn't be mapped to a Qt::Key."""


class KeyNamer(object):
    """
    Provide a means to get an IDA-compatible shortcut representation of
    a key event.

    See keyevent_to_shortcut()
    """

    def __init__(self):
        # Key names
        self._keymap = dict()
        for key, value in vars(Qt).items():
            if isinstance(value, Qt.Key):
                self._keymap[value] = key.partition('_')[2]

        # Modifier names. Note: insertion order matters, should match IDA's modifier order
        self._modmap = {
            Qt.ControlModifier: "Ctrl",
            Qt.AltModifier: "Alt",
            Qt.ShiftModifier: "Shift",
            Qt.MetaModifier: "Meta",
            Qt.GroupSwitchModifier: self._keymap[Qt.Key_AltGr],
            Qt.KeypadModifier: self._keymap[Qt.Key_NumLock],
        }

    def keyevent_to_shortcut(self, event) -> Optional[str]:
        """Attempt to produce IDA-compatible shortcut for keyevent."""
        text = event.text()
        if text and text in DIRECT_KEYS:
            key = text
        else:
            # Try to map the key, first using the native virtual key and only
            # if it's a legit key like "A" or "F1", not "guillemotleft" or "cent"
            native = self._keymap.get(event.nativeVirtualKey(), None)
            if native and len(native) > 2:
                native = None
            # If we don't have a simple key, try to map the actual event key and
            # if all else fails, use the event text
            key = native or self._keymap.get(event.key(), text)

        if key in [None, "Control", "Alt", "Shift", "Meta"]:
            return None

        if event.modifiers() == Qt.ShiftModifier and key in DIRECT_KEYS:
            # A bit hacky here.. IDA looks at '%' as a non-shifted shortcut, but
            # on US keyboards you would need shift to produce '%'. So if the only
            # modifier used with a 'direct key' is Shift, ignore it. This might or
            # might not work for other locales.
            return key

        sequence = []
        for modifier, text in self._modmap.items():
            if event.modifiers() & modifier:
                sequence.append(text)
        if key not in sequence:
            sequence.append(key)
        return '-'.join(sequence)


class ActionInfo(object):
    """Description of a registered action."""

    def __init__(self, label, action, shortcut, tooltip, state, visibility, checkable, checked):
        self.label = label
        self.action = action
        self.shortcut = shortcut
        self.tooltip = tooltip
        self.state = state
        self.visibility = visibility
        self.checkable = checkable
        self.checked = checked

    @classmethod
    def for_action(cls, action):
        """Construct an ActionInfo for the given action name."""

        def unw(ret):
            return ret[1] if ret and ret[0] else None

        shortcut = ida_kernwin.get_action_shortcut(action)
        label = ida_kernwin.get_action_label(action)
        tooltip = ida_kernwin.get_action_tooltip(action)
        state = unw(ida_kernwin.get_action_state(action))
        visibility = unw(ida_kernwin.get_action_visibility(action))
        checkable = unw(ida_kernwin.get_action_checkable(action))
        checked = unw(ida_kernwin.get_action_checked(action))

        return ActionInfo(label, action, shortcut, tooltip, state, visibility, checkable,
                          checked)


class DescribeKey(object):
    """
    A custom dialog that will show all actions registered with a
    shortcut when one is pressed.

    Actions matching a shortcut are displayed in a QTableWidget.

    Key-presses are intercepted by the table widget with a custom
    keyPressEvent() handler, and turned into shortcuts that can
    be used to look up the corresponding actions.

    ESC to exit.
    """

    def __init__(self):
        self._namer = KeyNamer()
        self._astmap = self._build_ast_map()
        self._shortcut_label = QLabel("Press a shortcut...")
        self._table = self._build_table()
        self._dialog = self._build_dialog()
        self._overlay = self._build_overlay()

    def _build_ast_map(self):
        """Build a Value->Name map of all AST_xxx enum values."""
        astmap = {
            ida_kernwin.AST_ENABLE: "Enable",
            ida_kernwin.AST_ENABLE_ALWAYS: "Enable Always",
            ida_kernwin.AST_ENABLE_FOR_WIDGET: "Enable for Widget",
            ida_kernwin.AST_ENABLE_FOR_IDB: "Enable for IDB",
            ida_kernwin.AST_DISABLE: "Disable",
            ida_kernwin.AST_DISABLE_ALWAYS: "Disable Always",
            ida_kernwin.AST_DISABLE_FOR_WIDGET: "Disable for Widget",
            ida_kernwin.AST_DISABLE_FOR_IDB: "Disable for IDB",
        }
        return astmap

    def _build_action_map(self):
        """Return a Shortcut->Action map of all registered actions."""
        action_map = defaultdict(list)

        actions = ida_kernwin.get_registered_actions()
        for name in actions:
            shortcut = ida_kernwin.get_action_shortcut(name)
            if shortcut:
                action_map[shortcut].append(name)

        return action_map

    def _handle_keyevent(self, event: QKeyEvent, action_map, fn):
        """Intercept key events and update UI with related actions."""
        # First, clear the overlay
        self._dismiss_overlay()

        shortcut = self._namer.keyevent_to_shortcut(event)
        self._set_shortcut(shortcut)
        if DEBUG:
            print("evt: {}  key: {:7X}  native key: {:2X}  native scancode: {:2X}  "
                  "text: {:1}  shortcut: {}".format(
                type(event), event.key(), event.nativeVirtualKey(), event.nativeScanCode(),
                event.text(), shortcut))

        actions = action_map.get(shortcut, []) if shortcut else []
        self._set_data(actions)

        # Only the ESC key goes through
        if event.matches(QKeySequence.Cancel):
            fn(event)

    def _build_table(self) -> QTableWidget:
        """Construct the table widget used to display actions."""
        table = QtWidgets.QTableWidget()
        table.setColumnCount(COLUMN_COUNT)
        table.setRowCount(0)

        table.setEditTriggers(QtWidgets.QAbstractItemView.NoEditTriggers)
        table.setSelectionBehavior(QtWidgets.QAbstractItemView.SelectRows)
        table.setRangeSelected(QTableWidgetSelectionRange(0, 0, 0, 0), True)
        table.setShowGrid(False)

        table.setHorizontalHeaderItem(COL_LABEL, QtWidgets.QTableWidgetItem("Label"))
        table.setHorizontalHeaderItem(COL_ACTION, QtWidgets.QTableWidgetItem("Action"))
        table.setHorizontalHeaderItem(COL_SHORTCUT, QtWidgets.QTableWidgetItem("Shortcut"))
        table.setHorizontalHeaderItem(COL_TOOLTIP, QtWidgets.QTableWidgetItem("Description"))
        table.setHorizontalHeaderItem(COL_STATE, QtWidgets.QTableWidgetItem("State"))
        table.setHorizontalHeaderItem(COL_VISIBILITY, QtWidgets.QTableWidgetItem("V"))
        table.setHorizontalHeaderItem(COL_CHECKABLE, QtWidgets.QTableWidgetItem("C"))
        table.setHorizontalHeaderItem(COL_CHECKED, QtWidgets.QTableWidgetItem("Cd"))

        # Magic calculation in attempt to size the State column more or less sensibly
        state_width = int(table.fontMetrics().width("Disable for Widget") * 1.2) + 5

        table.horizontalHeader().setStretchLastSection(False)
        table.horizontalHeader().setSectionResizeMode(COL_TOOLTIP, QHeaderView.Stretch)

        table.setColumnWidth(COL_SHORTCUT, 88)
        table.setColumnWidth(COL_STATE, state_width)
        table.setColumnWidth(COL_VISIBILITY, 25)
        table.setColumnWidth(COL_CHECKABLE, 25)
        table.setColumnWidth(COL_CHECKED, 25)

        for i in range(COLUMN_COUNT):
            table.horizontalHeaderItem(i).setTextAlignment(QtCore.Qt.AlignLeft)

        table.verticalHeader().setHidden(True)
        table.verticalHeader().setSectionResizeMode(QtWidgets.QHeaderView.ResizeToContents)
        table.verticalHeader().setMaximumSectionSize(19)   # TODO: '19' magic constant

        # We build the action map once for the lifetime of the dialog
        action_map = self._build_action_map()
        old_kp = table.keyPressEvent
        table.keyPressEvent = lambda evt: self._handle_keyevent(evt, action_map, old_kp)

        old_resize = table.resizeEvent
        def resizeEvent(evt: QResizeEvent):
            """Size LABEL and ACTION columns proportionally to table width."""
            width = evt.size().width()
            if width != evt.oldSize().width():
                self._table.setColumnWidth(COL_LABEL, width / 5)
                self._table.setColumnWidth(COL_ACTION, width / 5)
            old_resize(evt)

        table.resizeEvent = resizeEvent

        return table

    def _set_shortcut(self, shortcut):
        """Set the current shortcut being displayed by the UI."""
        self._shortcut_label.setText(shortcut)

    def _set_data(self, actions):
        """Set the actions being displayed by the UI (in the table)."""
        self._table.clearContents()
        self._table.setRowCount(len(actions))

        for i, action in enumerate(actions):
            info = ActionInfo.for_action(action)

            # Label
            item = QtWidgets.QTableWidgetItem()
            item.setData(QtCore.Qt.DisplayRole, info.label)
            self._table.setItem(i, COL_LABEL, item)

            # Action
            item = QtWidgets.QTableWidgetItem()
            item.setData(QtCore.Qt.DisplayRole, info.action)
            self._table.setItem(i, COL_ACTION, item)

            # Shortcut
            item = QtWidgets.QTableWidgetItem()
            item.setData(QtCore.Qt.DisplayRole, info.shortcut)
            self._table.setItem(i, COL_SHORTCUT, item)

            # Tooltip
            item = QtWidgets.QTableWidgetItem()
            item.setData(QtCore.Qt.DisplayRole, info.tooltip)
            self._table.setItem(i, COL_TOOLTIP, item)

            # State
            item = QtWidgets.QTableWidgetItem()
            item.setData(QtCore.Qt.DisplayRole,
                         self._astmap.get(info.state, None) or str(info.state))
            self._table.setItem(i, COL_STATE, item)

            # Visibility
            item = QtWidgets.QTableWidgetItem()
            item.setData(QtCore.Qt.DisplayRole, "V" if info.visibility else "")
            self._table.setItem(i, COL_VISIBILITY, item)

            # Checkable
            item = QtWidgets.QTableWidgetItem()
            item.setData(QtCore.Qt.DisplayRole, "Y" if info.checkable else "")
            self._table.setItem(i, COL_CHECKABLE, item)

            # Checked
            item = QtWidgets.QTableWidgetItem()
            item.setData(QtCore.Qt.DisplayRole, "Y" if info.checked else "")
            self._table.setItem(i, COL_CHECKED, item)

    def _build_status(self):
        """Construct the status line for the UI."""
        layout = QHBoxLayout()
        layout.addWidget(QLabel("Shortcut: "))
        layout.addWidget(self._shortcut_label)
        layout.addStretch()
        layout.setContentsMargins(5, 5, 5, 5)
        return layout

    def _build_dialog(self):
        """Construct the main UI dialog."""
        dialog = QDialog()
        dialog.setWindowTitle("Describe Key")
        dialog.resize(800, 200)

        layout = QVBoxLayout()
        layout.setContentsMargins(0, 0, 0, 0)
        layout.addWidget(self._table)
        layout.addLayout(self._build_status())
        dialog.setLayout(layout)

        return dialog

    def _build_overlay(self):
        """
        Construct the initial help overlay.

        The overlay is a 'floating' widget, child of the dialog. It
        will be raised on top of the other widgets (the table), and
        hidden upon the first keypress received.
        """
        label = QLabel("Press a shortcut...", self._dialog)
        label.setStyleSheet("color : #CC3030; font-size: 17px; font-style: italic")
        label.adjustSize()

        px = (self._dialog.width() - label.width()) / 2
        label.move(px, 50)
        label.show()
        label.raise_()

        return label

    def _dismiss_overlay(self):
        """Hide the help overlay, if it exists."""
        if self._overlay:
            self._overlay.hide()
            self._overlay = None

    def show(self):
        """Show the main UI dialog."""
        self._dialog.exec()


class KeyActionHandler(ida_kernwin.action_handler_t):

    def activate(self, ctx):
        # Build a new dialog for every invocation of the action
        dk = DescribeKey()
        dk.show()
        return False

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


def unregister_action():
    """Unregister the DescribeKey action from IDA."""
    ida_kernwin.unregister_action(DESCRIBE_KEY_ACTION)

def register_action():
    """Register the DescribeKey action with IDA."""
    unregister_action()
    desc = ida_kernwin.action_desc_t(
        DESCRIBE_KEY_ACTION,
        DESCRIBE_KEY_TEXT,
        KeyActionHandler(),
        DESCRIBE_KEY_SHORTCUT,
        DESCRIBE_KEY_TOOLTIP)
    return ida_kernwin.register_action(desc)


class DescribeKeyPlugin(ida_idaapi.plugin_t):
    flags = idaapi.PLUGIN_FIX  # Always stay loaded, even without an IDB
    wanted_name = DESCRIBE_KEY_TEXT
    wanted_hotkey = DESCRIBE_KEY_SHORTCUT
    comment = DESCRIBE_KEY_TOOLTIP
    help = ""
    version = VERSION

    def init(self):
        return ida_idaapi.PLUGIN_KEEP  # keep us in the memory

    def term(self):
        pass

    def run(self, arg):
        dk = DescribeKey()
        dk.show()


def PLUGIN_ENTRY():
    """Plugin entry point when loaded as a plugin by IDA."""
    return DescribeKeyPlugin()


repository: 20000s/app_practise
file: 企业微信/re_qywx/traceNatives.py
description: 我逆过的一些app
wanted_name: traceNatives
comment: traceNatives
created_at: 2022-02-15 11:05:15+00:00
pushed_at: 2022-11-02 13:29:11+00:00
url:
https://github.com/20000s/app_practise/blob/26e61f049b312627f332701a8847c18cd993f3e8/%E4%BC%81%E4%B8%9A%E5%BE%AE%E4%BF%A1/re_qywx/traceNatives.py

content:
# -*- coding:utf-8 -*-
import os
from idaapi import plugin_t
from idaapi import PLUGIN_PROC
from idaapi import PLUGIN_OK
import ida_nalt
import idaapi
import idautils
import idc
import time


# 获取SO文件名和路径
def getSoPathAndName():
    fullpath = ida_nalt.get_input_file_path()
    filepath,filename = os.path.split(fullpath)
    return filepath,filename

# 获取代码段的范围
def getSegAddr():
    textStart = []
    textEnd = []

    for seg in idautils.Segments():
        if (idc.get_segm_name(seg)).lower() == '.text' or (
        idc.get_segm_name(seg)).lower() == 'text':
            tempStart = idc.get_segm_start(seg)
            tempEnd = idc.get_segm_end(seg)

            textStart.append(tempStart)
            textEnd.append(tempEnd)

    return min(textStart), max(textEnd)


class traceNatives(plugin_t):
    flags = PLUGIN_PROC
    comment = "traceNatives"
    help = ""
    wanted_name = "traceNatives"
    wanted_hotkey = ""

    def init(self):
        print("traceNatives(v0.1) plugin has been loaded.")
        return PLUGIN_OK

    def run(self, arg):
        # 查找需要的函数
        ea, ed = getSegAddr()
        search_result = []
        for func in idautils.Functions(ea, ed):
            try:
                functionName = str(idaapi.ida_funcs.get_func_name(func))
                if len(list(idautils.FuncItems(func))) > 10:
                    # 如果是thumb模式，地址+1
                    arm_or_thumb = idc.get_sreg(func, "T")
                    if arm_or_thumb:
                        func += 1
                    search_result.append(hex(func))
            except:
                pass

        so_path, so_name = getSoPathAndName()
        search_result = [f"-a '{so_name}!{offset}'" for offset in search_result]
        search_result = " ".join(search_result)

        script_name = so_name.split(".")[0] + "_" + str(int(time.time())) +".txt"
        save_path = os.path.join(so_path, script_name)
        with open(save_path, "w", encoding="utf-8")as F:
            F.write(search_result)

        print("使用方法如下：")
        print(f"frida-trace -UF -O {save_path}")

    def term(self):
        pass


def PLUGIN_ENTRY():
    return traceNatives()


repository: benwoo1110/PMA-Notes
file: Plugins/ifl.py
comment: Interactive Functions List
created_at: 2022-09-27 01:21:48+00:00
pushed_at: 2022-12-10 06:36:38+00:00
url:
https://github.com/benwoo1110/PMA-Notes/blob/9d322ef0cbb99709f9227e322846b58dd1b1c9ab/Plugins/ifl.py

content:
"""
Adapted from https://github.com/hasherezade/ida_ifl
"""

#!/usr/bin/env python3
#
# IFL - Interactive Functions List
#
# how to install: copy the script into plugins directory, i.e: C:\Program Files\IDA <version>\plugins
# then:
# run from IDA menu: View -> PLUGIN_NAME
# or press: PLUGIN_HOTKEY
#
"""
CC-BY: hasherezade, run via IDA Pro >= 7.0
"""
__VERSION__ = '1.4.4'
__AUTHOR__ = 'hasherezade'

PLUGIN_NAME = "IFL - Interactive Functions List"
PLUGIN_HOTKEY = "Ctrl-Alt-F"

import idaapi  # type: ignore
import idc  # type: ignore
import ida_bytes
import ida_kernwin

from idaapi import BADADDR, jumpto, next_addr, o_void, prev_addr,\
    print_insn_mnem, set_cmt, set_name, SN_NOWARN
from idc import CIC_ITEM, demangle_name, FUNCATTR_END, FUNCATTR_FRAME,\
    get_func_attr, get_func_name, get_inf_attr, get_operand_type,\
    get_operand_value, get_type, GetDisasm, INF_SHORT_DN, set_color, \
    FUNCATTR_START
from idautils import Functions, XrefsFrom, XrefsTo  # type: ignore

from idaapi import PluginForm  # type: ignore
from PyQt5 import QtGui, QtCore, QtWidgets  # type: ignore
from PyQt5.QtCore import QObject, pyqtSignal  # type: ignore

from typing import Optional, List, Tuple, Any, Union

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

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

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

IS_ALTERNATE_ROW = False

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

# Theme

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

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

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

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

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

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


def _is_hex_str(s):
    # skip prefix:
    if s[:2] == "0x":
        s = s[2:]
    # validate hex charset:
    hex_digits = set("0123456789abcdefABCDEF")
    for val in s:
        if not (val in hex_digits):
            return False
    return True


# Functions and args 

def function_at(ea: int) -> Optional[int]:
    start = ea
    functions = Functions(start)
    for func in functions:
        return func
    return None


def parse_function_args(ea: int) -> str:
    local_variables = []
    arguments = []
    current = local_variables

    frame = idc.get_func_attr(ea, FUNCATTR_FRAME)
    arg_string = ""
    if frame is None:
        return ""

    start = idc.get_first_member(frame)
    end = idc.get_last_member(frame)
    count = 0
    max_count = 10000
    args_str = ""
    while start <= end and count <= max_count:
        size = idc.get_member_size(frame, start)
        count = count + 1
        if size is None:
            start = start + 1
            continue

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

        if name in [" r", " s"]:
            # Skip return address and base pointer
            current = arguments
            continue
        arg_string += f" {name}"
        current.append(name)
    args_str = ", ".join(arguments)
    if len(args_str) == 0:
        args_str = "void"
    return f"({args_str})"


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

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

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

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


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


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


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


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


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

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


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

    updateSignal = pyqtSignal()

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

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

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

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

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


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

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

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


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

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

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

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

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

    def _displayBackground(self, row: int, col: int) -> Any:
        curr_theme = get_theme()
        if curr_theme is not None:
            self.theme = curr_theme
      
        func_info = self.function_info_list[row]
        if col == self.COL_START or col == self.COL_END:
            return QtGui.QColor(self.theme[0])
        if col == self.COL_NAME:
            if func_info.is_import:
                return QtGui.QColor(self.theme[1])
            return QtGui.QColor(self.theme[2])
        return None

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

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

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

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

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

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

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

        func_info = self.function_info_list[row]

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

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

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

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


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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


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

    # private

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

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

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

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

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

        self._importsSet = set()
        self._importNamesSet = set()
        nimps = idaapi.get_import_module_qty()
        for i in range(0, nimps):
            idaapi.enum_import_names(i, self.imports_names_callback)

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

        if name in self._importNamesSet:
            return True
        return False

    def _listRefsTo(self, start: int) -> List[Tuple[int, int]]:
        """Make a list of all the references to the given function.
        Args:
          func : The function references to which we are searching.
          start : The function's start offset.
        Returns:
          list : A list of tuples.
            Each tuple represents: the offsets:
            0 : the offset from where the given function was referenced by the foreign function
            1 : the function's start address
        """

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

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

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

    def _listRefsFrom(self, func, start: int, end: int) -> List[Tuple[int, int]]:
        """Make a list of all the references made from the given function.
        Args:
          func : The function inside of which we are searching.
          start : The function's start offset.
          end : The function's end offset.
        Returns:
          list : A list of tuples. Each tuple represents:
            0 : the offset from where the given function referenced the other entity
            1 : the address that was referenced
        """

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

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

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

            is_import = self._isImportStart(start)

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

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

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

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


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

# private
    _LIVE_FILTER = True

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

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

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

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

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

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

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

        info = idaapi.get_inf_structure()
        if info.is_64bit():
            curr_val = idc.get_qword(start)
            if (curr_val == thunk_val):
                return ida_bytes.create_data(start, idaapi.FF_QWORD, 8, idaapi.BADADDR)
        elif info.is_32bit():
            curr_val = ida_bytes.get_dword(start)
            if (curr_val == thunk_val):
                return ida_bytes.create_data(start, idaapi.FF_DWORD, 4, idaapi.BADADDR)
        return False

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

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

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

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

                set_cmt(start, func_name, 1)  # set the name as a comment
                comments += 1
        return (functions, comments)

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

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

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

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

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

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

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

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

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

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

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

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

        data = g_DataManager.currentRva
        self.setRefOffset(data)

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

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

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

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

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

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

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

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

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

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

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

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

        # init data structures:
        self.funcMapper = FunctionsMapper_t()
        self.criterium_id = 0

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

        # Create models
        self.subDataManager = DataManager()

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

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

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

        font = self.refs_view.font()
        font.setPointSize(8)
        self.refs_view.setFont(font)

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

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

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

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

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

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

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

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

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

        filter_panel.setLayout(filter_layout)
        self.filter_edit.setFixedHeight(20)
        filter_panel.setFixedHeight(40)
        filter_panel.setAutoFillBackground(True)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


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

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

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

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


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

    m_functionInfoForm.Show()

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

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

    help = "Interactive Function List. Comments? Remarks? Mail to: hasherezade@gmail.com"
    wanted_name = PLUGIN_NAME
    wanted_hotkey = ''

    def init(self) -> None:
        idaapi.register_action(idaapi.action_desc_t(
            'ifl:open',  # action name
            PLUGIN_NAME,
            IFLMenuHandler(),
            PLUGIN_HOTKEY,
            'Opens Interactive Function List Pane')
        )

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

        return idaapi.PLUGIN_OK

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

    def term(self) -> None:
        pass

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


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


repository: clovme/WTools
wanted_name: autoinst
created_at: 2025-03-27 00:14:54+00:00
pushed_at: 2025-05-11 10:18:21+00:00
file:
software/Reverse Cracking/IDA Professional 9.0/python/examples/core/auto_instantiate_widget_plugin.py

url:
https://github.com/clovme/WTools/blob/f12f01c833a6a6241799a9d01e9117cc28a65dbc/software/Reverse%20Cracking/IDA%20Professional%209.0/python/examples/core/auto_instantiate_widget_plugin.py

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

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

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

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

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

keywords: desktop
"""

import ida_idaapi
import ida_kernwin

title = "Auto-instantiable at IDA startup"

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

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

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

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

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

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

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

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

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

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

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

    def run(self, arg):
        pass

    def term(self):
        pass


def PLUGIN_ENTRY():
    return auto_inst_plugin_t()

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


repository: iancohee/ghidra-sre
file: Extensions/IDAPro/Python/6xx/loaders/xmlldr.py
description: Patched NSA Reverse Engineering Tool
wanted_name: XML Importer
comment: Import XML PROGRAM file
created_at: 2019-03-06 14:47:12+00:00
pushed_at: 2019-03-06 15:12:02+00:00
url:
https://github.com/iancohee/ghidra-sre/blob/273bbb415b4eea13bcf1176b247bfbe11c038dca/Extensions/IDAPro/Python/6xx/loaders/xmlldr.py

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

import idaapi
import time
from xml.etree import cElementTree


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


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

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


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

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


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

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

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


    def term(self):
        pass


def PLUGIN_ENTRY():
    return XmlImporterPlugin()


class Cancelled(Exception):
    pass


class FileError(Exception):
    pass


class MultipleAddressSpacesNotSupported(Exception):
    pass


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

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


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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


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

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

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

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

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

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


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

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

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

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


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


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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


repository: tomer132246/IDATopaqueminator
file: Topaqueminator.py
description: IDA plugin for simple opaque predicates removal using symbolic execution with angr
wanted_name: Topaqueminator
comment: Topaqueminator Comment
created_at: 2023-08-10 06:27:46+00:00
pushed_at: 2023-09-19 16:52:15+00:00
url:
https://github.com/tomer132246/IDATopaqueminator/blob/7258d4cf09c5f639da11fe7e624124c2260a0bdb/Topaqueminator.py

content:
import idaapi, idautils, idc, ida_funcs, ida_auto, ida_bytes    
import sys
import os
sys.stdout.encoding = 'utf-8'
import angr
from keystone import *
from capstone import *
from idc import *
from ida_ua import *
from ida_bytes import *
from ida_allins import *
import logging
from enum import Enum
import shutil
import re
import six
import sark
import struct
####################################################################
#Constants.
####################################################################

"""
TODO:
-fix the plugin "reload" need
-move keypatcher to different class
-check if patching the binary ida is working on is enough, or does angr need to reload it every time.
-implement deflow algorithm https://ferib.dev/blog.php?l=post/Reversing_Common_Obfuscation_Techniques' (????)
heuristics:
    -think about checkign the damage an opaque predicate does to a binary as an heursitic.
    -think about state explosion heuristic.
    -think about doing dual jp-mov-jnp fixes ahead.
-fix non full scan
-add patterns for found opaques
-

-PROBLEMS:
1111111111111111111111111 ON RET_ADDR_CHECK
.text:000000014209C761                 jo      short loc_14209C76A
.text:000000014209C763                 mov     bl, bl
.text:000000014209C765                 jmp     short loc_14209C78E ; Keypatch modified this from:
.text:000000014209C765 ; END OF FUNCTION CHUNK FOR retAddrCheck ;   jno short loc_14209C78E
.text:000000014209C765 ; ---------------------------------------------------------------------------
.text:000000014209C767                 db 81h
.text:000000014209C768                 db 0C2h, 18h
.text:000000014209C76A ; ---------------------------------------------------------------------------
.text:000000014209C76A ; START OF FUNCTION CHUNK FOR retAddrCheck
.text:000000014209C76A
.text:000000014209C76A loc_14209C76A:                          ; CODE XREF: retAddrCheck+281↑j
.text:000000014209C76A                 nop
.text:000000014209C76B                 jo      short loc_14209C78E


22222222222222222222222222 ON EXEC MAYBE
0x1406272f3
0x1406272f6
1 jb 5
2 jnb 8
3 db db
4 db db
5 jb 8 
6 db db
7 db db

"""


T_VER = 21
DEBUG = True
DEBUG_VERBOSE = False
CONDITIONAL_JMP_SIZE_IN_BYTES = 2
DEPTH = 150
ACTIVE_THERSHOLD_FOR_LAST_PREDICATE_MODE = 700
TOPAQUE_ANGR_ADDITION = '_Topaqueminator_ver_' + str(T_VER) +'_patched'
APPLY_PATCHES_ON_INPUT_FILE = True
CONTINUOS_ANALYZING = False
TESTS_MODE = False


SINGLE_BLOCK_MODE = False
WHOLE_FUNCTION_MODE = True
FIX_DEFLOW_MODE = True


patch_info = []
last_found_opaque_predicate = None
addresses_avoided_set = set()
addresses_find_set = set()
opaque_jumpedto_addresses_set = set()
unfixable_deflow_addresses = set()
is_last_predicate_mode = False


##################py compitability funcs################################################################
def to_string(s):
  # python3 bytes
  if six.PY3 and isinstance(s, bytes):
      return s.decode('latin-1')
  # python2 unicode
  elif six.PY2 and isinstance(s, six.text_type):
      return s.encode('utf-8')
  return str(s)

def to_hexstr(buf, sep=' '):
    # for python3 bytes
    if six.PY3 and isinstance(buf, bytes):
        return sep.join("{0:02x}".format(c) for c in buf).upper()
    return sep.join("{0:02x}".format(ord(c)) for c in buf).upper()

def convert_hexstr(code):
    # normalize code
    code = code.lower()
    code = code.replace(' ', '')    # remove space
    code = code.replace('h', '')    # remove trailing 'h' in 90h
    code = code.replace('0x', '')   # remove 0x
    code = code.replace('\\x', '')  # remove \x
    code = code.replace(',', '')    # remove ,
    code = code.replace(';', '')    # remove ;
    code = code.replace('"', '')    # remove "
    code = code.replace("'", '')    # remove '
    code = code.replace("+", '')    # remove +

    # single-digit hexcode?
    if len(code) == 1 and ((code >= '0' and code <= '9') or (code >= 'a' and code <= 'f')):
        # stick 0 in front (so 'a' --> '0a')
        code = '0' + code

    # odd-length is invalid
    if len(code) % 2 != 0:
        return None

    try:
        hex_data = code.decode('hex')
        # we want a list of int
        return [ord(i) for i in hex_data]
    except:
        # invalid hex
        return None

#######################################################################################################

################################ IDA 6/7 Compatibility function #########################################
def get_dtype(ea, op_idx):
    if idaapi.IDA_SDK_VERSION >= 700:
        insn = idaapi.insn_t()
        idaapi.decode_insn(insn, ea)
        dtype = insn.ops[op_idx].dtype
        dtyp_size = idaapi.get_dtype_size(dtype)
    else:
        dtype = idaapi.cmd.Operands[op_idx].dtyp
        dtyp_size = idaapi.get_dtyp_size(dtype)
    return dtype, dtyp_size

def set_comment(ea, comment):
    if idaapi.IDA_SDK_VERSION >= 700:
        idc.set_cmt(ea, comment, 0)
    else:
        idc.MakeComm(ea, comment)

def get_comment(ea):
    if idaapi.IDA_SDK_VERSION >= 700:
        return idc.get_cmt(ea, 0)
    return idc.Comment(ea)

def read_range_selection():
    if idaapi.IDA_SDK_VERSION >= 700:
        return idaapi.read_range_selection(None)
    return idaapi.read_selection()
#########################################################################################################

class Keypatch_Asm:
    # supported architectures
    arch_lists = {
        "X86 16-bit": (KS_ARCH_X86, KS_MODE_16),                # X86 16-bit
        "X86 32-bit": (KS_ARCH_X86, KS_MODE_32),                # X86 32-bit
        "X86 64-bit": (KS_ARCH_X86, KS_MODE_64),                # X86 64-bit
        "ARM": (KS_ARCH_ARM, KS_MODE_ARM),                      # ARM
        "ARM Thumb": (KS_ARCH_ARM, KS_MODE_THUMB),              # ARM Thumb
        "ARM64 (ARMV8)": (KS_ARCH_ARM64, KS_MODE_LITTLE_ENDIAN),# ARM64
        "Hexagon": (KS_ARCH_HEXAGON, KS_MODE_BIG_ENDIAN),       # Hexagon
        "Mips32": (KS_ARCH_MIPS, KS_MODE_MIPS32),               # Mips32
        "Mips64": (KS_ARCH_MIPS, KS_MODE_MIPS64),               # Mips64
        "PowerPC 32": (KS_ARCH_PPC, KS_MODE_PPC32),             # PPC32
        "PowerPC 64": (KS_ARCH_PPC, KS_MODE_PPC64),             # PPC64
        "Sparc 32": (KS_ARCH_SPARC, KS_MODE_SPARC32),           # Sparc32
        "Sparc 64": (KS_ARCH_SPARC, KS_MODE_SPARC64),           # Sparc64
        "SystemZ": (KS_ARCH_SYSTEMZ, KS_MODE_BIG_ENDIAN),       # SystemZ
    }

    endian_lists = {
        "Little Endian": KS_MODE_LITTLE_ENDIAN,                 # little endian
        "Big Endian": KS_MODE_BIG_ENDIAN,                       # big endian
    }

    syntax_lists = {
        "Intel": KS_OPT_SYNTAX_INTEL,
        "Nasm": KS_OPT_SYNTAX_NASM,
        "AT&T": KS_OPT_SYNTAX_ATT
    }

    def __init__(self, arch=None, mode=None):
        # update current arch and mode
        self.update_hardware_mode()

        # override arch & mode if provided
        if arch is not None:
            self.arch = arch
        if mode is not None:
            self.mode = mode

        # IDA uses Intel syntax by default
        self.syntax = KS_OPT_SYNTAX_INTEL

    # return Keystone arch & mode (with endianess included)
    @staticmethod
    def get_hardware_mode():
        (arch, mode) = (None, None)

        # heuristically detect hardware setup
        info = idaapi.get_inf_structure()
        
        try:
            cpuname = info.procname.lower()
        except:
            cpuname = info.procName.lower()

        try:
            # since IDA7 beta 3 (170724) renamed inf.mf -> is_be()/set_be()
            is_be = idaapi.cvar.inf.is_be()
        except:
            # older IDA versions
            is_be = idaapi.cvar.inf.mf
        # print("Keypatch BIG_ENDIAN = %s" %is_be)
        
        if cpuname == "metapc":
            arch = KS_ARCH_X86
            if info.is_64bit():
                mode = KS_MODE_64
            elif info.is_32bit():
                mode = KS_MODE_32
            else:
                mode = KS_MODE_16
        elif cpuname.startswith("arm"):
            # ARM or ARM64
            if info.is_64bit():
                arch = KS_ARCH_ARM64
                if is_be:
                    mode = KS_MODE_BIG_ENDIAN
                else:
                    mode = KS_MODE_LITTLE_ENDIAN
            else:
                arch = KS_ARCH_ARM
                # either big-endian or little-endian
                if is_be:
                    mode = KS_MODE_ARM | KS_MODE_BIG_ENDIAN
                else:
                    mode = KS_MODE_ARM | KS_MODE_LITTLE_ENDIAN
        elif cpuname.startswith("sparc"):
            arch = KS_ARCH_SPARC
            if info.is_64bit():
                mode = KS_MODE_SPARC64
            else:
                mode = KS_MODE_SPARC32
            if is_be:
                mode |= KS_MODE_BIG_ENDIAN
            else:
                mode |= KS_MODE_LITTLE_ENDIAN
        elif cpuname.startswith("ppc"):
            arch = KS_ARCH_PPC
            if info.is_64bit():
                mode = KS_MODE_PPC64
            else:
                mode = KS_MODE_PPC32
            if cpuname == "ppc":
                # do not support Little Endian mode for PPC
                mode += KS_MODE_BIG_ENDIAN
        elif cpuname.startswith("mips"):
            arch = KS_ARCH_MIPS
            if info.is_64bit():
                mode = KS_MODE_MIPS64
            else:
                mode = KS_MODE_MIPS32
            if is_be:
                mode |= KS_MODE_BIG_ENDIAN
            else:
                mode |= KS_MODE_LITTLE_ENDIAN
        elif cpuname.startswith("systemz") or cpuname.startswith("s390x"):
            arch = KS_ARCH_SYSTEMZ
            mode = KS_MODE_BIG_ENDIAN

        return (arch, mode)

    def update_hardware_mode(self):
        (self.arch, self.mode) = self.get_hardware_mode()

    # normalize assembly code
    # remove comment at the end of assembly code
    @staticmethod
    def asm_normalize(text):
        text = ' '.join(text.split())
        if text.rfind(';') != -1:
            return text[:text.rfind(';')].strip()

        return text.strip()

    @staticmethod
    # check if input address is valid
    # return
    #       -1  invalid address at target binary
    #        0  type mismatch of input address
    #        1  valid address at target binary
    def check_address(address):
        try:
            if idc.is_mapped(address):
                return 1
            else:
                return -1
        except:
            # invalid type
            return 0

    ### resolve IDA names from input asm code
    # todo: a better syntax parser for all archs
    def ida_resolve(self, assembly, address=idc.BADADDR):
        def _resolve(_op, ignore_kw=True):
            names = re.findall(r"[\$a-z0-9_:\.]+", _op, re.I)

            # try to resolve all names
            for name in names:
                # ignore known keywords
                if ignore_kw and name in ('byte', 'near', 'short', 'word', 'dword', 'ptr', 'offset'):
                    continue

                sym = name

                # split segment reg
                parts = name.partition(':')
                if parts[2] != '':
                    sym = parts[2]

                (typ, value) = idaapi.get_name_value(address, sym)

                # skip if name doesn't exist or segment / segment registers
                if typ in (idaapi.NT_SEG, idaapi.NT_NONE):
                    continue

                _op = _op.replace(sym, '0x{0:X}'.format(value))

            return _op

        if self.check_address(address) == 0:
            print("Keypatch: WARNING: invalid input address {0}".format(address))
            return assembly

        # for now, we only support IDA name resolve for X86, ARM, ARM64, MIPS, PPC, SPARC
        if not (self.arch in (KS_ARCH_X86, KS_ARCH_ARM, KS_ARCH_ARM64, KS_ARCH_MIPS, KS_ARCH_PPC, KS_ARCH_SPARC)):
            return assembly

        _asm = assembly.partition(' ')
        mnem = _asm[0]
        opers = _asm[2].split(',')

        for idx, op in enumerate(opers):
            _op = list(op.partition('['))
            ignore_kw = True
            if _op[1] == '':
                _op[2] = _op[0]
                _op[0] = ''
            else:
                _op[0] = _resolve(_op[0], ignore_kw=True)
                ignore_kw = False

            _op[2] = _resolve(_op[2], ignore_kw=ignore_kw)

            opers[idx] = ''.join(_op)

        asm = "{0} {1}".format(mnem, ','.join(opers))
        return asm

    # return bytes of instruction or data
    # return None on failure
    def ida_get_item(self, address, hex_output=False):
        if self.check_address(address) != 1:
            # not a valid address
            return (None, 0)

        # return None if address is in the middle of instruction / data
        if address != idc.get_item_head(address):
            return (None, 0)

        size = idc.get_item_size(address)
        item = idc.get_bytes(address, size)

        if item is None:
            return (None, 0)

        if hex_output:
            item = to_hexstr(item)

        return (item, size)

    @staticmethod
    def get_op_dtype_name(ea, op_idx):
        dtyp_lists = {
            idaapi.dt_byte: 'byte',     #  8 bit
            idaapi.dt_word: 'word',     #  16 bit
            idaapi.dt_dword: 'dword',   #  32 bit
            idaapi.dt_float: 'dword',   #  4 byte
            idaapi.dt_double: 'dword',  #  8 byte
            #idaapi.dt_tbyte = 5        #  variable size (ph.tbyte_size)
            #idaapi.dt_packreal = 6         #  packed real format for mc68040
            idaapi.dt_qword: 'qword',   #  64 bit
            idaapi.dt_byte16: 'xmmword',#  128 bit
            #idaapi.dt_code = 9         #  ptr to code (not used?)
            #idaapi.dt_void = 10        #  none
            #idaapi.dt_fword = 11       #  48 bit
            #idaapi.dt_bitfild = 12     #  bit field (mc680x0)
            #idaapi.dt_string = 13      #  pointer to asciiz string
            #idaapi.dt_unicode = 14     #  pointer to unicode string
            #idaapi.dt_3byte = 15       #  3-byte data
            #idaapi.dt_ldbl = 16        #  long double (which may be different from tbyte)
            idaapi.dt_byte32: 'ymmword',# 256 bit
        }
        idc.get_dtype
        dtype, dtyp_size = get_dtype(ea, op_idx)
        if dtype == idaapi.dt_tbyte and dtyp_size == 10:
          return 'xword'

        dtyp_name = dtyp_lists.get(dtype, None)

        return dtyp_name

    # return asm instructions from start to end
    def ida_get_disasm_range(self, start, end):
        codes = []
        while start < end:
            asm = self.asm_normalize(idc.GetDisasm(start))
            if asm == None:
                asm = ''
            codes.append(asm)
            start = start + idc.get_item_size(start)

        return codes

    # get disasm from IDA
    # return '' on invalid address
    def ida_get_disasm(self, address, fixup=False):

        def GetMnem(asm):
            sp = asm.find(' ')
            if sp == -1:
                return asm
            return asm[:sp]

        if self.check_address(address) != 1:
            # not a valid address
            return ''

        # return if address is in the middle of instruction / data
        if address != idc.get_item_head(address):
            return ''

        asm = self.asm_normalize(idc.GetDisasm(address))
        # for now, only support IDA syntax fixup for Intel CPU
        if not fixup or self.arch != KS_ARCH_X86:
            return asm

        # KS_ARCH_X86 mode
        # rebuild disasm code from IDA
        i = 0
        mnem = GetMnem(asm)
        if mnem == '' or mnem in ('rep', 'repne', 'repe'):
            return asm

        opers = []
        while idc.get_operand_type(address, i) > 0 and i < 6:
            t = idc.get_operand_type(address, i)
            o = idc.print_operand(address, i)

            if t in (idc.o_mem, idc.o_displ):
                parts = list(o.partition(':'))
                if parts[2] == '':
                    parts[2] = parts[0]
                    parts[0] = ''

                if '[' not in parts[2]:
                    parts[2] = '[{0}]'.format(parts[2])

                o = ''.join(parts)

                if 'ptr ' not in o:
                    dtyp_name = self.get_op_dtype_name(address, i)
                    if dtyp_name != None:
                        o = "{0} ptr {1}".format(dtyp_name, o)

            opers.append(o)
            i += 1

        asm = mnem
        for o in opers:
            if o != '':
                asm = "{0} {1},".format(asm, o)

        asm = asm.strip(',')
        return asm

    # assemble code with Keystone
    # return (encoding, count), or (None, 0) on failure
    def assemble(self, assembly, address, arch=None, mode=None, syntax=None):

        # return assembly with arithmetic equation evaluated
        def eval_operand(assembly, start, stop, prefix=''):
            imm = assembly[start+1:stop]
            try:
                eval_imm = eval(imm)
                if eval_imm > 0x80000000:
                    eval_imm = 0xffffffff - eval_imm
                    eval_imm += 1
                    eval_imm = -eval_imm
                return assembly.replace(prefix + imm, prefix + hex(eval_imm))
            except:
                return assembly

        # IDA uses different syntax from Keystone
        # sometimes, we can convert code to be consumable by Keystone
        def fix_ida_syntax(assembly):

            # return True if this insn needs to be fixed
            def check_arm_arm64_insn(arch, mnem):
                if arch == KS_ARCH_ARM:
                    if mnem.startswith("ldr") or mnem.startswith("str"):
                        return True
                    return False
                elif arch == KS_ARCH_ARM64:
                    if mnem.startswith("ldr") or mnem.startswith("str"):
                        return True
                    return mnem in ("stp")
                return False

            # return True if this insn needs to be fixed
            def check_ppc_insn(mnem):
                return mnem in ("stw")

            # replace the right most string occurred
            def rreplace(s, old, new):
                li = s.rsplit(old, 1)
                return new.join(li)

            # convert some ARM pre-UAL assembly to UAL, so Keystone can handle it
            # example: streqb --> strbeq
            def fix_arm_ual(mnem, assembly):
                # TODO: this is not an exhaustive list yet
                if len(mnem) != 6:
                    return assembly

                if (mnem[-1] in ('s', 'b', 'h', 'd')):
                    #print(">> 222", mnem[3:5])
                    if mnem[3:5] in ("cc", "eq", "ne", "hs", "lo", "mi", "pl", "vs", "vc", "hi", "ls", "ge", "lt", "gt", "le", "al"):
                        return assembly.replace(mnem, mnem[:3] + mnem[-1] + mnem[3:5], 1)

                return assembly

            if self.arch != KS_ARCH_X86:
                assembly = assembly.lower()
            else:
                # Keystone does not support immediate 0bh, but only 0Bh
                assembly = assembly.upper()

            # however, 0X must be converted to 0x
            # Keystone should fix this limitation in the future
            assembly = assembly.replace("0X", " 0x")

            _asm = assembly.partition(' ')
            mnem = _asm[0]
            if mnem == '':
                return assembly

            # for PPC, Keystone does not accept registers with 'r' prefix,
            # but only the number behind. lets try to fix that here by
            # removing the prefix 'r'.
            if self.arch == KS_ARCH_PPC:
                for n in range(32):
                    r = " r%u," %n
                    if r in assembly:
                        assembly = assembly.replace(r, " %u," %n)
                for n in range(32):
                    r = "(r%u)" %n
                    if r in assembly:
                        assembly = assembly.replace(r, "(%u)" %n)
                for n in range(32):
                    r = ", r%u" %n
                    if assembly.endswith(r):
                        assembly = rreplace(assembly, r, ", %u" %n)

            if self.arch == KS_ARCH_X86:
                if mnem == "RETN":
                    # replace retn with ret
                    return assembly.replace('RETN', 'RET', 1)
                if 'OFFSET ' in assembly:
                    return assembly.replace('OFFSET ', ' ')
                if mnem in ('CALL', 'JMP') or mnem.startswith('LOOP'):
                    # remove 'NEAR PTR'
                    if ' NEAR PTR ' in assembly:
                        return assembly.replace(' NEAR PTR ', ' ')
                elif mnem[0] == 'J':
                    # JMP instruction
                    if ' SHORT ' in assembly:
                        # remove ' short '
                        return assembly.replace(' SHORT ', ' ')
            elif self.arch in (KS_ARCH_ARM, KS_ARCH_ARM64, KS_ARCH_PPC):
                # *** ARM
                # LDR     R1, [SP+rtld_fini],#4
                # STR     R2, [SP,#-4+rtld_fini]!
                # STR     R0, [SP,#fini]!
                # STR     R12, [SP,#4+var_8]!

                # *** ARM64
                # STP     X29, X30, [SP,#-0x10+var_150]!
                # STR     W0, [X29,#0x150+var_8]
                # LDR     X0, [X0,#(qword_4D6678 - 0x4D6660)]
                # TODO:
                # ADRP    X19, #interactive@PAGE

                # *** PPC
                # stw     r5, 0x120+var_108(r1)
                
                if self.arch == KS_ARCH_ARM and mode == KS_MODE_THUMB:
                    assembly = assembly.replace('movt.w', 'movt')

                if self.arch == KS_ARCH_ARM:
                    #print(">> before UAL fix: ", assembly)
                    assembly = fix_arm_ual(mnem, assembly)
                    #print(">> after UAL fix: ", assembly)

                if check_arm_arm64_insn(self.arch, mnem) or (("[" in assembly) and ("]" in assembly)):
                    bang = assembly.find('#')
                    bracket = assembly.find(']')
                    if bang != -1 and bracket != -1 and bang < bracket:
                        return eval_operand(assembly, bang, bracket, '#')
                    elif '+0x0]' in assembly:
                        return assembly.replace('+0x0]', ']')
                elif check_ppc_insn(mnem):
                    start = assembly.find(', ')
                    stop = assembly.find('(')
                    if start != -1 and stop != -1 and start < stop:
                        return eval_operand(assembly, start, stop)
            return assembly

        def is_thumb(address):
            return idc.get_sreg(address, 'T') == 1

        if self.check_address(address) == 0:
            return (None, 0)

        # use default syntax, arch and mode if not provided
        if syntax is None:
            syntax = self.syntax
        if arch is None:
            arch = self.arch
        if mode is None:
            mode = self.mode

        if arch == KS_ARCH_ARM and is_thumb(address):
            mode = KS_MODE_THUMB

        try:
            ks = Ks(arch, mode)
            if arch == KS_ARCH_X86:
                ks.syntax = syntax
            encoding, count = ks.asm(fix_ida_syntax(assembly), address)
        except KsError as e:
            # keep the below code for debugging
            #print("Keypatch Error: {0}".format(e))
            #print("Original asm: {0}".format(assembly))
            #print("Fixed up asm: {0}".format(fix_ida_syntax(assembly)))
            encoding, count = None, 0

        return (encoding, count)


    # patch at address, return the number of written bytes & original data
    # this process can fail in some cases
    @staticmethod
    def patch_raw(address, patch_data, size):
        ea = address
        orig_data = ''

        while ea < (address + size):

            if not idc.has_value(idc.get_full_flags(ea)):
                print("Keypatch: FAILED to read data at 0x{0:X}".format(ea))
                break

            orig_byte = idc.get_wide_byte(ea)
            orig_data += chr(orig_byte)
            patch_byte = ord(patch_data[ea - address])

            if patch_byte != orig_byte:
                # patch one byte
                if idaapi.patch_byte(ea, patch_byte) != 1:
                    print("Keypatch: FAILED to patch byte at 0x{0:X} [0x{1:X}]".format(ea, patch_byte))
                    break
            ea += 1
        return (ea - address, orig_data)

    # patch at address, return the number of written bytes & original data
    # on patch failure, we revert to the original code, then return (None, None)
    def patch(self, address, patch_data, size):
        # save original function end to fix IDA re-analyze issue after patching
        orig_func_end = idc.get_func_attr(address, idc.FUNCATTR_END)

        (patched_len, orig_data) = self.patch_raw(address, patch_data, size)

        if size != patched_len:
            # patch failure
            if patched_len > 0:
                # revert the changes
                (rlen, _) = self.patch_raw(address, orig_data, patched_len)
                if rlen == patched_len:
                    print("Keypatch: successfully reverted changes of {0:d} byte(s) at 0x{1:X} [{2}]".format(
                                        patched_len, address, to_hexstr(orig_data)))
                else:
                    print("Keypatch: FAILED to revert changes of {0:d} byte(s) at 0x{1:X} [{2}]".format(
                                        patched_len, address, to_hexstr(orig_data)))

            return (None, None)

        # ask IDA to re-analyze the patched area
        if orig_func_end == idc.BADADDR:
            # only analyze patched bytes, otherwise it would take a lot of time to re-analyze the whole binary
            idaapi.plan_and_wait(address, address + patched_len + 1)
        else:
            idaapi.plan_and_wait(address, orig_func_end)

            # try to fix IDA function re-analyze issue after patching
            idc.set_func_end(address, orig_func_end)

        return (patched_len, orig_data)

    # return number of bytes patched
    # return
    #    0  Invalid assembly
    #   -1  PatchByte failure
    #   -2  Can't read original data
    #   -3  Invalid address
    def patch_code(self, address, assembly, syntax, padding, save_origcode, orig_asm=None, patch_data=None, patch_comment=None, undo=False):
        global patch_info

        if self.check_address(address) != 1:
            # not a valid address
            return -3

        orig_comment = get_comment(address)
        if orig_comment is None:
            orig_comment = ''

        nop_comment = ""
        padding_len = 0
        if not undo:
            # we are patching via Patcher
            (orig_encoding, orig_len) = self.ida_get_item(address)
            if (orig_encoding, orig_len) == (None, 0):
                return -2

            (encoding, count) = self.assemble(assembly, address, syntax=syntax)
            if encoding is None:
                return 0

            patch_len = len(encoding)
            patch_data = ''.join(chr(c) for c in encoding)

            if patch_data == orig_encoding:
                #print("Keypatch: no need to patch, same encoding data [{0}] at 0x{1:X}".format(to_hexstr(orig_encoding), address))
                return orig_len

            # for now, only support NOP padding on Intel CPU
            if padding and self.arch == KS_ARCH_X86:
                if patch_len < orig_len:
                    padding_len = orig_len - patch_len
                    patch_len = orig_len
                    patch_data = patch_data.ljust(patch_len, "\x90")
                elif patch_len > orig_len:
                    patch_end = address + patch_len - 1
                    ins_end = idc.get_item_end(patch_end)
                    padding_len = ins_end - patch_end - 1

                    if padding_len > 0:
                        patch_len = ins_end - address
                        patch_data = patch_data.ljust(patch_len, "\x90")

                if padding_len > 0:
                    nop_comment = "\nKeypatch padded NOP to next boundary: {0} bytes".format(padding_len)

            orig_asm = self.ida_get_disasm_range(address, address + patch_len)
        else:
            # we are reverting the change via "Undo" menu
            patch_len = len(patch_data)

        (plen, p_orig_data) = self.patch(address, patch_data, patch_len)
        if plen is None:
            # failed to patch
            return -1

        if not undo: # we are patching
            new_patch_comment = None
            if save_origcode is True:
                # append original instruction to comments
                if orig_comment == '':
                    new_patch_comment = "Keypatch modified this from:\n  {0}{1}".format('\n  '.join(orig_asm), nop_comment)
                else:
                    new_patch_comment = "\nKeypatch modified this from:\n  {0}{1}".format('\n  '.join(orig_asm), nop_comment)

                new_comment = "{0}{1}".format(orig_comment, new_patch_comment)
                set_comment(address, new_comment)

            if padding_len == 0:
                print("Keypatch: successfully patched {0:d} byte(s) at 0x{1:X} from [{2}] to [{3}]".format(plen,
                                        address, to_hexstr(p_orig_data), to_hexstr(patch_data)))
            else:
                print("Keypatch: successfully patched {0:d} byte(s) at 0x{1:X} from [{2}] to [{3}], with {4} byte(s) NOP padded".format(plen,
                                        address, to_hexstr(p_orig_data), to_hexstr(patch_data), padding_len))
            # save this patching for future "undo"
            patch_info.append((address, assembly, p_orig_data, new_patch_comment))
        else:   # we are reverting
            if patch_comment:
                # clean previous IDA comment by replacing it with ''
                new_comment = orig_comment.replace(patch_comment, '')
                set_comment(address, new_comment)

            print("Keypatch: successfully reverted {0:d} byte(s) at 0x{1:X} from [{2}] to [{3}]".format(plen,
                                        address, to_hexstr(p_orig_data), to_hexstr(patch_data)))

        return plen

    # fill a range of code [addr_begin, addr_end].
    # return the length of patched area
    # on failure, return 0 = wrong input, -1 = failed to patch
    def fill_code(self, addr_begin, addr_end, assembly, syntax, padding, save_origcode, orig_asm=None):
        # treat input as assembly code first
        (encoding, _) =  self.assemble(assembly, addr_begin, syntax=syntax)

        if encoding is None:
            # input might be a hexcode string. try to convert it to raw bytes
            encoding = convert_hexstr(assembly)

        if encoding is None:
            # invalid input: this is neither assembly nor hexcode string
            return 0

        # save original assembly code before overwritting them
        orig_asm = self.ida_get_disasm_range(addr_begin, addr_end)

        # save original comment at addr_begin
        # TODO: save comments in this range, but how to interleave them?
        orig_comment = get_comment(addr_begin)
        if orig_comment is None:
            orig_comment = ''

        patch_data = ""
        assembly_new = []
        size = addr_end - addr_begin
        # calculate filling data
        encode_chr = ''.join(chr(c) for c in encoding)
        while True:
            if len(patch_data) + len(encode_chr) <= size:
                patch_data = patch_data + encode_chr
                assembly_new += [assembly.strip()]
            else:
                break

        # for now, only support NOP padding on Intel CPU
        if padding and self.arch == KS_ARCH_X86:
            for i in range(size -len(patch_data)):
                assembly_new += ["nop"]
            patch_data = patch_data.ljust(size, "\x90")

        (plen, p_orig_data) = self.patch(addr_begin, patch_data, len(patch_data))
        if plen is None:
            # failed to patch
            return -1

        new_patch_comment = ''
        # append original instruction to comments
        if save_origcode is True:
            if orig_comment == '':
                new_patch_comment = "Keypatch filled range [0x{0:X}:0x{1:X}] ({2} bytes), replaced:\n  {3}".format(addr_begin, addr_end - 1, addr_end - addr_begin, '\n  '.join(orig_asm))
            else:
                new_patch_comment = "\nKeypatch filled range [0x{0:X}:0x{1:X}] ({2} bytes), replaced:\n  {3}".format(addr_begin, addr_end - 1, addr_end - addr_begin, '\n  '.join(orig_asm))

            new_comment = "{0}{1}".format(orig_comment, new_patch_comment)
            set_comment(addr_begin, new_comment)

        print("Keypatch: successfully filled range [0x{0:X}:0x{1:X}] ({2} bytes) with \"{3}\", replaced \"{4}\"".format(
                    addr_begin, addr_end - 1, addr_end - addr_begin, assembly, '; '.join(orig_asm)))

        # save this modification for future "undo"
        patch_info.append((addr_begin, '\n  '.join(assembly_new), p_orig_data, new_patch_comment))

        return plen


    ### Form helper functions
    @staticmethod
    def dict_to_ordered_list(dictionary):
        l = sorted(list(dictionary.items()), key=lambda t: t[0], reverse=False)
        keys = [i[0] for i in l]
        values = [i[1] for i in l]

        return (keys, values)

    def get_value_by_idx(self, dictionary, idx, default=None):
        (keys, values) = self.dict_to_ordered_list(dictionary)

        try:
            val = values[idx]
        except IndexError:
            val = default

        return val

    def find_idx_by_value(self, dictionary, value, default=None):
        (keys, values) = self.dict_to_ordered_list(dictionary)

        try:
            idx = values.index(value)
        except:
            idx = default

        return idx

    def get_arch_by_idx(self, idx):
        return self.get_value_by_idx(self.arch_lists, idx)

    def find_arch_idx(self, arch, mode):
        return self.find_idx_by_value(self.arch_lists, (arch, mode))

    def get_syntax_by_idx(self, idx):
        return self.get_value_by_idx(self.syntax_lists, idx, self.syntax)

    def find_syntax_idx(self, syntax):
        return self.find_idx_by_value(self.syntax_lists, syntax)
    ### /Form helper functions

class OpaqueAnalyzeRetValues(Enum):
    NOT_OPAQUE = 1
    ALWAYS_JMP = 2
    NEVER_JMP = 3
    CHANGE_MODE = 4
    ERROR = 5

class OpaqueIdaPatcherRetValues(Enum):
    SKIP_BLOCK = 1
    RETURN_TO_MAINLOOP = 2
    PROCEED = 3


class FunctionOpaqueIdentifier:

    def __init__(self):
        if not DEBUG_VERBOSE:
            logger = logging.getLogger('angr')
            logger.propagate = False
            logger = logging.getLogger('claripy')
            logger.propagate = False
            logger = logging.getLogger('cle')
            logger.propagate = False
            logger = logging.getLogger('pyvex')
            logger.propagate = False
        
        logging.basicConfig(stream=sys.stderr, level=logging.DEBUG)

        self.current_func_name = idaapi.get_func_name(idaapi.get_screen_ea())
        self.current_func = ida_funcs.get_func(idaapi.get_screen_ea())
        self.ida_inputfile_path = idaapi.get_input_file_path()
        self.inputfile_path = idaapi.get_input_file_path().replace('.exe', '') + TOPAQUE_ANGR_ADDITION + '.exe'
    
        logging.debug("Reapplying IDA patches on the operating on exe..")
        if os.path.isfile(self.inputfile_path):
            try:
                os.remove(self.inputfile_path)
            except:
                self.inputfile_path = idaapi.get_input_file_path().replace('.exe', '') + TOPAQUE_ANGR_ADDITION + '_f.exe'
        shutil.copy(self.ida_inputfile_path, self.inputfile_path)
        self.apply_patch_bytes(self.current_func.start_ea, self.current_func.end_ea)
        if not TESTS_MODE and not FIX_DEFLOW_MODE:
            logging.debug("Loading angr...")
            self.angr_prj = angr.Project(self.inputfile_path, load_options={'auto_load_libs':False}) 
        else:
            logging.debug("WE ARE IN TESTS MODE. ANGR IS NOT LOADED.")
        self.kp_asm = Keypatch_Asm()
        print("Down the rabbit hole..")
        print("[Topaqueminator] Identifying opaque predicates on function:", self.current_func_name)
        
        self.run()



    def check_if_jmp_x86_64(self, inst_ea):
        inst = idautils.DecodeInstruction(inst_ea).get_canon_mnem()
        if 'jmp' in inst:
            return True
        return False
    


    def check_conditional_x86_64(self, inst_ea):
        inst = idautils.DecodeInstruction(inst_ea).get_canon_mnem()
        if 'j' in inst and 'jmp' not in inst:
            return True
        return False



    def convert_bytes_to_asm_instructions(self, bytes):
        md = Cs(CS_ARCH_X86, CS_MODE_64)
        asm_str = ""
        for inst in md.disasm(bytes, self.current_func.start_ea):
            print(str(hex(inst.address)) + " " + inst.mnemonic + " " + inst.op_str + ";")
            asm_str += str(hex(inst.address)) + " " + inst.mnemonic + " " + inst.op_str + ";" +"\n"
        return asm_str



    def get_the_location_of_the_jump(self, jmp_inst_ea):
        if self.check_if_jmp_x86_64(jmp_inst_ea) or self.check_conditional_x86_64(jmp_inst_ea):
            return idc.get_operand_value(jmp_inst_ea, 0)
        return None



    def fix_function_end(self, new_end_ea):
        logging.debug("Patching function end: {}".format(hex(new_end_ea)))
        #self.current_func.end_ea = new_end_ea
        func_ea = idc.next_head(idc.prev_head(new_end_ea))
        res = ida_funcs.set_func_end(func_ea, new_end_ea)
        logging.debug("set end result: {}".format(res))
        self.referesh_screen_and_reanalyze()



    def test_if_instruction_broken(self, inst_addr):
        if 'j' in idc.print_insn_mnem(inst_addr):
            if idc.generate_disasm_line(inst_addr, 0)[-2:-1] == "+" and idc.generate_disasm_line(inst_addr,0)[-1:].isdigit():
                logging.debug("Broken Instruction: {}".format(hex(inst_addr)))
                return True
        return False



    def test_if_instruction_jmp_qword_broken(self, inst_addr):
        if 'j' in idc.print_insn_mnem(inst_addr):
            if "+" in idc.generate_disasm_line(inst_addr, 0) and "qword" in idc.generate_disasm_line(inst_addr, 0):
                logging.debug("Broken qword instruction: {}".format(hex(inst_addr)))
                return True
        return False
    


    def fix_jmp_qword(self, inst_addr):
        code_addr = self.get_the_location_of_the_jump(inst_addr)
        ret = idc.create_insn(code_addr)
        self.referesh_screen_and_reanalyze()
        return ret



    def fix_jmp_inst_on_opaque_predicate(self, opaque_addr):
        if self.test_if_instruction_broken(opaque_addr) or self.test_if_instruction_jmp_qword_broken(opaque_addr):
            logging.debug("Fixing broken instruction...")
            code_addr = self.get_the_location_of_the_jump(opaque_addr)
            fix_addr = code_addr - int(idc.generate_disasm_line(opaque_addr,0)[-1:])
            ret_val_del = idc.del_items(fix_addr,1)
            self.referesh_screen_and_reanalyze()
            ret_val_create = idc.create_insn(code_addr)
            self.referesh_screen_and_reanalyze()
            if ret_val_create == 0 or ret_val_del == False:
                return False
            return True
        return False



    def referesh_screen_and_reanalyze(self):
        ida_auto.auto_wait()
        idaapi.reanalyze_function(self.current_func)
        ida_auto.auto_wait()



    def tompaque_patch_code(self, instruction_string, address):
        cmt = "Topaque Patched from: {}".format(idc.generate_disasm_line(address, 0)) + "\n" + "\tTopaque Patched to: " + instruction_string
        self.kp_asm.patch_code(address, instruction_string, KS_ARCH_X86, 1, 0, patch_comment=cmt)
        idc.set_cmt(address, cmt, False)
        self.referesh_screen_and_reanalyze()
        if APPLY_PATCHES_ON_INPUT_FILE:
            self.apply_patch_bytes(self.current_func.start_ea, self.current_func.end_ea)
        self.referesh_screen_and_reanalyze()



    def fix_opaque_predicate(self, opaque_addr, opaque_type_num, different_addr = None):
        if opaque_type_num == OpaqueAnalyzeRetValues.ALWAYS_JMP:
            if different_addr is not None:
                jmp_location_str = different_addr
            else:
                jmp_location_str = hex(self.get_the_location_of_the_jump(opaque_addr))
            raw_assembly = "jmp " + jmp_location_str
            self.tompaque_patch_code(raw_assembly, opaque_addr)
            self.referesh_screen_and_reanalyze()
        elif opaque_type_num == OpaqueAnalyzeRetValues.NEVER_JMP:
            raw_assembly = "nop"
            self.tompaque_patch_code(raw_assembly, opaque_addr)
            self.referesh_screen_and_reanalyze()



    def fix_undef_code_places(self, predicate_ea):
        idc.create_insn(predicate_ea)
        self.referesh_screen_and_reanalyze()



    def fix_using_single_block_analysis(self):
         #single block anaylsis.
        while(SINGLE_BLOCK_MODE and self.patch_opaque_predicates_loop(SINGLE_BLOCK_MODE)):
            if not CONTINUOS_ANALYZING:
                break
        return
    
    def fix_using_whole_function_analysis(self):
        #next iteration start each time from start (whol func analysis)
        while (self.patch_opaque_predicates_loop(False)):
             if not CONTINUOS_ANALYZING:
                break
    
    def deflow_main_loop(self):
        pass
        
    def fix_using_deflow(self):
       f = FixCode(self.current_func.start_ea)
       f.process()




    def unravel_jmp_chains(self, opaque_addr):
        if self.check_conditional_x86_64(opaque_addr) or self.check_if_jmp_x86_64(opaque_addr):
            next_jmp_loc = self.get_the_location_of_the_jump(opaque_addr)
            init_inst_string = idautils.DecodeInstruction(opaque_addr).get_canon_mnem()
            jmp_inst_string = idautils.DecodeInstruction(next_jmp_loc).get_canon_mnem()
            counter = 0
            while(self.check_conditional_x86_64(next_jmp_loc) and (init_inst_string == jmp_inst_string)):
                counter += 1
                next_jmp_loc = self.get_the_location_of_the_jump(next_jmp_loc)
                logging.debug(f"{next_jmp_loc}")
                if next_jmp_loc is None:
                    logging.debug("unravel_jmp_chains next_jmp_loc is None ")
                jmp_inst_string = idautils.DecodeInstruction(next_jmp_loc).get_canon_mnem()
                
            return (counter > 0), next_jmp_loc
        return False, None



    def ida_pre_explore_patch(self, block, is_single_block_mode):
        global last_found_opaque_predicate
        global is_last_predicate_mode
        logging.debug("INSIDE ida_pre_explore_patch(self, block, is_single_block_mode):")
        opaque_addr = idc.prev_head(block.end_ea)
        #fix idas stupid shit, than continue with trying to analyze opaque.
        is_unraveled, final_jmp_loc = self.unravel_jmp_chains(opaque_addr)
        
        if is_unraveled:
            inst_cmd_str = idautils.DecodeInstruction(opaque_addr).get_canon_mnem()
            logging
            self.tompaque_patch_code(inst_cmd_str + " {}".format(hex(final_jmp_loc)), opaque_addr)
            return OpaqueIdaPatcherRetValues.RETURN_TO_MAINLOOP
        
        if is_single_block_mode:
            #handle the case of single jump blocks. they are NOT opaque.
            if 'j' in idc.print_insn_mnem(block.start_ea):
                logging.debug("single block start from JMP instruction is NOT opaque. block_start_ea: {} ".format(hex(block.start_ea)))
                return OpaqueIdaPatcherRetValues.SKIP_BLOCK
        
        if is_last_predicate_mode:
            if(last_found_opaque_predicate == None):
                logging.debug("we are in last found mode, but last_found_opaq is None.")
                raise "Error: we are in last found mode, but last_found_opaq is None."
                
            if block.start_ea < self.get_the_location_of_the_jump(last_found_opaque_predicate):
                logging.debug("block.start_ea < self.get_the_location_of_the_jump(last_found_opaque_predicate), skipping block")
                return OpaqueIdaPatcherRetValues.SKIP_BLOCK

        if self.check_conditional_x86_64(opaque_addr) or self.check_if_jmp_x86_64(opaque_addr):
            jump_to_operand_location_value = self.get_the_location_of_the_jump(opaque_addr)
            if self.test_if_instruction_broken(opaque_addr) or self.test_if_instruction_jmp_qword_broken(opaque_addr):
                if idaapi.get_func(jump_to_operand_location_value) is None: #was jump + 0x15
                    logging.debug("Chunk does not belone to any function, create and set func end")
                    #func_ea = block.end_ea
                    if self.fix_jmp_inst_on_opaque_predicate(opaque_addr) == False:
                        logging.debug("UNABLE TO FIX BROKEN INSTRUCTION, PROCEEDING..")
                        return OpaqueIdaPatcherRetValues.PROCEED
                    #idc.set_func_end(func_ea, jump_to_operand_location_value + 0x10)
                    self.referesh_screen_and_reanalyze()
                    return OpaqueIdaPatcherRetValues.RETURN_TO_MAINLOOP
                if jump_to_operand_location_value < opaque_addr:
                    #broken instuction + opaque address
                    #TODO: if we are at the beginning, return false and let the user decide.
                    raise ValueError("SELF DECRYPTING CODE FOUND. PLEASE HANDLY MANUALLY.")
                    #self.tompaque_patch_code("nop", opaque_addr)
                    #return OpaqueIdaPatcherRetValues.RETURN_TO_MAINLOOP
            if jump_to_operand_location_value < opaque_addr:
                    #we are in a loop, analyze the next instruction. no point in analyzing this one.  
                    return OpaqueIdaPatcherRetValues.SKIP_BLOCK
            else:
                #cond jmp is 2 bytes
                if idc.next_head(opaque_addr) - (opaque_addr+CONDITIONAL_JMP_SIZE_IN_BYTES) == 1:
                    #is conditional, instruction not broken, next head is 1 bytes fucking close.
                    #they are doing that stupid shit again. ja jb [x] jb [x
                    logging.debug("is conditional, instruction not broken, next head is 1 bytes fucking close.")
                    logging.debug("they are doing that stupid shit again. ja jb [x] jb [x")
                    if self.check_if_jmp_x86_64(opaque_addr):
                        return OpaqueIdaPatcherRetValues.PROCEED
                    if final_jmp_loc is not None:
                        self.handle_found_opaque_predicate(OpaqueAnalyzeRetValues.ALWAYS_JMP, opaque_addr, block)
                    last_found_opaque_predicate = opaque_addr
                    return OpaqueIdaPatcherRetValues.RETURN_TO_MAINLOOP
                              
        return OpaqueIdaPatcherRetValues.PROCEED



    def handle_found_opaque_predicate(self, analyze_res, opaque_addr, block, different_addr = None):
        global is_last_predicate_mode
        if (analyze_res == OpaqueAnalyzeRetValues.ALWAYS_JMP) or (analyze_res == OpaqueAnalyzeRetValues.NEVER_JMP):
            #Opaque was found, restart the function analysis.
            #TODO: currently, angr is working on the binary, meaning if IDA patched the bytes it does not mean angr seen it.
            #think of something clever to solve that.
            print("######## - OPAQUE FOUND - ########")
            print("Opaque conditional jump address is: ", hex(opaque_addr))
            if analyze_res == OpaqueAnalyzeRetValues.ALWAYS_JMP:
                self.fix_jmp_inst_on_opaque_predicate(opaque_addr)
                if sum(1 for _ in (idautils.CodeRefsTo(idc.next_head(opaque_addr), 1))) == 1:
                    logging.debug("No xrefs, patching opaque func chunk end")
                    #if there are no xrefs to the next line, fix the function chunk end.
                    self.fix_function_end(block.end_ea)
                logging.debug("Opaque type is ALWAYS JMP")
            elif analyze_res == OpaqueAnalyzeRetValues.NEVER_JMP:
                logging.debug("Opaque type is NEVER JMP")
            else:
                    logging.error("PROBABLY SHOULDNT BE HERE \nError was returned, probably trying to go to a chunk that is not a part of fuction.\n")
                    raise "ERROR"
            print("Starting to patch the opaque")
            self.fix_opaque_predicate(opaque_addr, analyze_res, different_addr)
            print("Tomerminatored the opaque.")
            return True
        if analyze_res == OpaqueAnalyzeRetValues.CHANGE_MODE:
            logging.debug("Changing mode.. to many active shitty states.")
            is_last_predicate_mode = True
        if analyze_res == OpaqueAnalyzeRetValues.NOT_OPAQUE:
            if self.fix_jmp_inst_on_opaque_predicate(opaque_addr):
                return True

        return False


    def patch_opaque_predicates_loop(self, is_single_block_mode):
        global is_last_predicate_mode
        self.referesh_screen_and_reanalyze()
        fc = idaapi.FlowChart(self.current_func)
        counter = 0
        
        #MAIN LOOP
        for block in fc:
            print("======================================================================")
            
            self.referesh_screen_and_reanalyze()
            opaque_addr = idc.prev_head(block.end_ea)
            patch_retval = self.ida_pre_explore_patch(block, is_single_block_mode)
            
            if patch_retval == OpaqueIdaPatcherRetValues.SKIP_BLOCK:
                continue
            if patch_retval == OpaqueIdaPatcherRetValues.RETURN_TO_MAINLOOP:
                return True
            if patch_retval == OpaqueIdaPatcherRetValues.PROCEED:
                pass

            if (self.check_conditional_x86_64(opaque_addr)):
                print("Found conditional jump address in: ", hex(opaque_addr), ".\n Starting to explore.")
                analyze_res = self.explore_for_opaques(opaque_addr, is_single_block_mode, block.start_ea)
                if self.handle_found_opaque_predicate(analyze_res, opaque_addr, block):
                    return True
            else:
                print("conditional jump not found.")  

            print("Iteration Done.")
            counter += 1
            if counter == DEPTH:
                logging.debug("MAX ITERATIONS, STOPPING.")
                break
    
        return False



    def convert_asm_string_to_bytes(self, asm_str)->bytes:
        ks = Ks(KS_ARCH_X86, KS_MODE_64)
        all_shell_code = b""
        for inst in asm_str:
            encoding, count = ks.asm(inst)
            opcodes = b""
            for i in encoding:
                opcodes += bytes([i])
            all_shell_code += opcodes



    def pre_explore_get_ea_of_execution_start(self, is_single_block_mode, block_start_ea):
        global last_found_opaque_predicate
        global is_last_predicate_mode
        start_ea = 0x0
        if (last_found_opaque_predicate is not None) and is_last_predicate_mode:
            start_ea = self.get_the_location_of_the_jump(last_found_opaque_predicate)
            logging.debug("Start from last found predicate, not function start.")
        else:
            start_ea = self.current_func.start_ea

        if is_single_block_mode:
            start_ea = block_start_ea

        if (not is_last_predicate_mode) and (not is_single_block_mode):
            start_ea = self.current_func.start_ea
        
        return start_ea



    def apply_pre_explore_heursitcs(self, predicate_ea):
        global addresses_avoided_set
        global addresses_find_set
        predicate_jmp_address_operand = self.get_the_location_of_the_jump(predicate_ea)

        #Append invariants
        addresses_avoided_set.add(predicate_jmp_address_operand)
        addresses_find_set.add(idaapi.next_head(predicate_ea, predicate_ea+0xF))
        addresses_find_set.add(idc.next_addr(predicate_ea))

        logging.debug("PRE EXPLORE: Avoid set is: [{}]".format(', '.join(hex(x) for x in addresses_avoided_set)))
        logging.debug("PRE EXPLORE: Find set is: [{}]".format(', '.join(hex(x) for x in addresses_find_set)))
        


    def evaluate_explore_results(self, simMgr):
        logging.debug("SimMgr: {}".format(simMgr))
        total_stashes_len = 0

        for stash in simMgr.stashes:
            total_stashes_len += len(stash)

        if len(simMgr.unconstrained) == total_stashes_len:
            raise ValueError("All states ended unconstrained.")
            return OpaqueAnalyzeRetValues.NOT_OPAQUE
        
        if len(simMgr.found) == 0:
            logging.debug("len(simMgr.found) == 0 and len(simMgr.unconstrained) == 0")
            return OpaqueAnalyzeRetValues.ALWAYS_JMP
        else:
            if total_stashes_len - len(simMgr.found) == 0:
                logging.debug("len(simMgr.found) != 0 or len(simMgr.unconstrained) != 0, total_stashes_len = len(found)")
                return OpaqueAnalyzeRetValues.NEVER_JMP
            if total_stashes_len - len(simMgr.avoid) == 0:
                logging.debug("total_stashes_len = len(avoid)")
                return OpaqueAnalyzeRetValues.ALWAYS_JMP
        
        return OpaqueAnalyzeRetValues.NOT_OPAQUE



    def patch_bytes_on_position(self, ea, fpos, org_val, patch_val):
        if fpos != -1:
            self.inputfile_obj.seek(fpos)
            self.inputfile_obj.write(struct.pack('B', patch_val))
        return 0



    def apply_patch_bytes(self, start_ea, end_ea):
        self.inputfile_obj = open(self.inputfile_path, 'rb+')
        idaapi.visit_patched_bytes(start_ea, end_ea, self.patch_bytes_on_position)
        self.inputfile_obj.close()



    def apply_post_explore_heuristics(self, ret_code, predicate_ea, simMgr):     
        global last_found_opaque_predicate
        global addresses_avoided_set
        global addresses_find_set
        global opaque_jumpedto_addresses_set
        ret_code = ret_code
        predicate_jmp_address_operand = self.get_the_location_of_the_jump(predicate_ea)
        
        #restore invaraints.
        #TODO: think of heuristic.
        addresses_avoided_set.remove(predicate_jmp_address_operand)
        addresses_find_set.clear()
        if (ret_code != OpaqueAnalyzeRetValues.NOT_OPAQUE):
            last_found_opaque_predicate = predicate_ea
            opaque_jumpedto_addresses_set.add(predicate_jmp_address_operand)

        if len(simMgr.active) > ACTIVE_THERSHOLD_FOR_LAST_PREDICATE_MODE:
            ret_code = OpaqueAnalyzeRetValues.CHANGE_MODE
            
             

        logging.debug("POST EXPLORE: Avoid set is: [{}]".format(', '.join(hex(x) for x in addresses_avoided_set)))
        logging.debug("POST EXPLORE: Find set is: [{}]".format(', '.join(hex(x) for x in addresses_find_set)))
        return ret_code
    
    
    @staticmethod
    def avoid_function(state):
        return state.addr >= min(addresses_avoided_set)
   
   
    @staticmethod
    def find_function(state):
        return (state.addr >= min(addresses_find_set)) and (state.addr <= max(addresses_find_set))
    
    
    
    def explore_for_opaques(self, predicate_ea, is_single_block_mode, block_start_ea):
        start_ea = 0x0
        ret_code = OpaqueAnalyzeRetValues.NOT_OPAQUE

        start_ea = self.pre_explore_get_ea_of_execution_start(is_single_block_mode, block_start_ea)
        
        #ignore func call
        s = self.angr_prj.factory.blank_state(addr=start_ea)
        s.options.add(angr.options.CALLLESS)
        #s.options.add(angr.options.SYMBOL_FILL_UNCONSTRAINED_REGISTERS)
        #s.options.add(angr.options.SYMBOL_FILL_UNCONSTRAINED_MEMORY)
        #s.options.add(angr.options.SYMBOLIC)
        #s.options.add(angr.options.SYMBOLIC_INITIAL_VALUES)
        #s.options.add(angr.options.ABSTRACT_MEMORY)
        
        simMgr = self.angr_prj.factory.simgr(s)

        self.apply_pre_explore_heursitcs(predicate_ea)
        
        simMgr.explore(find=FunctionOpaqueIdentifier.find_function, avoid=FunctionOpaqueIdentifier.avoid_function)
        
        ret_code = self.evaluate_explore_results(simMgr)
    
        ret_code = self.apply_post_explore_heuristics(ret_code, predicate_ea, simMgr)
    
        return ret_code

   
   
    def patch_opaque_predicates(self):
        
        if(FIX_DEFLOW_MODE):
            logging.debug("Deflow fix running.")
            self.fix_using_deflow()
        elif(SINGLE_BLOCK_MODE):
            logging.debug("Single block level opaque removal fix running.")
            self.fix_using_single_block_analysis()
        elif(WHOLE_FUNCTION_MODE):
            logging.debug("Whole function level opaque removal fix running.")
            self.fix_using_whole_function_analysis()



    def run(self):
        self.patch_opaque_predicates()

        
        
jumps = [NN_jmp, NN_jmpfi, NN_jmpni, NN_jmpshort, NN_ja, NN_jae, NN_jb, NN_jbe, NN_jc, NN_jcxz, NN_jecxz, NN_jrcxz, NN_je,NN_jg,NN_jge,NN_jl,NN_jle,NN_jna,NN_jnae,NN_jnb,NN_jnbe,NN_jnc,NN_jne,NN_jng,NN_jnge,NN_jnl,NN_jnle,NN_jno,NN_jnp,NN_jns,NN_jnz,NN_jo,NN_jp,NN_jpe,NN_jpo,NN_js,NN_jz]


class FixCode:
    unexplored: set[int] = set()
    explored:  set[int] = set()

    def __init__(self, ea) -> None:
        self.unexplored.add(ea)

    @staticmethod
    def is_valid_jump(cmd):
        jump_to = cmd.Op1.addr
        item_start = get_item_head(jump_to)
        # print(f"{cmd.ea:x} - is_valid_jump: {jump_to:x} =? {get_item_head(jump_to):x}")
        return item_start == jump_to

    # makes code if needed on ea and returns next instr addr
    @staticmethod
    def make_code(ea) -> int:
        # skip if already code
        if is_code(get_flags(ea)) and get_item_head(ea) == ea:
            return get_item_head(ea)

        cmd = insn_t()
        auto_wait()
        if create_insn(ea, cmd) <= 0:
            # try to undef and retry
            del_items(ea, 0, 10)
            cmd = insn_t()
            if create_insn(ea, cmd) <= 0:
                print(f"create_insn(ea, cmd) failed {ea:x}, dont know what to do")
                return BADADDR
            auto_wait()

        return get_item_head(ea)

    @staticmethod
    def append_cmt(ea, cmt):
        e_cmt = get_cmt(ea, False) or ''
        set_cmt(ea, e_cmt + " " + cmt, 0)

    @staticmethod
    def fill_nop(cmd):
        # 2 bytes jump
        FixCode.append_cmt(cmd.ea, f"Patched jmp, original: {GetDisasm(cmd.ea)}")
        patch_byte(cmd.ea, 0xEB)
        for ea in range(get_item_end(cmd.ea), cmd.Op1.addr):
            patch_byte(ea, 0x90)
        # if cmd.itype in NN_call:

    # fix 'chunk' from ea and below, until ret or no code
    def fix_chunk(self, ea):
        start_chunk_ea = ea
        cmd = insn_t()
        print(f"fix_chunk starting ea: {ea:x}")
        FixCode.append_cmt(ea, f"chunk {ea:x} starts here")
        while True:
            ea = get_item_head(ea)
            new_ea = self.make_code(ea)
            if new_ea == BADADDR:
                break
            ea = new_ea
            size = decode_insn(cmd, ea)
            print(f"working on {ea:x}, size: {size} {GetDisasm(ea)}")
            if cmd.itype in jumps:

                # check if jump+X - not valid, obfuscation
                rc = self.is_valid_jump(cmd)
                if rc:
                    if cmd.Op1.addr not in self.explored:
                        FixCode.append_cmt(cmd.ea, "original good jump ")
                        self.explored.add(cmd.Op1.addr)
                        if cmd.itype == NN_jmp:
                            print(f"straight JMP to {cmd.Op1.addr:x} continue there")
                            FixCode.append_cmt(cmd.ea, "JMP, taking it to process")
                            ea = cmd.Op1.addr
                            continue

                        # seems to be valid conditional jump, add addr to process this branch
                        self.unexplored.add(cmd.Op1.addr)
                        print(f"valid jump at {cmd.ea:x}, adding {cmd.Op1.addr:x}")
                        FixCode.append_cmt(cmd.ea, f"good conditional jump, queue as chunk {cmd.Op1.addr:x}")
                    else:
                        print(f"already processed, valid jump at {cmd.ea:x}")
                else:
                    if size > 2:
                        print(f"WARNING, size {size} > 2, obfuscations jumps are 2 byte long")
                        return

                    print(f"obfuscation jump at {cmd.ea:x} to {cmd.Op1.addr:x}, nop it")
                    if cmd.ea > cmd.Op1.addr:
                        print(f"negative jump at {cmd.ea:x} to {cmd.Op1.addr:x}, stop chunk here, idk what to do")
                        return
                    # fill with nops
                    self.fill_nop(cmd)
                    # self.make_code(cmd.Op1.addr)
                    ea = cmd.Op1.addr
                    del_items(ea, 0, 10)
                    auto_wait()
                    continue

            if cmd.itype == NN_retn:
                print(f"reached return. done  {cmd.ea:x} ")
                break

            # advance to next instr
            ea = get_item_end(ea)
        print(f"fix_chunk ends: {start_chunk_ea:x}")
        FixCode.append_cmt(ea, f"chunk {start_chunk_ea:x} ends here")

    def process(self):
        loops = 0
        while len(self.unexplored) > 0:
            ea = self.unexplored.pop()
            self.fix_chunk(ea)
            loops += 1
            _addr = [f"{x:x} " for x in self.unexplored]
            _done = [f"{x:x} " for x in self.explored]
            print(f"addresses to process: {len(self.unexplored)} - {_addr}")
            print(f"addresses done: {len(self.explored)} - {_done}")
            # if loops > 50:
            #     break

        
        

def run_main_class():
    analyze = FunctionOpaqueIdentifier()
                
        


    
"""------------------------------"""
class Topaqueminator(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    wanted_name = "Topaqueminator"
    wanted_hotkey = "Ctrl-Alt-T"
    comment = 'Topaqueminator Comment'
    help = 'TBD'

    def init(self):
        print('--------------------------------------')
        print('Topaqueminator Starting. VER:' , T_VER)
        print('--------------------------------------')
        return idaapi.PLUGIN_OK
    
    def run(self, ctx):
        run_main_class()
       
        

    def term(self):
        print('Made by Tomerminator, use as you wish.')
        print('Topaqueminator finished.')
"""------------------------------"""
def PLUGIN_ENTRY():
    try:
        return Topaqueminator()
    except Exception as err:
        import traceback
        print('Error: %s\n%s' % str((err), traceback.format_exc()))
        raise
"""------------------------------"""


repository: LAC-Japan/IDA_Plugin_AntiDebugSeeker
file: sources/python/ver1.3/AntiDebugSeeker.py
description: Automatically identify and extract potential anti-debugging techniques used by malware.
wanted_name: AntiDebugSeeker
comment: Anti Debug Detection Plugin
created_at: 2023-09-11 02:29:26+00:00
pushed_at: 2024-11-22 03:45:20+00:00
url:
https://github.com/LAC-Japan/IDA_Plugin_AntiDebugSeeker/blob/070e26d115734a988c53a39bd8ee7bcec7e9a957/sources/python/ver1.3/AntiDebugSeeker.py

content:
import idaapi
import ida_ida
import ida_bytes
import idc
import idautils
import ida_kernwin
from ida_kernwin import PluginForm
import ida_lines
import os
import re
import json
import sys
from collections import defaultdict
from PyQt5 import QtWidgets, QtGui, QtCore
from PyQt5.QtGui import QColor
from PyQt5.QtWidgets import QApplication
import traceback

class AntiDebugResults(ida_kernwin.Choose):
    def __init__(self, title, items):
        ida_kernwin.Choose.__init__(self, title, [["Category Name", 20 | ida_kernwin.Choose.CHCOL_PLAIN],
                                                 ["Possible Anti-Debug API", 20 | ida_kernwin.Choose.CHCOL_PLAIN],
                                                 ["Address", 15 | ida_kernwin.Choose.CHCOL_HEX],
                                                 ["Possible Anti-Debug Technique", 25 | ida_kernwin.Choose.CHCOL_PLAIN],
                                                 ["Address", 15 | ida_kernwin.Choose.CHCOL_HEX]])
        self.items = items

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

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

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

    def OnGetLineAttr(self, n):
        black = 0x000000
        gray = 0xC0C0C0
        if self.items[n][0] == "" and self.items[n][1] == "":
            return [ida_kernwin.Choose.CHCOL_PLAIN | black, ida_kernwin.Choose.CHCOL_PLAIN | black]
        elif n % 2 == 0:
            return [ida_kernwin.Choose.CHCOL_PLAIN | gray, ida_kernwin.Choose.CHCOL_PLAIN | gray]
        else:
            return [ida_kernwin.Choose.CHCOL_PLAIN | gray, ida_kernwin.Choose.CHCOL_PLAIN | gray]

    def OnSelectLine(self, n):
        col2_addr = self.items[n][2]
        col4_addr = self.items[n][4]

        if col2_addr:
            idaapi.jumpto(int(col2_addr, 16))
        if col4_addr:
            idaapi.jumpto(int(col4_addr, 16))

    def OnDblClick(self, n):
        widget = ida_kernwin.find_widget("Anti Debug Detection Results")
        if not widget:
            return

        sel = self.GetEmbSelection()
        if not sel:
            return

        col = sel[0].y
        addr = self.items[n][col]

        if addr:
            idaapi.jumpto(int(addr, 16))
        return 1

    def show(self):
        self.Show()
        
class MyEmbeddedForm(PluginForm):
           def __init__(self, anti_debug_techniques_descriptions, function_display_results):
              super(MyEmbeddedForm, self).__init__()
              self.anti_debug_techniques_descriptions = anti_debug_techniques_descriptions
              self.function_display_results = function_display_results

           def OnCreate(self, form):

              self.qt_form = self.FormToPyQtWidget(form)
              
              self.viewer = DetectedFunctionListViewer(self.anti_debug_techniques_descriptions, self.function_display_results)
              self.viewer.setGeometry(100, 100, 800, 600)
              
              layout = QtWidgets.QVBoxLayout()
              layout.addWidget(self.viewer)
              self.qt_form.setLayout(layout)

           def Show(self):
              return PluginForm.Show(self, "Detected Function List")
        
class DetectedFunctionListViewer(QtWidgets.QWidget):
    highlighted_functions = None
    detected_all_functions = None

    def __init__(self, descriptions, results, parent=None):
        try:
            super(DetectedFunctionListViewer, self).__init__()
            self.descriptions = descriptions
            self.results = results
            
            if DetectedFunctionListViewer.highlighted_functions is None:
                DetectedFunctionListViewer.highlighted_functions = self.get_highlighted_functions()
                DetectedFunctionListViewer.detected_all_functions_dict = self.list_to_dict(self.results)
            
            self.init_ui()
            self.show()
        except Exception as e:
            import traceback
            print("Error in DetectedFunctionListViewer:", e)
            print(traceback.format_exc())

    def init_ui(self):
        self.layout = QtWidgets.QVBoxLayout(self)
        
        self.search_bar = QtWidgets.QLineEdit(self)
        self.search_bar.setPlaceholderText("Search...")
        self.search_bar.textChanged.connect(self.search_highlight)
        self.layout.addWidget(self.search_bar)
        
        self.list_widget = QtWidgets.QListWidget(self)
        
        font = self.list_widget.font()
        font.setPointSize(10)
        self.list_widget.setFont(font)
        
        for line in self.results:
            item = QtWidgets.QListWidgetItem(line)
            tooltip_text = []
            
            address_match = re.search(r"0x[0-9A-Fa-f]+", line)
            if address_match:
               item.setForeground(QtGui.QColor(0xFFD700))
            
            for key in self.descriptions.keys():
                if key in line:
                    item.setForeground(QtGui.QColor(0xFFA07A))
                    tooltip_text.append(self.descriptions[key])
            if tooltip_text:
               item.setToolTip('\n'.join(tooltip_text))
            self.list_widget.addItem(item)

        self.list_widget.itemDoubleClicked.connect(self.on_item_double_clicked)
        
        self.help_label = QtWidgets.QLabel("Double-click on a function name starting with 'sub' to investigate it recursively call.", self)
        self.help_label.setAlignment(QtCore.Qt.AlignCenter)
        self.help_label.setStyleSheet("color: gray;")

        self.layout.addWidget(self.help_label)
        self.layout.addWidget(self.list_widget)
        self.setLayout(self.layout)
        self.resize(400, 300)
        
    def list_to_dict(self, func_list):
        function_dict = {}
        current_key = None
        current_values = []
        first_key_set = False

        for item in func_list:
            if item == '':
                if current_key and current_values:
                    function_dict[current_key] = current_values
                current_key = None
                current_values = []
            elif current_key is None:
                current_key = item
                if not first_key_set:
                    first_key_set = True
            elif not any(marker in item for marker in ['(', ')', 'detected']):
                current_values.append(item)
        
        if current_key and current_values:
            function_dict[current_key] = current_values

        return function_dict
        
    def get_highlighted_functions(self):
        highlighted = []
        for line in self.results:
            match = re.search(r'\bsub_[0-9A-Fa-f]+\b', line)
            if match:
                highlighted.append(match.group(0))
        return highlighted
        
    def search_highlight(self):
            
        search_text = self.search_bar.text().lower()

        for i in range(self.list_widget.count()):
            item = self.list_widget.item(i)
            item_text = item.text()
             
            is_address = bool(re.search(r"0x[0-9A-Fa-f]+", item_text))
            is_description_key = any(key in item_text for key in self.descriptions.keys())

            item.setBackground(QColor(0xFFFFFF))
            item.setForeground(QColor(0x000000))

            if search_text and search_text in item_text.lower():
                item.setBackground(QColor(0xAED6F1)) 

            if is_address:
                item.setForeground(QColor(0xFFD700))
            elif is_description_key:
                item.setForeground(QColor(0xFFA07A))

    def on_item_double_clicked(self, item):
        clicked_line = item.text()
        clicked_func_name = self.get_function_name_from_line(clicked_line)
                
        if clicked_func_name:
           max_depth = self.get_max_depth(clicked_func_name)
           all_results = self.recursive_xrefs(clicked_func_name, max_depth)
           self.result_viewer = RecursiveViewer({}, all_results, clicked_func_name)
        
        match = re.search(r"\((0x[0-9A-Fa-f]+)\)", clicked_line)
        if match:
            address = int(match.group(1), 16)
            idaapi.jumpto(address)
            
    def get_function_name_from_line(self, line):
        match = re.search(r'\bsub_[0-9A-Fa-f]+\b', line)
        if match:
           func_name = match.group(0)
           print("Checking the recursive calls :",func_name)
           return func_name
           
        return None
        
    def get_max_depth(self, func_name, depth=0, visited=None):
        if visited is None:
            visited = set()

        max_depth_current = depth

        if func_name and func_name not in visited:
            visited.add(func_name)
            func_addr = idc.get_name_ea_simple(func_name)

            for xref in idautils.XrefsTo(func_addr):
                xref_func_name = idc.get_func_name(xref.frm)
                child_max_depth = self.get_max_depth(xref_func_name, depth+1, visited)
                max_depth_current = max(max_depth_current, child_max_depth)

        return max_depth_current
    
    def recursive_xrefs(self, func_name, max_depth, current_depth=0, visited=None):
        if visited is None:
            visited = set()

        results = []

        if func_name and func_name not in visited:
            visited.add(func_name)
            func_addr = idc.get_name_ea_simple(func_name)
            
            depth_display = max_depth - current_depth
            
            prefix = "|-- " if current_depth > 0 else ""
            indent = "  " * depth_display + prefix

            for xref in idautils.XrefsTo(func_addr):
                xref_func_name = idc.get_func_name(xref.frm)
                
                caller_addr = xref.frm
                formatted_output = f"{indent}{func_name}  called_addr ({caller_addr:08X})  (depth:{depth_display})"
                results.append(formatted_output)
                
                child_results = self.recursive_xrefs(xref_func_name, max_depth, current_depth+1, visited)
                results.extend(child_results)

        return results
        
class RecursiveViewer(DetectedFunctionListViewer):

    def __init__(self, descriptions, results, clicked_func_name, parent=None):
        super(RecursiveViewer, self).__init__(descriptions, results, parent)
        self.setWindowFlags(self.windowFlags() | QtCore.Qt.WindowStaysOnTopHint)
        
        self.setWindowTitle(f"Check the recursive calls : {clicked_func_name}")
        self.show()
        
    def search_highlight(self):
        pass

    def init_ui(self):
        self.layout = QtWidgets.QVBoxLayout(self)
        
        help_label = QtWidgets.QLabel("Items in gray indicate functions that match the Detected Function List.", self)
        help_label.setAlignment(QtCore.Qt.AlignCenter)
        help_label.setStyleSheet("color: gray;")
        
        self.layout.addWidget(help_label)
        
        self.list_widget = QtWidgets.QListWidget(self)
        
        font = self.list_widget.font()
        font.setPointSize(10)
        self.list_widget.setFont(font)
        
        for line in self.results:
            item = QtWidgets.QListWidgetItem(line)
            tooltip_text = []
            
            function_match = re.search(r'\bsub_[0-9A-Fa-f]+\b', line)
            if function_match:
                if function_match.group(0) in self.highlighted_functions:
                    item.setForeground(QtGui.QColor(0x808080))
                    function_list = self.detected_all_functions_dict.get(function_match.group(0), [])
                    if function_list:
                        tooltip_text.append(', '.join(function_list))
                
            address_match = re.search(r"0x[0-9A-Fa-f]+", line)
            if address_match:
               item.setForeground(QtGui.QColor(0xFFD700))
            
            for key in self.descriptions.keys():
                if key in line:
                    item.setForeground(QtGui.QColor(0xFFA07A))
                    tooltip_text.append(self.descriptions[key])
            if tooltip_text:
               item.setToolTip('\n'.join(tooltip_text))
            self.list_widget.addItem(item)

        self.list_widget.itemDoubleClicked.connect(self.on_item_double_clicked)

        self.layout.addWidget(self.list_widget)
        self.setLayout(self.layout)
        self.resize(400, 300)
        
    def on_item_double_clicked(self, item):
        clicked_line = item.text()
        
        match = re.search(r"\(([0-9A-Fa-f]+)\)", clicked_line)
        if match:
            address = int(match.group(1), 16)
            idaapi.jumpto(address)

class ConfigEditorForm(ida_kernwin.Form):
    def __init__(self, file_path):
        self.file_path = file_path
        try:
            with open(file_path, 'r') as f:
                self.file_contents = f.read()
        except Exception as e:
            print("Error reading file: %s" % str(e))
            self.file_contents = ""

        F = ida_kernwin.Form
        F.__init__(self, r"""STARTITEM {id:txtInput}
BUTTON YES* Save
BUTTON CANCEL Cancel
Edit configuration
{FormChangeCb}
<##File contents:{txtInput}>
        """, {
            'txtInput': F.MultiLineTextControl(text=self.file_contents),
            'FormChangeCb': F.FormChangeCb(self.OnFormChange),
        })

    def OnFormChange(self, fid):
        return 1

    def OnButtonOk(self, code=0):
        try:
            with open(self.file_path, 'w') as f:
                f.write(self.GetControlValue(self.txtInput).text)
            return 1
        except Exception as e:
            print("Error writing to file: %s" % str(e))
            return 0

    def OnFormChange(self, fid):
        if fid == -2:
            self.OnButtonOk()
        return 1

class AntiDebugPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = "Anti Debug Detection Plugin"
    help = "Detects the use of Anti Debug Technique"
    wanted_name = "AntiDebugSeeker"
    wanted_hotkey = "Ctrl-Shift-D"

    def init(self):
        idaapi.msg("AntiDebugSeeker initialized. Ctrl-Shift-D Start Analysis.\n")
        self.edit_config_hotkey_ctx = ida_kernwin.add_hotkey("Ctrl-Shift-E", self.edit_config)
        self.saved_results = []
        if self.edit_config_hotkey_ctx is None:
            idaapi.msg("Failed to register hotkey for config editor. It might be already in use.\n")
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        if arg == "edit":
            self.edit_config()
        elif arg == "details":
            self.show_anti_debug_details()
        else:
            idaapi.msg("AntiDebugSeeker is running...\n")
            self.detect_anti_debug_functions()

    def save_anti_debug_results(self, result_win):
        self.saved_results = result_win

    def collect_anti_debug_results(self):
        return self.saved_results

    def term(self):
        idaapi.msg("AntiDebugSeeker terminated.\n")
        idaapi.msg("Edit anti_debug.config : Switch Other tab and Press Ctrl+Shift+E.\n")

    def load_anti_debug_descriptions(self, file_name):
        
        plugin_dir = os.path.dirname(os.path.realpath(__file__))
        file_path = os.path.join(plugin_dir, file_name)
        with open(file_path, 'r') as f:
            return json.load(f)

    def edit_config(self):
        plugin_dir = os.path.dirname(os.path.realpath(__file__))
        file_path = os.path.join(plugin_dir, "anti_debug.config")
        editor = ConfigEditorForm(file_path)
        editor.Compile()
        editor.Execute()

    def load_config(self, file_name):
        config_data = {
            "Anti_Debug_API": {},
            "Anti_Debug_Technique": [],
            "default_search_range": None
        }

        plugin_dir = os.path.dirname(os.path.realpath(__file__))
        file_path = os.path.join(plugin_dir, file_name)

        current_section = None
        current_category = None

        with open(file_path, "r") as f:
            for line in f:
                line = line.strip()

                if line.startswith("###") and line.endswith("###"):
                    current_section = line[3:-3]
                    current_category = None
                elif line.startswith("[") and line.endswith("]"):
                    current_category = line[1:-1]
                    if current_section == "Anti_Debug_API":
                        config_data[current_section][current_category] = []
                    elif current_section == "Anti_Debug_Technique":
                        config_data["Anti_Debug_Technique"].append({
                            "name": current_category,
                            "search_keyword": None,
                            "nearby_keyword": None,
                            "nearby_keyword2": None,
                            "search_range": config_data["default_search_range"]
                        })
                elif current_category and line and current_section in ["Anti_Debug_API", "Anti_Debug_Technique"]:
                    if current_section == "Anti_Debug_API":
                        config_data[current_section][current_category].append(line)
                    elif current_section == "Anti_Debug_Technique":
                        entry = config_data["Anti_Debug_Technique"][-1]
                        if line.startswith("search_range="):
                            entry["search_range"] = int(line.split("=")[1])
                        else:
                            if entry["search_keyword"] is None:
                                entry["search_keyword"] = line
                            elif entry["nearby_keyword"] is None:
                                entry["nearby_keyword"] = line
                            elif entry["nearby_keyword2"] is None:
                                entry["nearby_keyword2"] = line
                elif line.startswith("default_search_range=") and current_section == "Anti_Debug_Technique":
                    config_data["default_search_range"] = int(line.split("=")[1])

        return config_data

    def search_asm(self, search_keyword, nearby_keyword=None, nearby_keyword2=None, search_range=None):
        api_names = set(name for _, name in idautils.Names())
        one_keyword_match = []
        matching_sets = []

        def search_keyword_in_code(keyword, ea, disasm_line, end_addr):
            found = False
            keyword_addr = None
            
            keyword_is_int = False
            try:
                keyword_as_int = int(keyword)
                keyword_is_int = True
            except ValueError:
                pass
            
            regex_matched = re.search(r'\b' + re.escape(str(keyword)) + r'(h)?\b', disasm_line)
            comment = ida_bytes.get_cmt(ea, 0)
            
            def check_found(keyword_is_int, regex_matched):
                return (not keyword_is_int) or (keyword_is_int and regex_matched)

            if keyword in api_names:
                keyword_addr = idc.get_name_ea_simple(keyword)

                if keyword_addr != idaapi.BADADDR and ea in list(idautils.CodeRefsTo(keyword_addr, 0)):
                    found = True
                    keyword_addr = ea

            if comment and keyword in comment:
               found = check_found(keyword_is_int, regex_matched)
                
            if keyword in disasm_line:
               found = check_found(keyword_is_int, regex_matched)

            if not found:
                ref_addr = idc.get_operand_value(ea, 0)
                if ref_addr != idaapi.BADADDR:
                    strtype = idaapi.get_str_type(ref_addr)
                    if strtype == 0xFFFFFFFF:
                        return False, None
                    ref_string = ida_bytes.get_strlit_contents(ref_addr, ida_bytes.get_item_size(ref_addr), int(strtype))
                    if ref_string and keyword in ref_string.decode('utf-8', errors='ignore'):
                        found = check_found(keyword_is_int, regex_matched)
                        
            if found:
               if keyword_addr != ea:
                  keyword_addr = ea
                        
            return found, keyword_addr
            
        entry_point = idc.get_inf_attr(idc.INF_START_IP)  
        entry_section = idaapi.getseg(entry_point)
        
        for seg_ea in idautils.Segments():
            if seg_ea != entry_section.start_ea:  
                continue
               
            ea = seg_ea
            end_addr = idc.get_segm_end(seg_ea)

            while ea < end_addr:
                disasm_line = idc.GetDisasm(ea)
                search_found, search_addr = search_keyword_in_code(search_keyword, ea, disasm_line, end_addr)

                if search_found:
                    base_addr = search_addr if search_keyword in api_names else ea
                    if not nearby_keyword and not nearby_keyword2:
                        one_keyword_match.append(hex(base_addr))
                    else:
                        next_ea = base_addr
                        while next_ea - base_addr < search_range and next_ea < end_addr:
                            disasm_next_line = idc.GetDisasm(next_ea)
                            found_nearby, nearby_addr = search_keyword_in_code(nearby_keyword, next_ea, disasm_next_line, end_addr)
                            if found_nearby:
                                if not nearby_keyword2:
                                    matching_sets.append((hex(base_addr), hex(next_ea), None))
                                else:
                                    next_ea2 = next_ea
                                    while next_ea2 - next_ea < search_range and next_ea2 < end_addr:
                                        disasm_next_line2 = idc.GetDisasm(next_ea2)
                                        found_nearby2, nearby_addr2 = search_keyword_in_code(nearby_keyword2, next_ea2, disasm_next_line2, end_addr)
                                        if found_nearby2:
                                            matching_sets.append((hex(base_addr), hex(next_ea), hex(next_ea2)))
                                        next_ea2 = idc.next_head(next_ea2)
                            next_ea = idc.next_head(next_ea)
                ea = idc.next_head(ea)

        if one_keyword_match:
           print(f"[Result] Found {len(one_keyword_match)} match(es) for First_keyword '{search_keyword}' under the rule of specifying one keyword.")
           return one_keyword_match
            
        if matching_sets:
           print(f"[Result] Found {len(matching_sets)} matching set(s).")
           return matching_sets
        else:
            return None, None, None

    def detect_anti_debug_functions(self):
        results = []
        color = 0x98FB98
        color2 = 0xAAD4FF
        detected_function_names = []

        def set_item_color(address, color):
            try:
                idaapi.set_item_color(int(address, 16), color)
            except TypeError as e:
                print(f"Error: {e}")

        def add_comment(address, comment_name):
            existing_comment = idaapi.get_cmt(address, 0)
            comment_description = anti_debug_techniques_descriptions.get(comment_name, "")
            comment_text = comment_name if existing_comment is None else existing_comment + " | " + comment_name
            if comment_description:
                comment_text += " - " + comment_description
            idaapi.set_cmt(address, comment_text, 0)

        config_data = self.load_config("anti_debug.config")
        anti_debug_functions = config_data["Anti_Debug_API"]
        search_patterns = config_data["Anti_Debug_Technique"]
        default_search_range = config_data["default_search_range"]
        anti_debug_techniques_descriptions = self.load_anti_debug_descriptions("anti_debug_techniques_descriptions.json")
        
        def is_address_in_target_sections(address, func, entry_section):
           if entry_section is not None:
               section_start = entry_section.start_ea
               section_end = entry_section.end_ea
               if section_start <= address <= section_end:
                   return True

           if func is None:
               return False

           func_start = func.start_ea
           func_end = func.end_ea

           code_section = idaapi.getseg(func_start)
           if code_section is not None and code_section.perm & idaapi.SEGPERM_EXEC:
               return True

           return False
            
        entry_point = idc.get_inf_attr(idc.INF_START_IP)
        entry_section = idaapi.getseg(entry_point)

        for category, functions in anti_debug_functions.items():
            for func_name in functions:
                func_addr = idc.get_name_ea_simple(func_name)

                if func_addr == idaapi.BADADDR:
                    continue

                xrefs = list(idautils.XrefsTo(func_addr))
                
                if xrefs:
                    idaapi.msg(f"{func_name} function found.\n")
                else:
                    continue

                processed_addresses = set()
                for xref in xrefs:
                    address = xref.frm
                    address_str = f"0x{address:X}"
                    containing_func = idaapi.get_func(xref.frm)
                    
                    if address_str in processed_addresses:
                       continue
                    
                    if is_address_in_target_sections(address, containing_func, entry_section):
                        function_name = idaapi.get_func_name(address)
                        detected_function_names.append(f"Function containing {func_name}: {function_name}")
                        idaapi.msg(f"  {address_str}\n")
                        set_item_color(address_str, color)
                        add_comment(xref.frm, category)
                        results.append([category, func_name, address_str, "", ""])
                        processed_addresses.add(address_str)

        start_addr = ida_ida.inf_get_max_ea()
        end_addr = ida_ida.inf_get_min_ea()

        for pattern in search_patterns:
            search_keyword = pattern["search_keyword"]
            nearby_keyword = pattern.get("nearby_keyword", None)
            nearby_keyword2 = pattern.get("nearby_keyword2", None)
            search_range = pattern.get("search_range", default_search_range)

            asm_address = nearby_address = nearby2_address = None

            if nearby_keyword2:
                result = self.search_asm(search_keyword, nearby_keyword, nearby_keyword2, search_range=search_range)
            elif nearby_keyword:
                result = self.search_asm(search_keyword, nearby_keyword, search_range=search_range)
            else:
                result = self.search_asm(search_keyword)
                
            if result:
                if isinstance(result, list):
                    print(f"Found for pattern {pattern['name']}.")
                    for element in result: 
                        if element is not None:
                           if isinstance(element, tuple):
                               asm_address, nearby_address, nearby2_address = element + (None,) * (3 - len(element))
                               for address in [asm_address, nearby_address, nearby2_address]:
                                   if address is not None:
                                       address_int = int(address, 16)
                                       set_item_color(address, color2) 
                                       
                               if asm_address is not None:
                                   function_name = idaapi.get_func_name(int(asm_address, 16))
                                   detected_function_names.append(f"Function containing {pattern['name']}: {function_name}")
                                   address_int = int(asm_address, 16)
                                   add_comment(address_int, pattern["name"])
                                   results.append(["", "", "", pattern["name"], str(asm_address), ""])

                           elif isinstance(element, str):
                              function_name = idaapi.get_func_name(int(element, 16))
                              detected_function_names.append(f"Function containing {pattern['name']}: {function_name}")
                              address_int = int(element, 16)
                              set_item_color(element, color2)
                              add_comment(address_int, pattern["name"])
                              results.append(["", "", "", pattern["name"], str(element), ""])
                              
                elif isinstance(result, tuple):
                    if result == (None, None, None):
                       print(f"Nothing Found for pattern {pattern['name']}.")
        
        function_map = defaultdict(list)
        function_address_map = {}

        for fn_name in detected_function_names:
           parts = fn_name.split(":")
           feature_name = parts[0].split()[-1]
           function_name = parts[1].strip()
           function_map[function_name].append(feature_name)
           
           function_address = idc.get_name_ea_simple(function_name)
           if function_address != idaapi.BADADDR:
              function_address_map[function_name] = f"0x{function_address:X}"
           else:
              function_address_map[function_name] = "Unknown Address"
              
        sorted_function_names = sorted(function_map.keys(), key=lambda x: int(function_address_map[x], 16) if function_address_map[x] != "Unknown Address" else float("inf"))

        function_display_results = []
        for function_name in sorted_function_names:
           features = function_map[function_name]
           address = function_address_map[function_name]
           
           display_line = f"{function_name}"
           function_display_results.append(display_line)
           
           address_line = f"({address})"
           function_display_results.append(address_line)
           
           for feature in features:
              function_display_results.append(f"    {feature}")
              
           detected_line = f"({len(features)}detected)"
           function_display_results.append(detected_line)
           function_display_results.append("")
           
        f = MyEmbeddedForm(anti_debug_techniques_descriptions, function_display_results)
        f.Show()
        
        sorted_results = sorted(results, key=lambda x: (int(x[2], 16) if x[2] else int(x[4], 16) if x[4] else float("inf")))
        self.save_anti_debug_results(sorted_results)

        result_win = AntiDebugResults("Anti Debug Detection Results", sorted_results)
        result_win.show()    
        
def PLUGIN_ENTRY():
    return AntiDebugPlugin()


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

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

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

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

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

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

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

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

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

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

def PLUGIN_ENTRY():
    return my_plugin_t()


repository: 19h/ida-plugins
file: frIDA/frIDA.py
url: https://github.com/19h/ida-plugins/blob/b045b6f6bc9ae2a53104986fb74f3fe1198e5d2e/frIDA/frIDA.py
comment: Runtime Pointer Resolution using Frida
created_at: 2025-04-08 22:54:27+00:00
pushed_at: 2025-04-08 23:07:17+00:00
content:
"""
frIDA: Runtime Pointer Resolution for IDA Pro
An isolated plugin for dynamic resolution of static pointers using Frida instrumentation.
Version: 1.1.1
"""

import os
import traceback
import time

# IDA API imports
import ida_idaapi
import ida_kernwin
import ida_bytes
import ida_name
import ida_nalt
import ida_segment
import ida_idp
import ida_typeinf
import ida_ua
import ida_xref

# Conditional Frida import
try:
    import frida
    HAS_FRIDA = True
except ImportError:
    HAS_FRIDA = False

#-------------------------------------------------------------------------------
# Plugin Constants and Configuration
#-------------------------------------------------------------------------------
PLUGIN_NAME = "frIDA"
PLUGIN_VERSION = "1.1.1" # Incremented version for fixes
PLUGIN_HOTKEY = "Ctrl+Alt+F"
DEBUG_MODE = True  # Set to True for verbose logging

#-------------------------------------------------------------------------------
# Logging and Utilities
#-------------------------------------------------------------------------------
def log_debug(msg):
    if DEBUG_MODE:
        print(f"[DEBUG] {PLUGIN_NAME}: {msg}")

def log_info(msg):
    print(f"[INFO] {PLUGIN_NAME}: {msg}")

def log_error(msg, exc=None):
    print(f"[ERROR] {PLUGIN_NAME}: {msg}")
    if exc:
        print(f"[ERROR] Exception: {exc}")
        if DEBUG_MODE:
            traceback.print_exc()

def get_imagebase():
    """Retrieve the rebased imagebase of the current IDB."""
    return ida_nalt.get_imagebase()

def get_file_name():
    """Extract the binary name from the current IDB."""
    root_filename = ida_nalt.get_root_filename()
    if root_filename:
        return root_filename
    input_path = ida_nalt.get_input_file_path()
    return os.path.basename(input_path) if input_path else "unknown_binary"

def get_bitness():
    """Determine whether the target is 32 or 64-bit."""
    return 64 if ida_idp.ph_get_flag() & ida_idp.PR_USE64 else 32

def is_in_executable_segment(ea):
    """Check if an address resides in an executable segment."""
    seg = ida_segment.getseg(ea)
    return seg and (seg.perm & ida_segment.SEGPERM_EXEC)

def format_address(addr, bitness=None):
    """Format an address according to architecture bitness."""
    if bitness is None:
        bitness = get_bitness()
    # Ensure addr is an integer before formatting
    try:
        addr_int = int(addr)
        return f"0x{addr_int:0{bitness//4}X}"
    except (ValueError, TypeError):
        return str(addr) # Return original string if conversion fails


#-------------------------------------------------------------------------------
# Memory Analysis Utilities
#-------------------------------------------------------------------------------
def analyze_pointer_type(ea):
    """
    Multi-layer memory type inference with context-specific heuristics.

    Args:
        ea: Effective address to analyze

    Returns:
        dict: Type metadata and structural characteristics
    """
    flags = ida_bytes.get_flags(ea)
    segment = ida_segment.getseg(ea)
    segment_name = ida_segment.get_segm_name(segment) if segment else "UNKNOWN"
    segment_class = ida_segment.get_segm_class(segment) if segment else "UNKNOWN"

    # Initialize result structure
    info = {
        'address': ea,
        'is_pointer': False,
        'size': 0,
        'type_name': 'unknown',
        'is_code_ref': False,
        'is_data_ref': False,
        'xrefs': [],
        'segment': segment_name,
        'segment_class': segment_class
    }

    # Type inference phase 1: Symbol name parsing
    name = ida_name.get_name(ea)
    if name:
        info['name'] = name
        # Parse IDA's type encoding in symbol names
        if name.startswith(('dword_', 'qword_', 'xmmword_', 'ptr_')):
            if name.startswith('dword_'):
                info['is_pointer'] = True
                info['size'] = 4
                info['type_name'] = 'dword'
            elif name.startswith('qword_'):
                info['is_pointer'] = True
                info['size'] = 8
                info['type_name'] = 'qword'
            elif name.startswith('xmmword_'):
                info['size'] = 16
                info['type_name'] = 'xmmword'
            elif name.startswith('ptr_'):
                info['is_pointer'] = True
                info['size'] = 8 if get_bitness() == 64 else 4
                info['type_name'] = 'pointer'
    else:
        # Generate synthetic name for unnamed locations
        info['name'] = f"loc_{ea:X}" if not ida_bytes.is_code(flags) else f"sub_{ea:X}"

    # Type inference phase 2: Data type flags
    if info['size'] == 0:
        if ida_bytes.is_qword(flags):
            info['is_pointer'] = True
            info['size'] = 8
            info['type_name'] = 'qword'
        elif ida_bytes.is_dword(flags):
            info['is_pointer'] = True
            info['size'] = 4
            info['type_name'] = 'dword'
        elif ida_bytes.is_word(flags):
            info['size'] = 2
            info['type_name'] = 'word'
        elif ida_bytes.is_byte(flags):
            info['size'] = 1
            info['type_name'] = 'byte'
        elif ida_bytes.is_oword(flags):
            info['size'] = 16
            info['type_name'] = 'oword'

    # Type inference phase 3: Segment-based heuristics
    if segment_name in ('.rdata', '.data', 'DATA', 'data', '.bss', 'BSS', 'CONST', '.rodata'):
        arch_ptr_size = 8 if get_bitness() == 64 else 4
        if info['size'] == arch_ptr_size and (ea % arch_ptr_size) == 0:
            info['is_pointer'] = True
            log_debug(f"Segment-based pointer inference: {info['name']} in {segment_name}")

    # Type inference phase 4: Cross-reference analysis
    xref = ida_xref.get_first_dref_to(ea)
    while xref != ida_idaapi.BADADDR:
        info['xrefs'].append(xref)
        xref_flags = ida_bytes.get_flags(xref)
        if ida_bytes.is_code(xref_flags):
            info['is_code_ref'] = True
        else:
            info['is_data_ref'] = True
        xref = ida_xref.get_next_dref_to(ea, xref)

    # Type inference phase 5: Advanced type information (if available)
    tif = ida_typeinf.tinfo_t()
    if ida_nalt.get_tinfo(tif, ea):
        if tif.is_ptr():
            info['is_pointer'] = True
            info['type_name'] = 'pointer'
            pointed_tif = tif.get_pointed_object()
            if pointed_tif:
                info['pointed_type'] = pointed_tif.dstr() or "unknown"
            if info['size'] == 0:
                info['size'] = 8 if get_bitness() == 64 else 4
        elif tif.get_size() > 0:
             if info['size'] == 0:
                 info['size'] = tif.get_size()
                 info['type_name'] = tif.dstr() or info['type_name']

    # Last resort: use item size if all else fails
    if info['size'] == 0:
        item_size = ida_bytes.get_item_size(ea)
        if item_size > 0:
            info['size'] = item_size
        else:
            info['size'] = 8 if get_bitness() == 64 else 4
            log_debug(f"Falling back to arch pointer size ({info['size']}) for {info['name']} at {hex(ea)}")

    # Final check: if size matches pointer size, mark as potential pointer
    if info['size'] == (8 if get_bitness() == 64 else 4):
        info['is_pointer'] = True

    log_debug(f"Memory type analysis for {info['name']} at {hex(ea)}: {info}")
    return info


#-------------------------------------------------------------------------------
# Frida Communication Core
#-------------------------------------------------------------------------------
class FridaMemoryResolver:
    """
    Memory introspection subsystem for dynamic binary analysis using Frida instrumentation.
    Handles connection, script execution, and result parsing for both pointer
    resolution and direct data reading.
    """
    def __init__(self):
        """Initialize the memory introspection subsystem."""
        self.device = None
        self.session = None
        self.script = None
        self.result = None
        self.error = None
        self.imagebase = get_imagebase() # Static base from IDA
        self.bitness = get_bitness()
        self.architecture = None
        self.target_module_name = get_file_name() # Store target module name

    def connect(self, connection_type="remote", host="127.0.0.1:27042"):
        """
        Establish communication channel with Frida instrumentation framework.

        Args:
            connection_type: Channel type ("local", "usb", "remote")
            host: Endpoint for remote connections (host:port)

        Returns:
            bool: True if connection succeeded, False otherwise
        """
        try:
            if not HAS_FRIDA:
                raise ImportError("Frida instrumentation dependency unresolved")

            self.architecture = {
                'pointer_size': 8 if self.bitness == 64 else 4,
                'endianness': 'little', # Assuming little-endian for common targets
                'register_width': 64 if self.bitness == 64 else 32
            }
            # TODO: Refine endianness check if possible/needed based on ida_idp flags

            if connection_type == "local":
                self.device = frida.get_local_device()
            elif connection_type == "usb":
                self.device = frida.get_usb_device()
            elif connection_type == "remote":
                self.device = frida.get_device_manager().add_remote_device(host)
            else:
                raise ValueError(f"Invalid channel vector: {connection_type}")

            log_debug(f"Successfully initialized {connection_type} instrumentation channel")
            return True

        except frida.ServerNotRunningError as e:
            log_error(f"Frida server not running or accessible on {host if connection_type=='remote' else connection_type}: {e}")
            return False
        except frida.TransportError as e:
            log_error(f"Transport layer error connecting via {connection_type}: {e}")
            return False
        except Exception as e:
            log_error(f"Unhandled exception in channel initialization: {e}", exc=e)
            return False

    def attach_to_process(self, target):
        """
        Establish memory introspection context in target process.
        Corrected to handle PID logging safely.

        Args:
            target: Process identifier (PID or name)

        Returns:
            bool: True if session established, False otherwise
        """
        if not self.device:
            log_error("Attempted session creation with uninitialized device context")
            return False

        try:
            target_pid = -1
            target_name = ""
            log_message_pid = "unknown" # Placeholder for logging

            if isinstance(target, int) or (isinstance(target, str) and target.isdigit()):
                # --- Attach via PID ---
                target_pid = int(target)
                log_message_pid = str(target_pid) # We know the PID we are using
                log_debug(f"Attempting attach via PID: {target_pid}")
                self.session = self.device.attach(target_pid)
            else:
                # --- Attach via Name ---
                target_name = target
                log_message_pid = f"'{target_name}'" # Log the name used for attaching
                log_debug(f"Attempting attach via name: {target_name}")
                self.session = self.device.attach(target_name)
                # NOTE: We don't reliably get the resolved PID back onto the session object easily here.

            # --- Check Session and Log ---
            if not self.session:
                 # This case is less likely as exceptions are usually thrown on failure, but check anyway.
                 log_error(f"Session object null after attach attempt to {log_message_pid}")
                 return False

            # Log success using the PID or Name we used for the attach call
            log_debug(f"Successfully attached to process {log_message_pid}")

            # --- Post-Attach Operations ---
            try:
                # Attempt operations that require a valid session
                self.session.enable_child_gating()
                log_debug("Enabled child process tracking")
            except frida.NotSupportedError:
                 log_debug("Child gating not supported on this target/platform.")
            except Exception as e:
                # Log other potential errors but don't necessarily fail the attach
                log_debug(f"Child gating initialization failed (non-fatal): {e}")

            return True # Attach succeeded

        except frida.ProcessNotFoundError as e:
            log_error(f"Target process not found: {target}")
            return False
        except frida.PermissionDeniedError as e:
            log_error(f"Insufficient privileges for process attachment: {e}")
            return False
        except frida.TransportError as e:
             log_error(f"Frida transport error during attach: {e}")
             return False
        except Exception as e:
            # Catch any other unexpected exceptions during attach
            log_error(f"Unhandled exception in session establishment: {e}", exc=e)
            return False

    def get_process_list(self):
        """
        Enumerate available processes on target device.

        Returns:
            list: Process descriptors (pid, name, etc.) or empty list on error.
        """
        if not self.device:
            log_error("Attempted process enumeration with uninitialized device")
            return []
        try:
            processes = self.device.enumerate_processes()
            log_debug(f"Enumerated {len(processes)} processes on target device")
            return processes
        except frida.TransportError as e:
            log_error(f"Transport error during process enumeration: {e}")
            return []
        except Exception as e:
            log_error(f"Unhandled exception in process enumeration: {e}", exc=e)
            return []

    def on_message(self, message, data):
        """
        Message handler for bidirectional communication with instrumentation script.

        Args:
            message: JSON message from script execution context
            data: Binary buffer for large data transfers (unused currently)
        """
        log_debug(f"Message received: Type={message.get('type')}")
        if message['type'] == 'send':
            payload = message['payload']
            if isinstance(payload, dict) and payload.get('type') == 'error':
                self.error = payload.get('message', 'Unspecified error in script execution')
                log_error(f"Script execution error: {self.error}")
                if 'stack' in payload:
                    log_error(f"Script stack trace:\n{payload['stack']}")
            else:
                self.result = payload
                log_debug(f"Received result data from script.")
        elif message['type'] == 'error':
            self.error = message.get('description', 'Unknown script runtime error')
            stack = message.get('stack', 'No stack trace available')
            log_error(f"Fatal error in script execution:\n{self.error}\nStack:\n{stack}")
        else:
            log_debug(f"Unhandled message type: {message['type']}")

    def resolve_pointer(self, pointer_ea, pointer_info, dereference_level=1, timeout=10):
        """
        Execute memory introspection operation to resolve pointer chains.
        Corrected logging to avoid session.pid access.

        Args:
            pointer_ea: Effective address of pointer in static binary
            pointer_info: Type metadata and structural characteristics
            dereference_level: Pointer chain traversal depth
            timeout: Maximum execution time in seconds

        Returns:
            dict: Memory introspection results or error descriptor
        """
        if not self.session:
            log_error("Attempted introspection with uninitialized session")
            return {'error': 'No active instrumentation session'}

        static_imagebase = get_imagebase()
        rva = pointer_ea - static_imagebase
        module_name = self.target_module_name

        log_debug(f"Resolving pointer: EA=0x{pointer_ea:X}, RVA=0x{rva:X}, Module='{module_name}', Levels={dereference_level}")

        self.result = None
        self.error = None

        script_source = self._generate_pointer_resolution_script(
            module_name, rva, pointer_info, dereference_level)

        if DEBUG_MODE:
             log_debug("Generated Frida script.")

        try:
            self.script = self.session.create_script(script_source)
            self.script.on('message', self.on_message)
            self.script.load()
            # Corrected logging - removed reference to self.session.pid
            log_debug(f"Deployed introspection script to target process.")

            start_time = time.time()
            processed = False
            while not processed and (time.time() - start_time < timeout):
                if self.result or self.error:
                    processed = True
                else:
                    time.sleep(0.05) # Yield control briefly

            # Final check after loop
            if self.error:
                log_error(f"Script execution failed: {self.error}")
                return {'error': self.error}
            elif self.result:
                log_debug("Script execution successful.")
                return self.result
            else:
                log_error(f"Script execution timed out after {timeout} seconds.")
                # Attempt to unload script even on timeout to prevent leaks
                if self.script:
                    try: self.script.unload()
                    except: pass # Ignore errors during timeout cleanup
                return {'error': 'Execution timeout in script runtime'}

        except frida.TransportError as e:
             log_error(f"Frida transport error during script execution: {e}", exc=e)
             return {'error': f"Transport error: {e}"}
        except frida.InvalidOperationError as e:
             log_error(f"Frida invalid operation during script execution: {e}", exc=e)
             return {'error': f"Invalid operation: {e}"}
        except Exception as e:
            log_error(f"Unhandled exception in script execution: {e}", exc=e)
            return {'error': f"General execution error: {e}"}
        finally:
            # Resource cleanup
            if self.script:
                try:
                    # Detach the message handler *before* unloading
                    self.script.off('message', self.on_message)
                    self.script.unload()
                    log_debug("Unloaded script execution context")
                except frida.InvalidOperationError:
                    log_debug("Script already unloaded or session detached.")
                except Exception as e:
                    log_debug(f"Non-fatal error during script unload: {e}")
                self.script = None # Clear reference

    def _generate_pointer_resolution_script(self, module_name, rva, pointer_info, dereference_level):
        """
        Generate Frida instrumentation script for pointer chain traversal
        with code pointer detection and module offset reporting.
        """
        memory_read_directive = "readU64" if self.bitness == 64 else "readU32"
        pointer_size = self.architecture['pointer_size']

        # Determine initial read operation based on static analysis info
        initial_read_op = memory_read_directive # Default to pointer size
        if pointer_info['size'] > 0:
            if pointer_info['size'] == 8: initial_read_op = "readU64"
            elif pointer_info['size'] == 4: initial_read_op = "readU32"
            elif pointer_info['size'] == 2: initial_read_op = "readU16"
            elif pointer_info['size'] == 1: initial_read_op = "readU8"
            # else keep default pointer size read

        # Escape module name for use in JS string literal
        js_module_name = module_name.replace('\\', '\\\\').replace('"', '\\"')

        # Note: This script focuses on pointer resolution. Data reading uses _generate_data_read_script.
        return f'''
        (function() {{
            // Helper function to get module info for an address
            function getModuleInfo(address) {{
                try {{
                    const modules = Process.enumerateModules();
                    for (const mod of modules) {{
                        // Use try-catch for base/size access as they might be invalid temporarily
                        try {{
                            if (address.compare(mod.base) >= 0 &&
                                address.compare(mod.base.add(mod.size)) < 0) {{
                                return {{
                                    name: mod.name,
                                    base: mod.base.toString(),
                                    offset: address.sub(mod.base).toString(16),
                                    size: mod.size.toString(16),
                                    path: mod.path || "unknown"
                                }};
                            }}
                        }} catch (e) {{ /* Ignore module if base/size invalid */ }}
                    }}
                }} catch (e) {{ console.log("Error in getModuleInfo for " + address + ": " + e); }}
                return null; // Not found in any module
            }}

            try {{
                const ptrSize = Process.pointerSize;
                let moduleBase = Module.findBaseAddress("{js_module_name}");
                if (!moduleBase) {{
                    // Try finding module by path/name if exact name fails
                    const modules = Process.enumerateModules();
                    const found = modules.find(m => m.name === "{js_module_name}" || (m.path && m.path.endsWith("{js_module_name}")));
                    if (found) {{
                         moduleBase = found.base;
                         console.log("Found module '{js_module_name}' via secondary lookup.");
                    }} else {{
                         send({{type: "error", message: "Module identification failed: {js_module_name}"}});
                         return;
                    }}
                }}

                const pointerAddress = moduleBase.add(ptr("0x{rva:X}"));
                let dereferenceChain = [];
                let currentValue = null; // Initialize to null
                let currentAddress = pointerAddress;
                let currentModuleInfo = getModuleInfo(currentAddress); // Get info for initial address

                // Initial memory read operation (Level 0)
                try {{
                    currentValue = currentAddress.{initial_read_op}();
                }} catch(e) {{
                    console.log("Initial read with specific size failed (addr: " + currentAddress + ", op: {initial_read_op}), falling back to pointer size read.");
                    try {{
                        currentValue = currentAddress.{memory_read_directive}();
                    }} catch (e2) {{
                        send({{
                            type: "error",
                            message: `Memory access violation at base address ${{currentAddress}}: ${{e2.message}}`,
                            stack: e2.stack
                        }});
                        return;
                    }}
                }}

                // Record base pointer state (Level 0)
                dereferenceChain.push({{
                    level: 0,
                    address: currentAddress.toString(),
                    value: currentValue.toString(16), // Store value as hex string
                    offset_from_base: pointerAddress.sub(moduleBase).toString(16),
                    memory_protection: (() => {{
                        try {{ return Process.findRangeByAddress(currentAddress)?.protection || "unknown"; }}
                        catch(e) {{ return "unreadable"; }}
                    }})(),
                    module: currentModuleInfo // Include module info for level 0
                }});

                // Traverse pointer chain
                for(let i = 0; i < {dereference_level}; i++) {{
                    let nextAddress = null;
                    let valueAsString = "invalid"; // Store value before ptr conversion attempt

                    try {{
                        // Check for NULL or invalid pointer value *before* trying ptr()
                        if (currentValue === null || currentValue === undefined ||
                            (typeof currentValue.isNull === 'function' && currentValue.isNull()) ||
                            (typeof currentValue.toUInt32 === 'function' && currentValue.toUInt32() === 0) ||
                             currentValue == 0) {{ // Handle primitive 0 too
                            dereferenceChain.push({{
                                level: i + 1,
                                address: "0x0", // Represents the value read previously
                                value: "NULL",
                                error: "null_pointer_value",
                                error_details: "Value read from previous level was NULL."
                            }});
                            break; // Stop traversal
                        }}

                        valueAsString = currentValue.toString(16); // Get hex string representation

                        // Construct pointer from previous value
                        nextAddress = ptr(currentValue.toString()); // Use original value for ptr()

                    }} catch(ptrError) {{
                        // Error converting the value to a pointer
                        dereferenceChain.push({{
                            level: i + 1,
                            address: "invalid_value", // Indicate the value was bad
                            value: valueAsString, // Show the value that failed
                            error: "invalid_pointer_value",
                            error_details: `Value 0x${{valueAsString}} could not be converted to a valid pointer: ${{ptrError.message}}`
                        }});
                        break; // Stop traversal
                    }}

                    // Now we have a potentially valid nextAddress, analyze it
                    currentAddress = nextAddress; // Update currentAddress for this level (i+1)
                    let memoryProtection = "unknown";
                    let moduleInfo = null;
                    let isReadable = false;
                    let isExecutable = false;

                    try {{
                        const range = Process.findRangeByAddress(currentAddress);
                        if (range) {{
                            memoryProtection = range.protection;
                            isReadable = memoryProtection.includes('r');
                            isExecutable = memoryProtection.includes('x');
                            moduleInfo = getModuleInfo(currentAddress); // Get module info for this address
                        }} else {{
                            // Address might be valid but outside known ranges
                            memoryProtection = "unknown_range";
                            // Try a probe read to check readability (use with caution)
                            try {{ Memory.readU8(currentAddress); isReadable = true; }} catch(probeError) {{ isReadable = false; }}
                        }}
                    }} catch(e) {{
                        // Error during range lookup (less common)
                        memoryProtection = "lookup_error";
                        isReadable = false;
                        console.log("Error finding range for " + currentAddress + ": " + e);
                    }}

                    // *** Check for Executable Memory ***
                    if (isExecutable) {{
                        dereferenceChain.push({{
                            level: i + 1,
                            address: currentAddress.toString(),
                            value: "N/A", // No value read from code
                            memory_protection: memoryProtection,
                            module: moduleInfo,
                            termination_reason: "points_to_executable_code"
                        }});
                        break; // Stop traversal
                    }}

                    // Check for Non-Readable Memory
                    if (!isReadable) {{
                        dereferenceChain.push({{
                            level: i + 1,
                            address: currentAddress.toString(),
                            value: "N/A", // Cannot read value
                            error: "non_readable",
                            error_details: `Memory region not readable (Protection: ${{memoryProtection}})`,
                            memory_protection: memoryProtection,
                            module: moduleInfo
                        }});
                        break; // Stop traversal
                    }}

                    // Perform next level dereference (read pointer-sized value)
                    try {{
                        currentValue = currentAddress.{memory_read_directive}(); // Read value at currentAddress
                    }} catch(readError) {{
                        dereferenceChain.push({{
                            level: i + 1,
                            address: currentAddress.toString(), // Address where read failed
                            value: "N/A", // Read failed
                            error: "access_violation",
                            error_details: `Memory access violation reading at address: ${{readError.message}}`,
                            memory_protection: memoryProtection,
                            module: moduleInfo
                        }});
                        break; // Stop traversal
                    }}

                    // Record successful dereference state
                    dereferenceChain.push({{
                        level: i + 1,
                        address: currentAddress.toString(),
                        value: currentValue.toString(16), // Store value as hex string
                        memory_protection: memoryProtection,
                        module: moduleInfo
                    }});

                    // Check if we reached the end of requested depth
                    if (i + 1 >= {dereference_level}) {{
                         break;
                    }}

                }} // End of for loop

                // Memory context analysis for the *last valid address* in the chain
                let memoryContext = null;
                let lastValidLink = null;
                for (let j = dereferenceChain.length - 1; j >= 0; j--) {{
                    if (!dereferenceChain[j].error && !dereferenceChain[j].termination_reason && dereferenceChain[j].address !== '0x0' && dereferenceChain[j].address !== 'invalid_value') {{
                        lastValidLink = dereferenceChain[j];
                        break;
                    }}
                }}

                if (lastValidLink) {{
                    try {{
                        const finalAddress = ptr(lastValidLink.address);
                        // Read up to 16 bytes for context, but not more than ptrSize if smaller
                        const contextSize = Math.min(ptrSize, 16);
                        const memBefore = finalAddress.sub(contextSize).readByteArray(contextSize);
                        // Read the actual value location again for context
                        const memExact = finalAddress.readByteArray(contextSize);
                        const memAfter = finalAddress.add(contextSize).readByteArray(contextSize);

                        // Helper for hex conversion needed here too
                        const bytesToHexCtx = (buffer) => {{
                           if (!buffer) return "N/A";
                           return Array.from(new Uint8Array(buffer))
                               .map(b => b.toString(16).padStart(2, '0'))
                               .join(' ');
                        }};

                        memoryContext = {{
                            context_address: finalAddress.toString(),
                            before: bytesToHexCtx(memBefore),
                            exact: bytesToHexCtx(memExact),
                            after: bytesToHexCtx(memAfter)
                        }};
                    }} catch(e) {{
                        console.log(`Memory context analysis error at ${{lastValidLink.address}}: ${{e.message}}`);
                    }}
                }}

                // Return complete introspection results
                send({{
                    type: "success",
                    module_base: moduleBase.toString(),
                    pointer_name: "{pointer_info.get('name', 'unnamed')}",
                    pointer_type: "{pointer_info.get('type_name', 'unknown')}",
                    pointer_size: {pointer_info['size']},
                    module_name: "{js_module_name}",
                    rva: "0x{rva:X}",
                    chain: dereferenceChain,
                    memory_context: memoryContext,
                    architecture: {{
                        bitness: {self.bitness},
                        pointer_size: {pointer_size},
                        endianness: "{self.architecture['endianness']}"
                    }}
                }});
            }} catch(e) {{
                // Global exception handler for script initialization failures
                send({{
                    type: "error",
                    message: `Script initialization failed: ${{e.message}}`,
                    stack: e.stack
                }});
            }}
        }})();
        '''

    def read_data(self, data_ea, data_info, timeout=5):
        """
        Reads data from a specific address at runtime based on type info.

        Args:
            data_ea: Effective address of the data in static binary
            data_info: Type metadata (size, type_name)
            timeout: Maximum execution time in seconds

        Returns:
            dict: Data read results or error descriptor
        """
        if not self.session:
            log_error("Attempted data read with uninitialized session")
            return {'error': 'No active instrumentation session'}

        static_imagebase = get_imagebase()
        rva = data_ea - static_imagebase
        module_name = self.target_module_name

        log_debug(f"Reading data: EA=0x{data_ea:X}, RVA=0x{rva:X}, Module='{module_name}', Size={data_info.get('size', '?')}")

        self.result = None
        self.error = None

        # Call the data reading script generator
        script_source = self._generate_data_read_script(
            module_name, rva, data_info)
        
        if DEBUG_MODE:
             log_debug("Generated Frida data read script.")

        try:
            self.script = self.session.create_script(script_source)
            self.script.on('message', self.on_message)
            self.script.load()
            log_debug(f"Deployed data read script to target process.")

            start_time = time.time()
            processed = False
            while not processed and (time.time() - start_time < timeout):
                if self.result or self.error:
                    processed = True
                else:
                    time.sleep(0.05)

            if self.error:
                log_error(f"Script execution failed: {self.error}")
                return {'error': self.error}
            elif self.result:
                log_debug("Script execution successful.")
                return self.result
            else:
                log_error(f"Script execution timed out after {timeout} seconds.")
                if self.script:
                    try: self.script.unload()
                    except: pass
                return {'error': 'Execution timeout in script runtime'}

        except frida.TransportError as e:
             log_error(f"Frida transport error during script execution: {e}", exc=e)
             return {'error': f"Transport error: {e}"}
        except frida.InvalidOperationError as e:
             log_error(f"Frida invalid operation during script execution: {e}", exc=e)
             return {'error': f"Invalid operation: {e}"}
        except Exception as e:
            log_error(f"Unhandled exception in script execution: {e}", exc=e)
            return {'error': f"General execution error: {e}"}
        finally:
            if self.script:
                try:
                    self.script.off('message', self.on_message)
                    self.script.unload()
                    log_debug("Unloaded script execution context")
                except frida.InvalidOperationError:
                    log_debug("Script already unloaded or session detached.")
                except Exception as e:
                    log_debug(f"Non-fatal error during script unload: {e}")
                self.script = None

    def _generate_data_read_script(self, module_name, rva, data_info):
        """
        Generates Frida script to read data based on size/type.
        Includes specific handling for signed types.
        Uses simplified JS generation and robust value formatting.
        Corrected module info lookup.
        """
        # ... (Python logic to determine read_op, etc. remains the same) ...
        size = data_info.get('size', 0)
        type_name = data_info.get('type_name', 'unknown').lower().replace('_', '').replace(' ', '')
        read_op = None
        read_arg = ""
        is_byte_array = False
        is_signed = False
        is_float = False

        if size == 1:
            if 'signed' in type_name or type_name in ('int8', 'sbyte', 'char'): read_op = "readS8"; is_signed = True
            else: read_op = "readU8"
        elif size == 2:
            if 'signed' in type_name or type_name in ('int16', 'short'): read_op = "readS16"; is_signed = True
            else: read_op = "readU16"
        elif size == 4:
            if 'float' in type_name: read_op = "readFloat"; is_float = True
            elif 'signed' in type_name or type_name in ('int32', 'int', 'long'): read_op = "readS32"; is_signed = True
            else: read_op = "readU32"
        elif size == 8:
            if 'double' in type_name: read_op = "readDouble"; is_float = True
            elif 'signed' in type_name or type_name in ('int64', '__int64', 'longlong'): read_op = "readS64"; is_signed = True
            else: read_op = "readU64"
        elif size > 0:
             read_op = "readByteArray"; is_byte_array = True; read_arg = str(size)
        else:
            read_op = "readByteArray"; is_byte_array = True; size = 16; read_arg = str(size)
            log_debug(f"Unsupported or zero size, falling back to readByteArray({size}).")

        if not read_op:
             log_error("Failed to determine read operation!")
             return "send({type: 'error', message: 'Internal plugin error: Could not determine read operation'});"

        js_module_name = module_name.replace('\\', '\\\\').replace('"', '\\"')
        ptr_size = self.architecture['pointer_size']
        read_call = f"dataAddress.{read_op}({read_arg})"

        # --- Generate JavaScript ---
        return f'''
        (function() {{
            // Helper to convert ArrayBuffer to hex string
            function bytesToHex(buffer) {{
                if (!buffer) return "N/A";
                try {{
                    return Array.from(new Uint8Array(buffer))
                        .map(b => b.toString(16).padStart(2, '0'))
                        .join(' ');
                }} catch (e) {{ return "Error converting bytes"; }}
            }}

            // *** Re-introduce getModuleInfo helper ***
            function getModuleInfoForAddress(address) {{
                try {{
                    const modules = Process.enumerateModules();
                    for (const mod of modules) {{
                        try {{
                            if (address.compare(mod.base) >= 0 &&
                                address.compare(mod.base.add(mod.size)) < 0) {{
                                return {{
                                    name: mod.name,
                                    base: mod.base.toString(),
                                    offset: address.sub(mod.base).toString(16),
                                    size: mod.size.toString(16),
                                    path: mod.path || "unknown"
                                }};
                            }}
                        }} catch (e) {{ /* Ignore module if base/size invalid */ }}
                    }}
                }} catch (e) {{ console.log("Error in getModuleInfoForAddress for " + address + ": " + e); }}
                return null; // Not found in any module
            }}

            try {{
                let moduleBase = Module.findBaseAddress("{js_module_name}");
                if (!moduleBase) {{
                    const modules = Process.enumerateModules();
                    const found = modules.find(m => m.name === "{js_module_name}" || (m.path && m.path.endsWith("{js_module_name}")));
                    if (found) {{ moduleBase = found.base; }}
                    else {{
                         send({{type: "error", message: "Module identification failed: {js_module_name}"}});
                         return;
                    }}
                }}

                const dataAddress = moduleBase.add(ptr("0x{rva:X}"));
                let rawValue = null;
                let valueStr = "Error";
                let protection = "unknown";
                let moduleInfo = null; // Will be populated by helper
                let specificReadError = null;

                try {{
                    // Get protection and module info first
                    const range = Process.findRangeByAddress(dataAddress);
                    if (range) {{ protection = range.protection; }}
                    // *** CORRECTED: Use helper function to find module info ***
                    moduleInfo = getModuleInfoForAddress(dataAddress);

                    // --- Perform the read operation ---
                    try {{
                        console.log("Attempting read: {read_call}");
                        rawValue = {read_call};
                        console.log("Raw value read:", rawValue);
                    }} catch(readErr) {{
                        console.error("!!! Specific Read Error Caught: " + readErr.message + "\\nStack:" + readErr.stack);
                        specificReadError = readErr;
                    }}
                    // --- End read operation ---

                    if (specificReadError !== null) {{
                        send({{
                            type: "error",
                            message: `Failed to read data at ${{dataAddress}}: ${{specificReadError.message}}`,
                            stack: specificReadError.stack
                        }});
                        return;
                    }}

                    // --- Format the successfully read value ---
                    // (Formatting logic remains the same as previous correct version)
                    if (rawValue !== null && rawValue !== undefined) {{
                        if ({str(is_byte_array).lower()}) {{
                            valueStr = bytesToHex(rawValue);
                        }} else if (typeof rawValue === 'number') {{
                            if ({str(is_float).lower()}) {{ valueStr = rawValue.toString(10); }}
                            else if ({str(is_signed).lower()}) {{ valueStr = rawValue.toString(10); }}
                            else {{ valueStr = "0x" + rawValue.toString(16); }}
                        }} else if (typeof rawValue === 'object' && rawValue.toString) {{
                             valueStr = rawValue.toString();
                             if (!valueStr.startsWith("0x") && !isNaN(parseInt(valueStr, 16))) {{ valueStr = "0x" + valueStr; }}
                        }} else {{ valueStr = String(rawValue); }}
                        console.log("Formatted value:", valueStr);
                    }} else if (rawValue === null) {{ valueStr = "NULL_READ_UNEXPECTED"; }}

                }} catch (e) {{
                    if (!specificReadError) {{
                        send({{
                            type: "error",
                            message: `Error processing data at ${{dataAddress}}: ${{e.message}}`,
                            stack: e.stack
                        }});
                        return;
                    }}
                }}

                // Memory context analysis (only if read was successful)
                let memoryContext = null;
                if (specificReadError === null) {{
                    try {{
                        const contextSize = Math.min({ptr_size}, 16);
                        const readSize = Math.max(1, {size if size > 0 else 1});
                        const memBefore = dataAddress.sub(contextSize).readByteArray(contextSize);
                        const memExact = dataAddress.readByteArray(readSize);
                        const memAfter = dataAddress.add(readSize).readByteArray(contextSize);
                        memoryContext = {{
                            context_address: dataAddress.toString(),
                            before: bytesToHex(memBefore),
                            exact: bytesToHex(memExact),
                            after: bytesToHex(memAfter)
                        }};
                    }} catch(e) {{ console.log(`Memory context analysis error: ${{e.message}}`); }}
                }}

                // Send success message
                send({{
                    type: "success",
                    runtime_address: dataAddress.toString(),
                    value: valueStr,
                    raw_value_type: typeof rawValue,
                    size: {size},
                    type_name: "{data_info.get('type_name', 'unknown')}",
                    memory_protection: protection,
                    module_info: moduleInfo, // Send the result from the helper function
                    memory_context: memoryContext
                }});

            }} catch(e) {{
                send({{
                    type: "error",
                    message: `Script execution failed: ${{e.message}}`,
                    stack: e.stack
                }});
            }}
        }})();
        '''

    def cleanup(self):
        """Release instrumentation resources and communication channels."""
        log_debug("Starting cleanup...")
        # Script context cleanup
        if self.script:
            try:
                # Check if script is loaded before trying to unload
                # Note: Frida doesn't offer a direct is_loaded check easily accessible here.
                # Rely on try-except.
                self.script.unload()
                log_debug("Unloaded script execution context")
            except frida.InvalidOperationError:
                 log_debug("Script already unloaded or session detached during unload.")
            except Exception as e:
                log_debug(f"Non-fatal error during script unload: {e}")
            self.script = None

        # Session termination
        if self.session and not self.session.is_detached: # Check if attached before detaching
            try:
                self.session.detach()
                log_debug("Detached from instrumentation session")
            except frida.InvalidOperationError:
                 log_debug("Session already detached.")
            except Exception as e:
                log_debug(f"Non-fatal error during session detach: {e}")
            self.session = None

        # Reference clearing
        self.device = None
        self.result = None
        self.error = None
        log_debug("Cleanup complete.")


#-------------------------------------------------------------------------------
# IDA UI Integration
#-------------------------------------------------------------------------------
def is_in_data_segment(ea):
    """Determine if an address belongs to a data segment."""
    seg = ida_segment.getseg(ea)
    if not seg: return False
    seg_class = ida_segment.get_segm_class(seg)
    if seg_class in ('DATA', 'CONST', 'BSS'): return True
    seg_name = ida_segment.get_segm_name(seg)
    if seg_name in ('.data', '.rdata', '.bss', 'data', 'DATA', 'CONST', '.rodata'): return True
    if (seg.perm & ida_segment.SEGPERM_READ) and not (seg.perm & ida_segment.SEGPERM_EXEC): return True
    return False

# (Ensure necessary imports like ida_kernwin, ida_idaapi, ida_bytes, ida_name,
#  ida_segment, ida_idp, ida_ua, ida_xref, os, traceback, time are present)
# (Ensure helper functions like HAS_FRIDA, log_debug, log_error, analyze_pointer_type,
#  FridaMemoryResolver, is_in_data_segment, get_file_name, get_bitness,
#  get_module_info_for_addr_py are defined above)
# (Ensure PLUGIN_NAME constant is defined)

class ResolvePointerAction(ida_kernwin.action_handler_t):
    """
    IDA action handler for runtime pointer resolution.
    Uses a subclassed Choose for operand selection and Form for process selection.
    """
    def __init__(self):
        ida_kernwin.action_handler_t.__init__(self)

    def activate(self, ctx):
        """Action handler entry point."""
        try:
            if not HAS_FRIDA:
                ida_kernwin.warning(f"{PLUGIN_NAME} requires the Frida Python package.\n"
                                  f"Install with: pip install frida frida-tools")
                return 0

            target_ea = self._get_target_ea_from_context(ctx)
            if target_ea == ida_idaapi.BADADDR:
                log_debug("No valid target address found in context.")
                return 0

            log_debug(f"Target EA resolved: {hex(target_ea)}")
            item_info = analyze_pointer_type(target_ea) # Renamed to item_info for clarity
            log_debug(f"Memory type analysis result: {item_info}")

            # --- Decision Point: Pointer Resolution vs. Data Reading ---
            if item_info['is_pointer']:
                # Looks like a pointer, proceed with standard resolution
                log_debug("Target identified as pointer, proceeding with resolution.")
                if self._resolve_pointer_at_runtime(target_ea, item_info):
                    return 1 # Success
                else:
                    return 0 # Failed or cancelled during resolution
            else:
                # Doesn't look like a pointer, ask to read data instead
                log_debug("Target identified as data, asking user to read value.")
                prompt = (f"HIDECANCEL\n{PLUGIN_NAME}: Read Runtime Data?\n"
                          f"The selected location {item_info.get('name', hex(target_ea))} "
                          f"(Type: {item_info.get('type_name', 'unknown')}, Size: {item_info.get('size', '?')}) "
                          f"appears to be data, not a pointer.\n\n"
                          f"Do you want to read its current value at runtime?")
                if ida_kernwin.ask_yn(1, prompt) == 1: # User wants to read data
                    log_debug("User chose to read runtime data value.")
                    # Call the new data reading function
                    if self._read_data_at_runtime(target_ea, item_info):
                         return 1 # Success
                    else:
                         return 0 # Failed or cancelled during data read
                else:
                    log_debug("User cancelled data reading operation.")
                    return 0 # User cancelled

        except Exception as e:
            log_error(f"Unhandled exception in action handler: {e}", exc=e)
            ida_kernwin.warning(f"{PLUGIN_NAME} Error:\nOperation failed unexpectedly.\n\nDetails: {e}\n\nCheck console for more info.")
            return 0
        
    def _visualize_runtime_data(self, target_ea, item_info, result):
        """Displays the read runtime data value."""
        lines = []
        lines.append(f"=== {PLUGIN_NAME}: Runtime Data Read ===")
        lines.append(f"Target: {item_info.get('name', 'unknown')} at static EA 0x{target_ea:X}")
        lines.append(f"Static Type: {item_info.get('type_name', 'unknown')} ({item_info.get('size', '?')} bytes)")

        runtime_addr = result.get('runtime_address', 'unknown')
        # Attempt to format runtime address with module+offset
        display_rt_addr = runtime_addr # Default
        mod_info_rt = result.get('module_info')
        if mod_info_rt and 'base' in mod_info_rt and 'name' in mod_info_rt:
             try:
                 rt_addr_int = int(runtime_addr, 16)
                 base_addr_int = int(mod_info_rt['base'], 16)
                 offset = rt_addr_int - base_addr_int
                 mod_name = mod_info_rt['name']
                 if len(mod_name) > 30: mod_name = mod_name[:15] + "..." + mod_name[-12:]
                 display_rt_addr = f"{mod_name}+0x{offset:X} ({runtime_addr})"
             except:
                 pass # Keep default if conversion fails

        lines.append(f"Runtime Address: {display_rt_addr}")
        lines.append(f"Memory Protection: {result.get('memory_protection', 'unknown')}")
        lines.append(f"--- Runtime Value ({result.get('type_name', '?')} / {result.get('size', '?')} bytes) ---")
        lines.append(f"  {result.get('value', 'N/A')}") # Display the formatted value from Frida

        mem_ctx = result.get('memory_context')
        if mem_ctx:
            lines.append("\n--- Memory Context ---")
            lines.append(f"Center Address: {display_rt_addr}") # Use same formatted address
            lines.append(f"  Before: {mem_ctx.get('before', 'N/A')}")
            lines.append(f"  Exact:  {mem_ctx.get('exact', 'N/A')}") # Shows the actual bytes read
            lines.append(f"  After:  {mem_ctx.get('after', 'N/A')}")

        print("\n".join(lines))
        output_widget = ida_kernwin.find_widget("Output window")
        if output_widget:
            ida_kernwin.activate_widget(output_widget, True)
        else:
            log_debug("Could not find 'Output window' to activate.")

        # --- Annotation ---
        prompt = f"HIDECANCEL\n{PLUGIN_NAME}: Annotation\nAnnotate static location 0x{target_ea:X} with runtime data?"
        if ida_kernwin.ask_yn(1, prompt) == 1:
            # Create annotation like "frIDA[Data]: 0xDEADBEEF (Runtime)"
            annotation = f"frIDA[Data]: {result.get('value', 'N/A')} (Runtime)"
            ida_bytes.set_cmt(target_ea, annotation, False)
            log_debug(f"Added annotation at 0x{target_ea:X}: {annotation}")
        
    def _read_data_at_runtime(self, target_ea, item_info):
        """Handles Frida connection and calls the data reading method."""
        resolver = FridaMemoryResolver()

        # --- Connection --- (Identical to _resolve_pointer_at_runtime)
        conn_choice = ida_kernwin.ask_buttons("Local", "Remote", "USB", 1, f"{PLUGIN_NAME}: Connection Type")
        if conn_choice < 0: return False
        conn_type = ["local", "remote", "usb"][conn_choice]
        host = "127.0.0.1:27042"
        if conn_type == "remote":
            host_input = ida_kernwin.ask_text(
                0,
                host,
                "Enter remote Frida server address (host:port):"
            )
            if not host_input: return False
            host = host_input

        ida_kernwin.show_wait_box(f"Connecting to Frida ({conn_type})...")
        if not resolver.connect(conn_type, host):
            ida_kernwin.hide_wait_box()
            ida_kernwin.warning(f"{PLUGIN_NAME} Error:\nFailed to connect via {conn_type}" + (f" to {host}" if conn_type == "remote" else ""))
            return False

        # --- Process Selection --- (Identical to _resolve_pointer_at_runtime)
        ida_kernwin.replace_wait_box("Enumerating processes...")
        processes = resolver.get_process_list()
        if not processes:
            ida_kernwin.hide_wait_box()
            ida_kernwin.warning(f"{PLUGIN_NAME} Error:\nFailed to enumerate processes on the target device.")
            resolver.cleanup()
            return False

        process_items = [f"{p.pid}: {p.name}" for p in processes]
        current_bin_name = get_file_name()
        default_sel = -1
        for i, item in enumerate(process_items):
             proc_name = item.split(': ', 1)[-1]
             if current_bin_name and (current_bin_name == proc_name or f"{current_bin_name}.exe" == proc_name):
                 default_sel = i
                 break
        if default_sel == -1:
            for i, item in enumerate(process_items):
                if current_bin_name and current_bin_name.lower() in item.lower():
                    default_sel = i
                    break
        if default_sel == -1: default_sel = 0

        ida_kernwin.hide_wait_box()

        form_str = f"""STARTITEM 0
{PLUGIN_NAME}: Select Target Process

<#Select the target process to attach to#Process:{{ProcessDropdown}}>

"""
        process_dropdown = ida_kernwin.Form.DropdownListControl(
            items=process_items,
            readonly=True,
            selval=default_sel
        )
        f = ida_kernwin.Form(form_str, { 'ProcessDropdown': process_dropdown })
        f.Compile()
        ok = f.Execute()

        if not ok:
            log_debug("Process selection cancelled.")
            resolver.cleanup()
            return False

        choice_idx = process_dropdown.value
        if choice_idx < 0 or choice_idx >= len(processes):
             log_error(f"Invalid process selection index: {choice_idx}")
             resolver.cleanup()
             return False
        pid = processes[choice_idx].pid
        log_debug(f"User selected process PID: {pid}")

        # --- Attach --- (Identical to _resolve_pointer_at_runtime)
        ida_kernwin.show_wait_box(f"Attaching to process {pid}...")
        if not resolver.attach_to_process(pid):
            ida_kernwin.hide_wait_box()
            ida_kernwin.warning(f"{PLUGIN_NAME} Error:\nFailed to attach to process {pid}.\nCheck permissions and if Frida is injected/running.")
            resolver.cleanup()
            return False

        # --- Read Data ---
        ida_kernwin.replace_wait_box(f"Reading runtime data at 0x{target_ea:X}...") # Update wait box text
        read_success = False
        try:
            # *** Call the NEW resolver method ***
            result = resolver.read_data(target_ea, item_info)

            ida_kernwin.hide_wait_box()

            if result and 'error' not in result:
                log_info("Runtime data read successful.")
                # *** Call the NEW visualization method ***
                self._visualize_runtime_data(target_ea, item_info, result)
                read_success = True
            elif result and 'error' in result:
                ida_kernwin.warning(f"{PLUGIN_NAME} Data Read Error:\n{result['error']}")
            else:
                ida_kernwin.warning(f"{PLUGIN_NAME} Error:\nRuntime data read failed (no result or unexpected error).")

        except Exception as e:
            ida_kernwin.hide_wait_box()
            log_error(f"Exception during read_data call: {e}", exc=e)
            ida_kernwin.warning(f"{PLUGIN_NAME} Error:\nAn unexpected error occurred during data read.\n\nDetails: {e}")
        finally:
            resolver.cleanup()
            if ida_kernwin.is_idaq() and ida_kernwin.find_widget("Wait box"):
                 try: ida_kernwin.hide_wait_box()
                 except: pass

        return read_success

    def update(self, ctx):
        """Enable action based on context and attach to popup."""
        widget_type = ctx.widget_type
        if widget_type in (ida_kernwin.BWN_DISASM, ida_kernwin.BWN_PSEUDOCODE, ida_kernwin.BWN_HEXVIEW):
            ida_kernwin.attach_action_to_popup(
                ctx.widget,
                None,
                f"{PLUGIN_NAME}:resolve_pointer", # Registered action name
                f"{PLUGIN_NAME}/",                # Submenu path
                ida_kernwin.SETMENU_INS           # Insertion flag
            )
            return ida_kernwin.AST_ENABLE_FOR_WIDGET
        else:
            return ida_kernwin.AST_DISABLE_FOR_WIDGET

    def _get_target_ea_from_context(self, ctx):
        """Get the effective address under the cursor based on the view type."""
        widget_type = ctx.widget_type

        if widget_type == ida_kernwin.BWN_DISASM:
            instr_ea = ctx.cur_ea
            if instr_ea == ida_idaapi.BADADDR: return ida_idaapi.BADADDR

            refs = self._extract_operand_references(instr_ea)
            if refs:
                if len(refs) == 1:
                    return list(refs)[0][0]
                else:
                    # --- Multiple References: Use Subclassed Choose ---
                    options = []
                    ref_list = sorted(list(refs), key=lambda x: x[0])
                    for ref_ea, is_src in ref_list:
                        name = ida_name.get_ea_name(ref_ea, ida_name.GN_VISIBLE) or f"0x{ref_ea:X}"
                        options.append(f"{name}")

                    # --- CORRECT Choose Implementation using Subclassing ---
                    class ReferenceChooser(ida_kernwin.Choose):
                        def __init__(self, title, items_list, flags=0):
                            cols = [ ["Reference Target", 40 | ida_kernwin.Choose.CHCOL_PLAIN] ]
                            ida_kernwin.Choose.__init__(
                                self,
                                title,
                                cols,
                                flags=flags)
                            self.items = items_list

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

                        def OnGetLine(self, n):
                            if n < 0 or n >= len(self.items): return []
                            return [self.items[n]]

                    chooser = ReferenceChooser(
                        f"{PLUGIN_NAME}: Select Reference Target",
                        options,
                        flags=ida_kernwin.Choose.CH_MODAL | ida_kernwin.Choose.CH_NOBTNS
                    )
                    choice_idx = chooser.Show()
                    # --- End of CORRECT Choose Implementation ---

                    if choice_idx >= 0 and choice_idx < len(ref_list):
                        return ref_list[choice_idx][0]
                    else:
                        return ida_idaapi.BADADDR

            else: # No operand references found
                return instr_ea

        elif widget_type == ida_kernwin.BWN_PSEUDOCODE:
            # TODO: Future improvement - investigate vu.get_selected_item()
            return ctx.cur_ea

        elif widget_type == ida_kernwin.BWN_HEXVIEW:
            return ctx.cur_ea

        return ida_idaapi.BADADDR

    def _extract_operand_references(self, instr_ea):
        """Extract memory references from instruction operands."""
        refs = set()
        insn = ida_ua.insn_t()
        if ida_ua.decode_insn(insn, instr_ea) <= 0:
            return refs

        for i in range(ida_ua.UA_MAXOP):
            op = insn.ops[i]
            if op.type == ida_ua.o_void: break

            ref_ea = ida_idaapi.BADADDR
            is_data_target = False

            if op.type == ida_ua.o_mem:
                ref_ea = op.addr
                if ida_bytes.is_data(ida_bytes.get_flags(ref_ea)) or is_in_data_segment(ref_ea):
                    is_data_target = True

            elif op.type == ida_ua.o_displ:
                ref_ea = op.addr
                if ida_bytes.is_data(ida_bytes.get_flags(ref_ea)) or is_in_data_segment(ref_ea):
                    is_data_target = True
                elif ida_idp.ph.id == ida_idp.PLFM_386 and get_bitness() == 64 and op.specflag1 == ida_ua.X86_OP_ADDR_RIP:
                     rip_target_ea = instr_ea + insn.size + op.addr
                     if is_in_data_segment(rip_target_ea) or ida_bytes.is_data(ida_bytes.get_flags(rip_target_ea)):
                         ref_ea = rip_target_ea
                         is_data_target = True
                     else:
                         ref_ea = ida_idaapi.BADADDR

            elif op.type == ida_ua.o_imm:
                imm_val = op.value
                if ida_segment.getseg(imm_val) is not None:
                    if is_in_data_segment(imm_val) or ida_bytes.is_data(ida_bytes.get_flags(imm_val)):
                        ref_ea = imm_val
                        is_data_target = True

            if ref_ea != ida_idaapi.BADADDR and is_data_target:
                 refs.add((ref_ea, i != 0))
                 log_debug(f"Extracted operand reference: EA=0x{ref_ea:X}, Type={op.type}, IsData={is_data_target}, OpIdx={i}")

        return refs

    def _resolve_pointer_at_runtime(self, target_ea, pointer_info):
        """Core logic to connect, attach, and resolve using Frida."""
        resolver = FridaMemoryResolver()

        # --- Connection ---
        conn_choice = ida_kernwin.ask_buttons("Local", "Remote", "USB", 1, f"{PLUGIN_NAME}: Connection Type")
        if conn_choice < 0: return False
        conn_type = ["local", "remote", "usb"][conn_choice]
        host = "127.0.0.1:27042"
        if conn_type == "remote":
            host_input = ida_kernwin.ask_text(
                0,
                host,
                "Enter remote Frida server address (host:port):"
            )
            if not host_input: return False
            host = host_input

        ida_kernwin.show_wait_box(f"Connecting to Frida ({conn_type})...")
        if not resolver.connect(conn_type, host):
            ida_kernwin.hide_wait_box()
            ida_kernwin.warning(f"{PLUGIN_NAME} Error:\nFailed to connect via {conn_type}" + (f" to {host}" if conn_type == "remote" else ""))
            return False

        # --- Process Selection (Using Form with DropdownListControl) ---
        ida_kernwin.replace_wait_box("Enumerating processes...")
        processes = resolver.get_process_list()
        if not processes:
            ida_kernwin.hide_wait_box()
            ida_kernwin.warning(f"{PLUGIN_NAME} Error:\nFailed to enumerate processes on the target device.")
            resolver.cleanup()
            return False

        process_items = [f"{p.pid}: {p.name}" for p in processes]
        current_bin_name = get_file_name()
        default_sel = -1
        for i, item in enumerate(process_items):
             proc_name = item.split(': ', 1)[-1]
             if current_bin_name and (current_bin_name == proc_name or f"{current_bin_name}.exe" == proc_name):
                 default_sel = i
                 break
        if default_sel == -1:
            for i, item in enumerate(process_items):
                if current_bin_name and current_bin_name.lower() in item.lower():
                    default_sel = i
                    break
        if default_sel == -1: default_sel = 0

        ida_kernwin.hide_wait_box()

        form_str = f"""STARTITEM 0
{PLUGIN_NAME}: Select Target Process

<#Select the target process to attach to#Process:{{ProcessDropdown}}>

"""
        process_dropdown = ida_kernwin.Form.DropdownListControl(
            items=process_items,
            readonly=True,
            selval=default_sel
        )
        f = ida_kernwin.Form(form_str, {
            'ProcessDropdown': process_dropdown
        })

        f.Compile()
        ok = f.Execute()

        if not ok:
            log_debug("Process selection cancelled.")
            resolver.cleanup()
            return False

        choice_idx = process_dropdown.value
        if choice_idx < 0 or choice_idx >= len(processes):
             log_error(f"Invalid process selection index: {choice_idx}")
             resolver.cleanup()
             return False

        pid = processes[choice_idx].pid
        log_debug(f"User selected process PID: {pid}")
        # --- End of Process Selection using Form ---

        # --- Attach ---
        ida_kernwin.show_wait_box(f"Attaching to process {pid}...")
        if not resolver.attach_to_process(pid):
            ida_kernwin.hide_wait_box()
            ida_kernwin.warning(f"{PLUGIN_NAME} Error:\nFailed to attach to process {pid}.\nCheck permissions and if Frida is injected/running.")
            resolver.cleanup()
            return False

        # --- Dereference Level ---
        ida_kernwin.hide_wait_box() # Hide before asking level

        # Corrected ask_long call with only 2 arguments: default value and prompt
        default_deref_level = 1 # Define the default value
        prompt_message = (f"{PLUGIN_NAME}: Traversal Depth\n\n" # Combine title/prompt
                          "Enter pointer chain dereference depth (e.g., 1, 2, ...):")

        deref_level = ida_kernwin.ask_long(default_deref_level, prompt_message)

        # Handle cancellation or invalid input (ask_long returns None on cancel)
        if deref_level is None: # Check specifically for None
            log_debug("Dereference level selection cancelled.")
            resolver.cleanup()
            return False

        # Optional: Add validation if needed (e.g., ensure > 0)
        if deref_level < 1:
             log_debug(f"Invalid dereference level entered: {deref_level}. Defaulting to 1.")
             deref_level = 1 # Or ask again, or cancel

        # --- Resolve ---
        ida_kernwin.show_wait_box(f"Performing runtime resolution (Depth: {deref_level})...")
        resolution_success = False
        try:
            result = resolver.resolve_pointer(target_ea, pointer_info, deref_level)
            ida_kernwin.hide_wait_box()

            if result and 'error' not in result:
                log_info("Runtime resolution successful.")
                self._visualize_memory_introspection(target_ea, result)
                resolution_success = True
            elif result and 'error' in result:
                ida_kernwin.warning(f"{PLUGIN_NAME} Resolution Error:\n{result['error']}")
            else:
                ida_kernwin.warning(f"{PLUGIN_NAME} Error:\nMemory introspection failed (no result or unexpected error).")

        except Exception as e:
            ida_kernwin.hide_wait_box()
            log_error(f"Exception during resolve_pointer call: {e}", exc=e)
            ida_kernwin.warning(f"{PLUGIN_NAME} Error:\nAn unexpected error occurred during resolution.\n\nDetails: {e}")
        finally:
            resolver.cleanup()
            if ida_kernwin.is_idaq() and ida_kernwin.find_widget("Wait box"):
                 try: ida_kernwin.hide_wait_box()
                 except: pass

        return resolution_success

    def _format_address_for_display(self, link):
        """Formats an address using module+offset if available."""
        addr_str = link.get('address', 'unknown')
        module_info = link.get('module')
        if not isinstance(addr_str, str): addr_str = str(addr_str)

        if not addr_str.startswith("0x"):
             if addr_str.lower() in ("invalid_value", "0x0", "null"):
                 return addr_str
             else:
                 try:
                     int(addr_str, 16)
                     addr_str = f"0x{addr_str}"
                 except ValueError:
                     return addr_str

        if module_info and 'name' in module_info and 'offset' in module_info:
            offset_str = module_info['offset']
            if not offset_str.startswith('0x'): offset_str = f"0x{offset_str}"
            mod_name = module_info['name']
            if len(mod_name) > 30: mod_name = mod_name[:15] + "..." + mod_name[-12:]
            return f"{mod_name}+{offset_str} ({addr_str})"
        else:
            return addr_str

    def _visualize_memory_introspection(self, target_ea, result):
        """Display the resolution results in IDA's output window."""
        chain = result.get('chain', [])
        if not chain:
            ida_kernwin.warning(f"{PLUGIN_NAME}: Invalid result format (missing chain).")
            return

        lines = []
        lines.append(f"=== {PLUGIN_NAME}: Runtime Memory Analysis ===")
        lines.append(f"Target: {result.get('pointer_name', 'unknown')} at static EA 0x{target_ea:X} (RVA: {result.get('rva', 'unknown')})")
        lines.append(f"Static Type: {result.get('pointer_type', 'unknown')} ({result.get('pointer_size', '?')} bytes)")
        lines.append(f"Runtime Module Base ({result.get('module_name', 'unknown')}): {result.get('module_base', 'unknown')}")
        arch_info = result.get('architecture')
        if arch_info:
            lines.append(f"Architecture: {arch_info.get('bitness', '?')}-bit, {arch_info.get('endianness', '?')} endian")

        lines.append("\n--- Memory Traversal Path ---")

        for i, link in enumerate(chain):
            level = link.get('level', -1)
            indent = "  " * level

            display_addr = self._format_address_for_display(link)
            lines.append(f"{indent}[L{level}] Address: {display_addr}")

            if 'value' in link and link['value'] != "N/A":
                 value_str = link['value']
                 if not isinstance(value_str, str): value_str = str(value_str)
                 if not value_str.startswith("0x") and value_str.lower() != "null": value_str = f"0x{value_str}"
                 lines.append(f"{indent}  -> Value: {value_str}")

            if 'memory_protection' in link:
                lines.append(f"{indent}     Prot: {link.get('memory_protection', 'unknown')}")

            if 'termination_reason' in link:
                reason = link['termination_reason'].replace('_', ' ').title()
                lines.append(f"{indent}     Stop: {reason}")

            if 'error' in link:
                error_type = link.get('error', 'unknown').replace('_', ' ').title()
                lines.append(f"{indent}     ERROR: {error_type}")
                if 'error_details' in link:
                    lines.append(f"{indent}       Details: {link.get('error_details', 'none')}")

            if level >= 0 and i < len(chain) - 1:
                 lines.append(f"{indent}  --------------------")

        mem_ctx = result.get('memory_context')
        if mem_ctx:
            lines.append("\n--- Memory Context (Around Last Valid Address) ---")
            ctx_addr_str = mem_ctx.get('context_address', '?')
            ctx_module_info = get_module_info_for_addr_py(ctx_addr_str)
            display_ctx_addr = self._format_address_for_display({'address': ctx_addr_str, 'module': ctx_module_info})
            lines.append(f"Center Address: {display_ctx_addr}")
            lines.append(f"  Before: {mem_ctx.get('before', 'N/A')}")
            lines.append(f"  Exact:  {mem_ctx.get('exact', 'N/A')}")
            lines.append(f"  After:  {mem_ctx.get('after', 'N/A')}")

        print("\n".join(lines))

        # Find the widget by its title
        output_widget = ida_kernwin.find_widget("Output window")

        # Check if the widget was found before trying to activate it
        if output_widget:
            ida_kernwin.activate_widget(output_widget, True) # Pass the widget object
        else:
            log_debug("Could not find 'Output window' to activate.")

        prompt = f"HIDECANCEL\n{PLUGIN_NAME}: Annotation\nAnnotate static location 0x{target_ea:X} with runtime result?"

        if ida_kernwin.ask_yn(1, prompt) == 1:
            if len(chain) > 0:
                # Format L0 value
                l0_val_str = chain[0].get('value', '?')
                if not isinstance(l0_val_str, str): l0_val_str = str(l0_val_str)
                if not l0_val_str.startswith("0x"): l0_val_str = f"0x{l0_val_str}"
                annotation = f"frIDA[L0]: {l0_val_str}"

                last_link = chain[-1]
                last_level = last_link.get('level', 0)

                if last_level > 0:
                    annotation += " -> "
                    runtime_addr_str = last_link.get('address', 'unknown') # Runtime address as string
                    display_target = runtime_addr_str # Default display is runtime address

                    # --- Attempt Symbol Resolution ---
                    module_info = last_link.get('module')
                    static_ea = ida_idaapi.BADADDR

                    # Try to resolve only if module info is present and seems valid
                    if module_info and 'name' in module_info and 'offset' in module_info:
                        try:
                            offset = int(module_info['offset'], 16)
                            module_name = module_info['name']
                            main_module_name = result.get('module_name', '')

                            # Heuristic: Assume it's the main module if names match
                            # TODO: Enhance this to handle other loaded modules if needed
                            if module_name == main_module_name:
                                static_base = get_imagebase()
                                if static_base != ida_idaapi.BADADDR:
                                    static_ea = static_base + offset
                                    log_debug(f"Calculated static EA: 0x{static_ea:X} (Base: 0x{static_base:X}, Offset: 0x{offset:X})")
                            else:
                                # Handling other modules requires finding their static base in IDA,
                                # which can be complex. Log for now.
                                log_debug(f"Target address in different module '{module_name}', symbol resolution might be incomplete.")
                                # As a fallback, try using the segment base if available
                                seg = ida_segment.getseg(target_ea) # Use original target_ea segment as hint? No, use runtime addr seg
                                if runtime_addr_str.startswith("0x"):
                                     rt_addr_int = int(runtime_addr_str, 16)
                                     seg_rt = ida_segment.getseg(rt_addr_int)
                                     if seg_rt:
                                         # This is still heuristic, might not map correctly statically
                                         log_debug(f"Using segment base 0x{seg_rt.start_ea:X} as potential static base for module {module_name}")
                                         # static_ea = seg_rt.start_ea + offset # This is likely incorrect mapping

                        except ValueError:
                            log_debug(f"Could not parse offset: {module_info.get('offset')}")
                        except Exception as e:
                            log_error(f"Error during static EA calculation: {e}", exc=e)

                    # Get symbol name if we found a valid static EA
                    symbol_name = None
                    if static_ea != ida_idaapi.BADADDR:
                        # Use GN_VISIBLE to get user-friendly names
                        symbol_name = ida_name.get_ea_name(static_ea, ida_name.GN_VISIBLE)
                        log_debug(f"Symbol lookup for 0x{static_ea:X}: '{symbol_name}'")

                    # Use symbol name if found, otherwise format module+offset
                    if symbol_name:
                        display_target = symbol_name
                    elif module_info and 'name' in module_info and 'offset' in module_info:
                         # Fallback to module+offset if symbol not found but info exists
                         display_target = self._format_address_for_display(last_link).split(' ')[0] # Get "module+offset" part
                    # --- End Symbol Resolution ---


                    # Construct the final part of the annotation
                    if 'error' in last_link:
                        error_type = last_link.get('error', 'unknown').replace('_',' ').title()
                        # Include runtime address in error for context
                        annotation += f"[L{last_level} ERROR: {error_type} @ {display_target} ({runtime_addr_str})]"
                    elif 'termination_reason' in last_link:
                        reason = last_link['termination_reason'].replace('_',' ').title()
                        annotation += f"[L{last_level} Stop: {reason} @ {display_target} ({runtime_addr_str})]"
                    elif 'value' in last_link and last_link['value'] != 'N/A':
                         value_str = last_link['value']
                         if not isinstance(value_str, str): value_str = str(value_str)
                         if not value_str.startswith("0x") and value_str.lower() != "null": value_str = f"0x{value_str}"
                         # Show resolved value and the symbol/location it points to
                         annotation += f"[L{last_level}]: {value_str} @ {display_target} ({runtime_addr_str})"
                    else:
                         # Fallback if no value/error/term
                         annotation += f"[L{last_level} @ {display_target} ({runtime_addr_str})]"

                # Add comment (non-repeatable)
                ida_bytes.set_cmt(target_ea, annotation, False)
                log_debug(f"Added annotation at 0x{target_ea:X}: {annotation}")

# End of ResolvePointerAction class


# Helper function (Python side) to get module info - used for context display
def get_module_info_for_addr_py(addr_str):
    try:
        if not addr_str or addr_str == '?': return None
        addr = int(addr_str, 16)
        # Prioritize Frida's module list if available (more accurate runtime info)
        # This requires passing the full module list back or querying Frida again,
        # which adds complexity. Stick to IDA's view for now.

        # Use IDA's module info first
        mod_info = ida_nalt.get_module_info(addr)
        if mod_info and mod_info.base != ida_idaapi.BADADDR:
            offset = addr - mod_info.base
            mod_name = os.path.basename(mod_info.name or "unknown")
            return {'name': mod_name, 'offset': f"0x{offset:X}"}

        # Fallback to segment info
        seg = ida_segment.getseg(addr)
        if seg:
             seg_name = ida_segment.get_segm_name(seg)
             # Heuristic: if segment base is near imagebase, assume it's part of the main module
             # This is less reliable with ASLR but better than nothing
             static_imagebase = get_imagebase()
             if abs(seg.start_ea - static_imagebase) < 0x200000: # Wider threshold
                  offset = addr - static_imagebase
                  return {'name': get_file_name(), 'offset': f"0x{offset:X}"}
             else:
                  # Treat as offset from segment start if not near imagebase
                  offset = addr - seg.start_ea
                  return {'name': f"Seg:{seg_name}", 'offset': f"0x{offset:X}"}

    except ValueError:
        pass # Invalid address string
    except Exception as e:
        log_debug(f"Error getting module info for context addr {addr_str}: {e}")
    return None


#-------------------------------------------------------------------------------
# Plugin Registration
#-------------------------------------------------------------------------------
class frIDAPlugin(ida_idaapi.plugin_t):
    """Main plugin registration class for IDA Pro."""
    flags = ida_idaapi.PLUGIN_KEEP
    comment = "Runtime Pointer Resolution using Frida"
    help = "Right-click pointers/data and select 'Resolve Pointer at Runtime' (Ctrl+Alt+F)"
    wanted_name = PLUGIN_NAME
    wanted_hotkey = "" # Hotkey registered via action

    def init(self):
        """Initialize the plugin."""
        global HAS_FRIDA

        try:
            import frida
            HAS_FRIDA = True
        except ImportError:
            HAS_FRIDA = False
            print(f"[{PLUGIN_NAME}] Warning: Frida Python package not found.")
            print(f"[{PLUGIN_NAME}] Install with: pip install frida frida-tools")
            print(f"[{PLUGIN_NAME}] Plugin will load but functionality will be disabled.")

        action_desc = ida_kernwin.action_desc_t(
            f"{PLUGIN_NAME}:resolve_pointer",
            "Resolve Pointer at Runtime",
            ResolvePointerAction(),
            PLUGIN_HOTKEY,
            "Connect to Frida and resolve this pointer's runtime value chain",
            -1
        )

        if not ida_kernwin.register_action(action_desc):
            log_error("Failed to register plugin action")
            return ida_idaapi.PLUGIN_SKIP

        ida_kernwin.attach_action_to_menu(
            f"Edit/Plugins/{PLUGIN_NAME}/Resolve Pointer at Runtime",
            f"{PLUGIN_NAME}:resolve_pointer",
            ida_kernwin.SETMENU_APP
        )

        log_info(f"Plugin v{PLUGIN_VERSION} initialized. Hotkey: {PLUGIN_HOTKEY}")
        print(f"[{PLUGIN_NAME}] Ready. Right-click on pointers/data or use {PLUGIN_HOTKEY}.")

        return ida_idaapi.PLUGIN_KEEP

    def run(self, arg):
        """Called when plugin is selected from the main menu."""
        ida_kernwin.info(
            f"{PLUGIN_NAME} v{PLUGIN_VERSION}\n\n"
            f"This plugin resolves static pointers/data at runtime using Frida.\n\n"
            f"Usage:\n"
            f"1. Right-click on a pointer, data reference, or immediate value\n"
            f"   in Disassembly, Pseudocode, or Hex View.\n"
            f"2. Select '{PLUGIN_NAME}/Resolve Pointer at Runtime' from the context menu.\n"
            f"3. Alternatively, place the cursor and press {PLUGIN_HOTKEY}.\n\n"
            f"Requires:\n"
            f"- Frida Python package (`pip install frida frida-tools`)\n"
            f"- Target process running with frida-server or frida-gadget."
        )

    def term(self):
        """Terminate the plugin."""
        ida_kernwin.detach_action_from_menu(
             f"Edit/Plugins/{PLUGIN_NAME}/Resolve Pointer at Runtime",
             f"{PLUGIN_NAME}:resolve_pointer"
        )
        ida_kernwin.unregister_action(f"{PLUGIN_NAME}:resolve_pointer")
        log_info("Plugin terminated")

#-------------------------------------------------------------------------------
# Plugin Entry Point
#-------------------------------------------------------------------------------
def PLUGIN_ENTRY():
    """IDA Pro plugin entry point."""
    return frIDAPlugin()


repository: runope/reserve_android_notes
file: ida/scripts/traceNatives.py
wanted_name: traceNatives
comment: traceNatives
created_at: 2022-04-25 09:06:38+00:00
pushed_at: 2022-08-09 06:46:31+00:00
url:
https://github.com/runope/reserve_android_notes/blob/9f436a36245b6c6e0ca0d2b93ee1872960b83d7e/ida/scripts/traceNatives.py

content:
# -*- coding:utf-8 -*-
import os
from idaapi import plugin_t
from idaapi import PLUGIN_PROC
from idaapi import PLUGIN_OK
import ida_nalt
import idaapi
import idautils
import idc
import time


# 获取SO文件名和路径
def getSoPathAndName():
    fullpath = ida_nalt.get_input_file_path()
    filepath,filename = os.path.split(fullpath)
    return filepath,filename

# 获取代码段的范围
def getSegAddr():
    textStart = []
    textEnd = []

    for seg in idautils.Segments():
        if (idc.get_segm_name(seg)).lower() == '.text' or (
        idc.get_segm_name(seg)).lower() == 'text':
            tempStart = idc.get_segm_start(seg)
            tempEnd = idc.get_segm_end(seg)

            textStart.append(tempStart)
            textEnd.append(tempEnd)

    return min(textStart), max(textEnd)


class traceNatives(plugin_t):
    flags = PLUGIN_PROC
    comment = "traceNatives"
    help = ""
    wanted_name = "traceNatives"
    wanted_hotkey = ""

    def init(self):
        print("traceNatives(v0.1) plugin has been loaded.")
        return PLUGIN_OK

    def run(self, arg):
        # 查找需要的函数
        ea, ed = getSegAddr()
        search_result = []
        for func in idautils.Functions(ea, ed):
            try:
                functionName = str(idaapi.ida_funcs.get_func_name(func))
                if len(list(idautils.FuncItems(func))) > 100:
                    # 如果是thumb模式，地址+1
                    arm_or_thumb = idc.get_sreg(func, "T")
                    if arm_or_thumb:
                        func += 1
                    search_result.append(hex(func))
            except:
                pass

        so_path, so_name = getSoPathAndName()
        search_result = [f"-a '{so_name}!{offset}'" for offset in search_result]
        search_result = " ".join(search_result)

        script_name = so_name.split(".")[0] + "_" + str(int(time.time())) +".txt"
        save_path = os.path.join(so_path, script_name)
        with open(save_path, "w", encoding="utf-8")as F:
            F.write(search_result)

        print("使用方法如下：")
        print(f"frida-trace -UF -O {save_path}")

    def term(self):
        pass


def PLUGIN_ENTRY():
    return traceNatives()


repository: jonpalmisc/ida_graphviz_export
file: graphviz.py
description: Export control flow graphs from IDA as Graphviz DOT code
wanted_name: graphviz
comment: Graphviz exporter
created_at: 2024-05-10 23:48:09+00:00
pushed_at: 2024-05-11 00:12:14+00:00
url:
https://github.com/jonpalmisc/ida_graphviz_export/blob/3438a946348c349be6449d6e67f6d04f168657d7/graphviz.py

content:
import ida_idaapi
import ida_kernwin
import ida_funcs
import ida_gdl
import ida_lines
import ida_name
import ida_ua


def _node_name(block: ida_gdl.BasicBlock) -> str:
    return f"bb_{block.start_ea}"


def _label_escape(line: str) -> str:
    return line.replace("\n", "\\l").replace('"', '\\"')


NODE_STYLE = f"shape=box fontname=Courier"


def to_dot(func: ida_funcs.func_t) -> str:
    nodes = []
    edges = []

    blocks = ida_gdl.FlowChart(func)
    for block in blocks:
        body = ""

        if label := ida_name.get_short_name(block.start_ea):
            body += label + ":\n"

        cur_ea = block.start_ea
        while cur_ea < block.end_ea:
            insn = ida_ua.insn_t()
            insn_len = ida_ua.decode_insn(insn, cur_ea)
            if not insn:
                print("Failed to decode instruction while making graph!")
                break

            line = ida_lines.generate_disasm_line(cur_ea)
            body += ida_lines.tag_remove(line) + "\n"

            cur_ea += insn_len

        nodes.append(
            f'  {_node_name(block)}[{NODE_STYLE} label="{_label_escape(body)}"];'
        )

        for next_block in block.succs():
            edges.append(f"  {_node_name(block)} -> {_node_name(next_block)};")

    return "\n".join(["digraph {"] + edges + [""] + nodes + ["}"])


class export_dot_handler_t(ida_kernwin.action_handler_t):
    dump_only: bool

    def __init__(self, dump_only: bool = False):
        ida_kernwin.action_handler_t.__init__(self)
        self.dump_only = dump_only

    def activate(self, ctx):  # pyright: ignore
        dot = to_dot(ctx.cur_func)
        if self.dump_only:
            print(dot)
            return 1

        path = ida_kernwin.ask_file(True, "graph.dot", "Save graph as DOT file")
        if path:
            with open(path, "w") as f:
                f.write(dot)

        return 1

    def update(self, ctx):
        return ida_kernwin.AST_ENABLE if ctx.cur_func else ida_kernwin.AST_DISABLE


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

    def finish_populating_widget_popup(self, widget, popup):  # pyright: ignore
        if ida_kernwin.get_widget_type(widget) != ida_kernwin.BWN_DISASM:
            return

        ida_kernwin.attach_action_to_popup(
            widget,
            popup,
            "graphviz:ProduceDOT",
            f"&Graphviz/",
            ida_kernwin.SETMENU_APP,
        )
        ida_kernwin.attach_action_to_popup(
            widget,
            popup,
            "graphviz:DumpDOT",
            f"&Graphviz/",
            ida_kernwin.SETMENU_APP,
        )


class graphviz_plugin_t(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_DRAW | ida_idaapi.PLUGIN_HIDE
    help = ""
    comment = "Graphviz exporter"
    wanted_name = "graphviz"
    wanted_hotkey = ""

    ui_hooks: graphviz_ui_hooks_t

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

        ida_kernwin.register_action(
            ida_kernwin.action_desc_t(
                "graphviz:ProduceDOT",
                "~C~reate Graphviz DOT file...",
                export_dot_handler_t(False),
                None,
                "Export the current function's CFG as Graphviz DOT code",
            )
        )
        ida_kernwin.register_action(
            ida_kernwin.action_desc_t(
                "graphviz:DumpDOT",
                "~D~ump Graphviz DOT code",
                export_dot_handler_t(True),
                None,
                "Print the current function's CFG as Graphviz DOT code in the command line window",
            )
        )

        ida_kernwin.attach_action_to_menu(
            "File/Produce file/",
            "graphviz:ProduceDOT",
            ida_kernwin.SETMENU_APP | ida_kernwin.SETMENU_ENSURE_SEP,
        )

        self.ui_hooks.hook()
        return ida_idaapi.PLUGIN_KEEP

    def run(self):  # pyright: ignore
        pass

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


def PLUGIN_ENTRY():
    return graphviz_plugin_t()


repository: zhouzhicun/ida-scripts
file: python/IDA/IDAScripts/plugin_zzSearchCode.py
description: ida相关脚本，包括字符串解密，BR/BLR混淆, 常见加密算法实现。
created_at: 2024-09-10 08:31:38+00:00
pushed_at: 2024-09-18 03:02:56+00:00
url:
https://github.com/zhouzhicun/ida-scripts/blob/fa533f6a4130921a0013133a9777847842d87c63/python/IDA/IDAScripts/plugin_zzSearchCode.py

content:
import idaapi
from ida_idaapi import plugin_t

import zzPluginBase.utils as utils
import zzPluginBase.funcUtils as funcUtils
import zzPluginBase.pluginUtils as pluginUtils



################################ 插件逻辑 ##########################################

svc0_code = "010000D4"       #svc 0
svc80_code = "011000D4"      #svc 0x80

def searchAllSVC():
    segNames = None
    svc0Arr = utils.searchCode(segNames, svc0_code)
    svc80Arr = utils.searchCode(segNames, svc80_code)
    printResult("found SVC0 at all: \n", svc0Arr)
    printResult("found SVC80 at all: \n", svc80Arr)

    segNames = ['.text']
    svc0Arr = utils.searchCode(segNames, svc0_code)
    svc80Arr = utils.searchCode(segNames, svc80_code)
    printResult("found SVC0 at .text: \n", svc0Arr)
    printResult("found SVC80 at .text: \n", svc80Arr)


def searchBR():
    brCodeArr = funcUtils.get_all_instructions("BR")
    print("found BR at: \n" + formatCode(brCodeArr))

def searchBLR():
    brCodeArr = funcUtils.get_all_instructions("BLR")
    print("found BLR at: \n" + formatCode(brCodeArr))


def searchCSEL():
    cselCodeArr = funcUtils.get_all_instructions("CSEL")
    ssetCodeArr = funcUtils.get_all_instructions("CSET")
    print("found CSEL at: \n" + formatCode(cselCodeArr))
    print("found CSET at: \n" + formatCode(ssetCodeArr))


def searchBR_CSEL():
    br_csel_codeArr = funcUtils.get_all_BR_CSEL()

    desc = '-----------------------  dump br-csel -------------------------------\n'
    for code in br_csel_codeArr:
        br_addr = code[0]
        csel_addr = code[1]
        csel_asm = code[2]
        desc += f"0x{br_addr:08X} 0x{csel_addr:08X}: {csel_asm}" + '\n'
    print(desc)


def formatCode(codeArr):
    result = ""
    for code in codeArr:
        address = code[0]
        disasm = code[1]
        result += f"0x{address:08X}: {disasm}" + '\n'
    return result


def printResult(tip, matchAddrs):
    str = tip
    for addr in matchAddrs:
        str += hex(addr) + ","
    print(str)  



#################################### 插件配置 ##################################################


ZZSearchCode_wanted_name = 'ZZSearchCode'
ZZSearchCode_comment = ''
ZZSearchCode_help = ''
ZZSearchCode_wanted_hotkey = ''
ZZSearchCode_flags = idaapi.PLUGIN_KEEP

ZZSearchCodeMenuConfig = pluginUtils.PluginMenuConfig("ZZSearchCode/", [
    pluginUtils.PluginSubMenu('my:search_svc', '搜索SVC指令', searchAllSVC),
    pluginUtils.PluginSubMenu('my:search_br', '搜索BR指令', searchBR),
    pluginUtils.PluginSubMenu('my:search_blr', '搜索BLR指令', searchBLR),
    pluginUtils.PluginSubMenu('my:search_csel', '搜索CSEL/CSET指令', searchCSEL),
    pluginUtils.PluginSubMenu('my:search_br_csel', '搜索 BR + CSEL/CSET指令', searchBR_CSEL)
])


#################################### 插件框架 ##################################################

class ZZSearchCodeUIHooks(idaapi.UI_Hooks):
    def finish_populating_widget_popup(self, widget, popup_handle, ctx):
        pluginUtils.registerUIHook(widget, popup_handle, ZZSearchCodeMenuConfig)


class ZZSearchCode(plugin_t):

    wanted_name = ZZSearchCode_wanted_name
    comment = ZZSearchCode_comment
    help = ZZSearchCode_help
    wanted_hotkey = ZZSearchCode_wanted_hotkey
    flags = ZZSearchCode_flags


    def init(self):

        pluginUtils.registerAction(ZZSearchCodeMenuConfig)

        global my_ui_hooks
        my_ui_hooks = ZZSearchCodeUIHooks()
        my_ui_hooks.hook()

        return idaapi.PLUGIN_KEEP

    # 插件运行中 这里是主要逻辑
    def run(self, arg: int):
        pass

    def term(self):
        pass 


# 注册插件
def PLUGIN_ENTRY():
    return ZZSearchCode()


repository: clovme/WTools
file: software/Reverse Cracking/IDA Professional 9.0/python/examples/hexrays/vds10.py
wanted_name: Optimize DbgRaiseAssertionFailure (IDAPython)
comment: Sample plugin10 for Hex-Rays decompiler
created_at: 2025-03-27 00:14:54+00:00
pushed_at: 2025-05-11 10:18:21+00:00
url:
https://github.com/clovme/WTools/blob/f12f01c833a6a6241799a9d01e9117cc28a65dbc/software/Reverse%20Cracking/IDA%20Professional%209.0/python/examples/hexrays/vds10.py

content:
"""
summary: a custom microcode instruction optimization rule

description:
  Installs a custom microcode instruction optimization rule,
  to transform:

      call   !DbgRaiseAssertionFailure <fast:>.0

  into

      call   !DbgRaiseAssertionFailure <fast:"char *" "assertion text">.0

  To see this plugin in action please use arm64_brk.i64
"""

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

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

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

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

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

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

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

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

def PLUGIN_ENTRY():
    return my_plugin_t()



repository: sys4/tlsrpt-reporter
file: tests/test_entrypoints.py
description: An application suite to receive TLSRPT datagrams and to generate and deliver TLSRPT reports
created_at: 2023-10-25 11:32:11+00:00
pushed_at: 2025-02-23 23:42:16+00:00
url:
https://github.com/sys4/tlsrpt-reporter/blob/32d00c13508dd7f9695b77e253e88c88dc838fbd/tests/test_entrypoints.py

content:
#
#    Copyright (C) 2024-2025 sys4 AG
#    Author Boris Lohner bl@sys4.de
#
#    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 <http://www.gnu.org/licenses/>.
#

import unittest
from tlsrpt_reporter import plugins

class MyTestCase(unittest.TestCase):
    def plugin_entrypoint(self, group, url, expected_typename):
        """
        Test existance of an entrypoint to be loaded as storage back-end
        :param group: name of the group
        :param url: URL describing the storage back-end
        :param expected_typename: expected typename of the loaded storage back-end
        """
        cls = plugins.get_plugin(group, url)
        self.assertEqual(cls.__name__, expected_typename)

    def test_plugin_entrypoints(self):
        """
        Test existance and loadability of the defined plugin entrypoints
        """
        self.plugin_entrypoint("tlsrpt.collectd", "sqlite:///tmp/test-collectd.sqlite", "TLSRPTCollectdSQLite")
        self.plugin_entrypoint("tlsrpt.collectd", "dummy://", "DummyCollectd")
        self.plugin_entrypoint("tlsrpt.fetcher", "sqlite:///tmp/test-collectd.sqlite", "TLSRPTFetcherSQLite")


if __name__ == '__main__':
    unittest.main()


repository: jiqiu2022/Tenet-IDA9.0
file: tenet/integration/ida_loader.py
description: Tenet的修复版本，支持了ida9.0,增加了类似LLDB的调试逻辑，优化了界面，增加了arm64的支持，以及ALSR部分的重写
wanted_name: Tenet
comment: Trace Explorer
created_at: 2025-04-13 02:25:49+00:00
pushed_at: 2025-04-17 08:29:57+00:00
url:
https://github.com/jiqiu2022/Tenet-IDA9.0/blob/969f649772f123ba52319cacf382e2ad00054133/tenet/integration/ida_loader.py

content:
import time
import logging

import ida_idaapi
import ida_kernwin

from tenet.util.log import pmsg
from tenet.integration.ida_integration import TenetIDA

logger = logging.getLogger("Tenet.IDA.Loader")

#------------------------------------------------------------------------------
# IDA Plugin Loader
#------------------------------------------------------------------------------
#
#    This file contains a stub 'plugin' class for the plugin as required by
#    IDA Pro. Practically speaking, there should be little to *no* logic placed
#    in this file because it is disassembler-specific.
#
#    When IDA Pro is starting up, it will import all python files placed in its
#    root plugin folder. It will then attempt to call PLUGIN_ENTRY() on each of
#    the imported 'plugins'. We import PLUGIN_ENTRY into tenet_plugin.py
#    so that IDA can see it.
#
#    PLUGIN_ENTRY() is expected to return a plugin object (TenetIDAPlugin)
#    derived from ida_idaapi.plugin_t. IDA will register the plugin, and
#    interface with the plugin object to load / unload the plugin at certain
#    times, per its configuration (flags, hotkeys).
#
#    There should be virtually no reason for you to modify this file.
#

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

class TenetIDAPlugin(ida_idaapi.plugin_t):
    """
    The IDA plugin stub for Tenet.
    """

    #
    # Plugin flags:
    # - PLUGIN_MOD:  The plugin may modify the database
    # - PLUGIN_PROC: Load/unload the plugin when an IDB opens / closes
    # - PLUGIN_HIDE: Hide the plugin from the IDA plugin menu
    #

    flags = ida_idaapi.PLUGIN_PROC | ida_idaapi.PLUGIN_MOD | ida_idaapi.PLUGIN_HIDE
    comment = "Trace Explorer"
    help = ""
    wanted_name = "Tenet"
    wanted_hotkey = ""

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

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

        try:
            self.core = TenetIDA()
            self.core.load()
        except Exception as e:
            pmsg("Failed to initialize Tenet")
            logger.exception("Exception details:")

        #
        # we return PLUGIN_KEEP here regardless of success/failure. this is to
        # ensure that IDA will not try to reload the plugin again.
        #

        return ida_idaapi.PLUGIN_KEEP

    def run(self, arg):
        """
        This is called by IDA when this file is loaded as a script.
        """
        ida_kernwin.warning("Tenet cannot be run as a script in IDA.")

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

        start = time.time()
        logger.debug("-"*50)

        try:
            self.core.unload()
            self.core = None
        except Exception as e:
            logger.exception("Failed to cleanly unload Tenet from IDA.")

        end = time.time()
        logger.debug("-"*50)

        logger.debug("IDA term done... (%.3f seconds...)" % (end-start))


repository: userlandkernel/iOS-Firmware-Changes
file: iPhone11Pro/aarch64_pac.py
wanted_name: Aarch64 PAC
comment: ARM v8.3-A Pointer Authentication extension
created_at: 2018-12-15 14:08:29+00:00
pushed_at: 2020-04-30 13:50:27+00:00
url:
https://github.com/userlandkernel/iOS-Firmware-Changes/blob/d6329a39346a9dbac0c27b2dd90c379a1542290f/iPhone11Pro/aarch64_pac.py

content:
# AArch64 8.3-A Pointer Authentication extension
# Enhanced and Python version of xerub C++ plugin available there: https://github.com/xerub/idastuff/blob/master/arm64/aarch64_pac

# Copyright (c) 2018 Eloi Benoist-Vanderbeken - Synacktiv
# Copyright (c) 2018 xerub

# This program is free software; you can redistribute it and/or 
# modify it under the terms of the GNU General Public License version 
# 2 as published by the Free Software Foundation. 

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

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

import idaapi
import ida_hexrays

PAC_NONE = 0
PAC_PACIASP = PAC_NONE + 1
PAC_PACIBSP = PAC_PACIASP + 1
PAC_AUTIASP = PAC_PACIBSP + 1
PAC_AUTIBSP = PAC_AUTIASP + 1
PAC_PACIAZ = PAC_AUTIBSP + 1
PAC_PACIBZ = PAC_PACIAZ + 1
PAC_AUTIAZ = PAC_PACIBZ + 1
PAC_AUTIBZ = PAC_AUTIAZ + 1
PAC_PACIA1716 = PAC_AUTIBZ + 1
PAC_PACIB1716 = PAC_PACIA1716 + 1
PAC_AUTIA1716 = PAC_PACIB1716 + 1
PAC_AUTIB1716 = PAC_AUTIA1716 + 1
PAC_PACIA = PAC_AUTIB1716 + 1
PAC_PACIB = PAC_PACIA + 1
PAC_PACDA = PAC_PACIB + 1
PAC_PACDB = PAC_PACDA + 1
PAC_AUTIA = PAC_PACDB + 1
PAC_AUTIB = PAC_AUTIA + 1
PAC_AUTDA = PAC_AUTIB + 1
PAC_AUTDB = PAC_AUTDA + 1
PAC_PACIZA = PAC_AUTDB + 1
PAC_PACIZB = PAC_PACIZA + 1
PAC_PACDZA = PAC_PACIZB + 1
PAC_PACDZB = PAC_PACDZA + 1
PAC_AUTIZA = PAC_PACDZB + 1
PAC_AUTIZB = PAC_AUTIZA + 1
PAC_AUTDZA = PAC_AUTIZB + 1
PAC_AUTDZB = PAC_AUTDZA + 1
PAC_PACGA = PAC_AUTDZB + 1
PAC_XPACLRI = PAC_PACGA + 1
PAC_XPACI = PAC_XPACLRI + 1
PAC_XPACD = PAC_XPACI + 1
PAC_RETAA = PAC_XPACD + 1
PAC_RETAB = PAC_RETAA + 1
PAC_BRAA = PAC_RETAB + 1
PAC_BRAB = PAC_BRAA + 1
PAC_BRAAZ = PAC_BRAB + 1
PAC_BRABZ = PAC_BRAAZ + 1
PAC_BLRAA = PAC_BRABZ + 1
PAC_BLRAB = PAC_BLRAA + 1
PAC_BLRAAZ = PAC_BLRAB + 1
PAC_BLRABZ = PAC_BLRAAZ + 1
PAC_ERETAA = PAC_BLRABZ + 1
PAC_ERETAB = PAC_ERETAA + 1
PAC_LDRAA = PAC_ERETAB + 1
PAC_LDRAB = PAC_LDRAA + 1

OP_NAMES = {
	PAC_PACIASP: "PACIASP",
	PAC_PACIBSP: "PACIBSP",
	PAC_AUTIASP: "AUTIASP",
	PAC_AUTIBSP: "AUTIBSP",
	PAC_PACIAZ: "PACIAZ",
	PAC_PACIBZ: "PACIBZ",
	PAC_AUTIAZ: "AUTIAZ",
	PAC_AUTIBZ: "AUTIBZ",
	PAC_PACIA1716: "PACIA1716",
	PAC_PACIB1716: "PACIB1716",
	PAC_AUTIA1716: "AUTIA1716",
	PAC_AUTIB1716: "AUTIB1716",
	PAC_PACIA: "PACIA",
	PAC_PACIB: "PACIB",
	PAC_PACDA: "PACDA",
	PAC_PACDB: "PACDB",
	PAC_AUTIA: "AUTIA",
	PAC_AUTIB: "AUTIB",
	PAC_AUTDA: "AUTDA",
	PAC_AUTDB: "AUTDB",
	PAC_PACIZA: "PACIZA",
	PAC_PACIZB: "PACIZB",
	PAC_PACDZA: "PACDZA",
	PAC_PACDZB: "PACDZB",
	PAC_AUTIZA: "AUTIZA",
	PAC_AUTIZB: "AUTIZB",
	PAC_AUTDZA: "AUTDZA",
	PAC_AUTDZB: "AUTDZB",
	PAC_PACGA: "PACGA",
	PAC_XPACLRI: "XPACLRI",
	PAC_XPACI: "XPACI",
	PAC_XPACD: "XPACD",
	PAC_RETAA: "RETAA",
	PAC_RETAB: "RETAB",
	PAC_BRAA: "BRAA",
	PAC_BRAB: "BRAB",
	PAC_BRAAZ: "BRAAZ",
	PAC_BRABZ: "BRABZ",
	PAC_BLRAA: "BLRAA",
	PAC_BLRAB: "BLRAB",
	PAC_BLRAAZ: "BLRAAZ",
	PAC_BLRABZ: "BLRABZ",
	PAC_ERETAA: "ERETAA",
	PAC_ERETAB: "ERETAB",
	PAC_LDRAA: "LDRAA",
	PAC_LDRAB: "LDRAB"
}

def decode_PAC(d, insn):
	if (d & 0xffffc000) == 0xdac10000: 
		m = (d >> 10) & 7
		Z = (d >> 13) & 1
		Xn = (d >> 5) & 0x1F
		Xd = d & 0x1F
		if Z == 0: 
			insn.itype = idaapi.ARM_hlt
			insn.segpref = 14
			insn.Op1.type = idaapi.o_reg
			insn.Op1.reg = Xd + 129
			insn.Op1.dtype = idaapi.dt_qword
			insn.Op3.type = idaapi.o_reg
			insn.Op3.reg = Xd + 129
			insn.Op3.dtype = idaapi.dt_qword
			insn.Op3.flags = 0
			insn.Op2.type = idaapi.o_reg
			insn.Op2.reg = Xn + 129
			insn.Op2.dtype = idaapi.dt_qword
			insn.Op2.flags = idaapi.OF_SHOW
			insn.insnpref = PAC_PACIA + m
			insn.size = 4
			return True
		elif Xn == 31: 
			insn.itype = idaapi.ARM_hlt
			insn.segpref = 14
			insn.Op1.type = idaapi.o_reg
			insn.Op1.reg = Xd + 129
			insn.Op1.dtype = idaapi.dt_qword
			insn.Op3.type = idaapi.o_reg
			insn.Op3.reg = Xd + 129
			insn.Op3.dtype = idaapi.dt_qword
			insn.Op3.flags = 0
			insn.insnpref = PAC_PACIZA + m
			insn.size = 4
			return True
	if (d & 0xfffffd1f) == 0xd503211f: 
		m = (d >> 6) & 3
		CRm = (d >> 9) & 1
		op2 = (d >> 5) & 1
		if CRm == 0: 
			insn.itype = idaapi.ARM_hlt
			insn.segpref = 14
			insn.Op1.type = idaapi.o_void
			insn.insnpref = PAC_PACIA1716 + m
		elif op2: 
			insn.itype = idaapi.ARM_hlt
			insn.segpref = 14
			insn.Op1.type = idaapi.o_void
			insn.insnpref = PAC_PACIASP + m
		else:
			insn.itype = idaapi.ARM_hlt
			insn.segpref = 14
			insn.Op1.type = idaapi.o_void
			insn.insnpref = PAC_PACIAZ + m
		insn.size = 4
		return True
	if (d & 0xffe0fc00) == 0x9ac03000: 
		Xm = (d >> 16) & 0x1F
		Xn = (d >> 5) & 0x1F
		Xd = d & 0x1F
		insn.itype = idaapi.ARM_hlt
		insn.segpref = 14
		insn.Op1.type = idaapi.o_reg
		insn.Op1.reg = Xd + 129
		insn.Op1.dtype = idaapi.dt_qword
		insn.Op2.type = idaapi.o_reg
		insn.Op2.reg = Xn + 129
		insn.Op2.dtype = idaapi.dt_qword
		insn.Op3.type = idaapi.o_reg
		insn.Op3.reg = Xm + 129
		insn.Op3.dtype = idaapi.dt_qword
		insn.insnpref = PAC_PACGA
		insn.size = 4
		return True
	if (d & 0xfffffbe0) == 0xdac143e0: 
		D = (d >> 10) & 1
		Xd = d & 0x1F
		insn.itype = idaapi.ARM_hlt
		insn.segpref = 14
		insn.Op1.type = idaapi.o_reg
		insn.Op1.reg = Xd + 129
		insn.Op1.dtype = idaapi.dt_qword
		insn.Op3.type = idaapi.o_reg
		insn.Op3.reg = Xd + 129
		insn.Op3.dtype = idaapi.dt_qword
		insn.Op3.flags = 0
		insn.insnpref = PAC_XPACI + D
		insn.size = 4
		return True
	if d == 0xd50320ff: 
		insn.itype = idaapi.ARM_hlt
		insn.segpref = 14
		insn.Op1.type = idaapi.o_void
		insn.insnpref = PAC_XPACLRI
		insn.size = 4
		return True
	if (d & 0xfffffbff) == 0xd65f0bff: 
		M = (d >> 10) & 1
		insn.insnpref = PAC_RETAA + M
		insn.itype = idaapi.ARM_ret
		insn.segpref = 14
		insn.Op1.type = idaapi.o_reg
		insn.Op1.reg = 30 + 129
		insn.Op1.dtype = idaapi.dt_qword
		insn.Op1.flags = 0
		insn.size = 4
		return True
	if (d & 0xfedff800) == 0xd61f0800: 
		is_blr = (d >> 19) & 4
		Z = (d >> 24) & 1
		M = (d >> 10) & 1
		Xn = (d >> 5) & 0x1F
		Xm = d & 0x1F
		if Z == 0 and Xm == 31: 
			insn.itype =  idaapi.ARM_blr if is_blr else idaapi.ARM_br
			insn.segpref = 14
			insn.Op1.type = idaapi.o_reg
			insn.Op1.reg = Xn + 129
			insn.Op1.dtype = idaapi.dt_qword
			insn.insnpref = PAC_BRAAZ + M + is_blr
			insn.size = 4
			return True
		elif Z: 
			insn.itype = idaapi.ARM_blr if is_blr else idaapi.ARM_br
			insn.segpref = 14
			insn.Op1.type = idaapi.o_reg
			insn.Op1.reg = Xn + 129
			insn.Op1.dtype = idaapi.dt_qword
			insn.Op2.type = idaapi.o_reg
			insn.Op2.reg = Xm + 129
			insn.Op2.dtype = idaapi.dt_qword
			insn.Op2.flags = idaapi.OF_SHOW
			insn.insnpref = PAC_BRAA + M + is_blr
			insn.size = 4
			return True
	if (d & 0xfffffbff) == 0xd69f0bff: 
		M = (d >> 10) & 1
		insn.insnpref = PAC_ERETAA + M
		insn.itype = idaapi.ARM_eret
		insn.segpref = 14
		insn.size = 4
		return True
	if (d & 0xff200400) == 0xf8200400: 
		M = (d >> 23) & 1
		imm10 = ((d & 0x400000) << 9) | ((d & 0x1ff000) << 10)
		offset = imm10 >> 19
		W = (d >> 11) & 1
		Xn = (d >> 5) & 0x1F
		Xt = d & 0x1F
		insn.itype = idaapi.ARM_ldr
		insn.segpref = 14
		insn.Op1.type = idaapi.o_reg
		insn.Op1.reg = Xt + 129
		insn.Op1.dtype = idaapi.dt_qword
		insn.Op2.type = idaapi.o_displ
		insn.Op2.reg = Xn + 129
		insn.Op2.dtype = idaapi.dt_qword
		insn.Op2.addr = offset
		if W != 0: 
			insn.auxpref = 0x20
		insn.insnpref = PAC_LDRAA + M
		insn.size = 4
		return True
	return False

class Aarch64PACHook(idaapi.IDP_Hooks):
	CUSTOM_INSTRUCTIONS = {idaapi.ARM_hlt, idaapi.ARM_ret, idaapi.ARM_blr, idaapi.ARM_br, idaapi.ARM_eret,idaapi.ARM_ldr}
	INDENT = 16
	def ev_ana_insn(self, outctx):
		return outctx.size if decode_PAC(idaapi.get_dword(outctx.ea), outctx) else 0

	def ev_emu_insn(self, insn):
		if insn.itype != idaapi.ARM_brk:
			return False
		return True

	def ev_out_mnem(self, outctx):
		if outctx.insn.itype in self.CUSTOM_INSTRUCTIONS:
			mnem = OP_NAMES.get(ord(outctx.insn.insnpref), None)
			if mnem is not None:
				outctx.out_custom_mnem(mnem, self.INDENT)
				return 1
		return 0

class Aarch64PACPlugin(idaapi.plugin_t):
	flags = idaapi.PLUGIN_PROC | idaapi.PLUGIN_HIDE
	comment = "ARM v8.3-A Pointer Authentication extension"
	wanted_hotkey = ""
	help = "Runs transparently"
	wanted_name = "Aarch64 PAC"
	hook = None
	enabled = 1

	def init(self):
		if idaapi.ph_get_id() != idaapi.PLFM_ARM or idaapi.BADADDR <= 0xFFFFFFFF:
			return idaapi.PLUGIN_SKIP
		print "%s init"%self.comment
		self.hook = Aarch64PACHook()
		self.hook.hook()
		return idaapi.PLUGIN_KEEP

	def run():
		pass

	def term(self):
		if self.hook is not None:
			self.hook.unhook()
		print "%s unloaded"%self.comment

def PLUGIN_ENTRY():
	return Aarch64PACPlugin()


repository: BraveCattle/ctf
file: tools/ida_pro/IDA_Pro_7.7/plugins/BinCAT/python/idabincat/bcplugin.py
wanted_name: BinCAT
comment: Interface to the BinCAT analyzer
created_at: 2024-01-23 08:17:51+00:00
pushed_at: 2024-10-23 09:31:30+00:00
url:
https://github.com/BraveCattle/ctf/blob/5858d0da5c048eb8468ff8be1d63fbc50d61b601/tools/ida_pro/IDA_Pro_7.7/plugins/BinCAT/python/idabincat/bcplugin.py

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

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

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

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

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

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

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

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

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


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

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

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


class AnalyzerUnavailable(Exception):
    pass


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

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

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

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

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

        bc_log.info("Autostarting")

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

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

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


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

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

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

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

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

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

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

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

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

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

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


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

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

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

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

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

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

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

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

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

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

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

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


class WebAnalyzer(Analyzer):
    API_VERSION = "1.2"

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

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

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

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

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

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

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

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

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

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

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


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

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

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

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

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

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

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

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

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

        self.gui.after_change_ea()

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

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

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

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

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

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

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

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

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

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

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


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

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

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

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

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

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

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

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

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


def PLUGIN_ENTRY():
    return BincatPlugin()


