# plugins

repository: hyp3ri0n-ng/DriverFuckBuddy
file: DriverBuddy.py
description: DriverBuddy for IDA 7.5. HG-style.
wanted_name: Driver Buddy
created_at: 2021-01-04 04:20:39+00:00
pushed_at: 2021-01-04 05:25:17+00:00
url:
https://github.com/hyp3ri0n-ng/DriverFuckBuddy/blob/b10007bc128f84982e147a0069e953f48b3f3016/DriverBuddy.py

content:
from idaapi import *
from idautils import *
from idc import *
from DriverBuddy import data
from DriverBuddy import ioctl
'''#######################################################################################

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

Written by Braden Hollembaek and Adam Pond of NCC Group
#######################################################################################'''

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

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

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

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

    def term(self):
        pass

def PLUGIN_ENTRY():
    return DriverBuddyPlugin()


repository: BinaryAnalysisPlatform/bap-ida-python
file: plugins/bap/plugins/bap_trace.py
description: integration with IDA
wanted_name: BAP: Load Observations
comment: Load Primus Observations
created_at: 2016-06-10 18:00:33+00:00
pushed_at: 2021-05-27 15:13:33+00:00
url:
https://github.com/BinaryAnalysisPlatform/bap-ida-python/blob/d8d4679de2f50bb75f556419565821d95404034e/plugins/bap/plugins/bap_trace.py

content:
# noqa: ignore=F811

import idaapi

import os


from bap.utils import trace
from PyQt5 import QtWidgets
from PyQt5 import QtGui

from PyQt5.QtCore import (
    Qt,
    QFile,
    QIODevice,
    QCryptographicHash as QCrypto,
    QRegExp,
    QTimer,
    QAbstractItemModel,
    QModelIndex,
    QVariant,
    pyqtSignal,
    QSortFilterProxyModel)


def add_insn_to_trace_view(ea, tid=1):
    idaapi.dbg_add_tev(1, tid, ea)


@trace.handler('pc-changed', requires=['machine-id', 'pc'])
def tev_insn(state, ev):
    "stores each visited instruction to the IDA Trace Window"
    add_insn_to_trace_view(state['pc'], tid=state['machine-id'])


@trace.handler('pc-changed', requires=['pc'])
def tev_insn0(state, ev):
    """stores each visited instruction to the IDA Trace Window.

    But doesn't set the pid/tid field, and keep it equal to 0
    (This enables interoperation with the debugger)
    """
    add_insn_to_trace_view(state['pc'])


@trace.handler('call', requires=['machine-id', 'pc'])
def tev_call(state, call):
    "stores call events to the IDA Trace Window"
    caller = state['pc']
    callee = idaapi.get_name_ea(0, call[0])
    idaapi.dbg_add_call_tev(state['machine-id'], caller, callee)


incidents = []
locations = {}


@trace.handler('incident')
def incident(state, data):
    incidents.append(Incident(data[0], [int(x) for x in data[1:]]))


@trace.handler('incident-location')
def incident_location(state, data):
    id = int(data[0])
    locations[id] = [parse_point(p) for p in data[1]]


# we are using PyQt5 here, because IDAPython relies on a system
# openssl 0.9.8 which is quite outdated and not available on most
# modern installations
def md5sum(filename):
    """computes md5sum of a file with the given ``filename``

    The return value is a 32 byte hexadecimal ASCII representation of
    the md5 sum (same value as returned by the ``md5sum filename`` command)
    """
    stream = QFile(filename)
    if not stream.open(QIODevice.ReadOnly | QIODevice.Text):
        raise IOError("Can't open file: " + filename)
    hasher = QCrypto(QCrypto.Md5)
    if not hasher.addData(stream):
        raise ValueError('Unable to hash file: ' + filename)
    stream.close()
    return str(hasher.result().toHex())


class HandlerSelector(QtWidgets.QGroupBox):
    def __init__(self, parent=None):
        super(HandlerSelector, self).__init__("Trace Event Processors", parent)
        self.setFlat(True)
        box = QtWidgets.QVBoxLayout(self)
        self.options = {}
        for name in trace.handlers:
            btn = QtWidgets.QCheckBox(name)
            btn.setToolTip(trace.handlers[name].__doc__)
            box.addWidget(btn)
            self.options[name] = btn
        box.addStretch(1)
        self.setCheckable(True)
        self.setChecked(True)
        self.setLayout(box)


class MachineSelector(QtWidgets.QWidget):
    def __init__(self, parent=None):
        super(MachineSelector, self).__init__(parent)
        box = QtWidgets.QHBoxLayout(self)
        label = MonitoringLabel('List of &machines (threads)')
        self.is_ready = label.is_ready
        self.updated = label.updated
        box.addWidget(label)
        self._machines = QtWidgets.QLineEdit('all')
        self._machines.setToolTip('an integer, \
        a comma-separated list of integers, or "all"')
        grammar = QRegExp(r'\s*(all|\d+\s*(,\s*\d+\s*)*)\s*')
        valid = QtGui.QRegExpValidator(grammar)
        self._machines.setValidator(valid)
        label.setBuddy(self._machines)
        box.addWidget(self._machines)
        box.addStretch(1)
        self.setLayout(box)

    def selected(self):
        if not self._machines.hasAcceptableInput():
            raise ValueError('invalid input')
        data = self._machines.text().strip()
        if data == 'all':
            return None
        else:
            return [int(x) for x in data.split(',')]


class MonitoringLabel(QtWidgets.QLabel):
    "a label that will monitors the validity of its buddy"

    updated = pyqtSignal()

    def __init__(self, text='', buddy=None, parent=None):
        super(MonitoringLabel, self).__init__(parent)
        self.setText(text)
        if buddy:
            self.setBuddy(buddy)

    def setText(self, text):
        super(MonitoringLabel, self).setText(text)
        self.text = text

    def setBuddy(self, buddy):
        super(MonitoringLabel, self).setBuddy(buddy)
        buddy.textChanged.connect(lambda x: self.update())
        self.update()

    def is_ready(self):
        return not self.buddy() or self.buddy().hasAcceptableInput()

    def update(self):
        self.updated.emit()
        if self.is_ready():
            super(MonitoringLabel, self).setText(self.text)
        else:
            super(MonitoringLabel, self).setText(
                '<font color=red>'+self.text+'</font>')


class ExistingFileValidator(QtGui.QValidator):
    def __init__(self, parent=None):
        super(ExistingFileValidator, self).__init__(parent)

    def validate(self, name, pos):
        if os.path.isfile(name):
            return (self.Acceptable, name, pos)
        else:
            return (self.Intermediate, name, pos)


class TraceFileSelector(QtWidgets.QWidget):

    def __init__(self, parent=None):
        super(TraceFileSelector, self).__init__(parent)
        box = QtWidgets.QHBoxLayout(self)
        label = MonitoringLabel('Trace &file:')
        self.is_ready = label.is_ready
        self.updated = label.updated
        box.addWidget(label)
        self.location = QtWidgets.QLineEdit('incidents')
        self.text = self.location.text
        must_exist = ExistingFileValidator()
        self.location.setValidator(must_exist)
        label.setBuddy(self.location)
        box.addWidget(self.location)
        openfile = QtWidgets.QPushButton(self)
        openfile.setIcon(self.style().standardIcon(
            QtWidgets.QStyle.SP_DialogOpenButton))
        dialog = QtWidgets.QFileDialog(self)
        openfile.clicked.connect(dialog.open)
        dialog.fileSelected.connect(self.location.setText)
        box.addWidget(openfile)
        box.addStretch(1)
        self.setLayout(box)


class IncidentView(QtWidgets.QWidget):
    def __init__(self, parent=None):
        super(IncidentView, self).__init__(parent)
        self.view = QtWidgets.QTreeView()
        self.view.setAllColumnsShowFocus(True)
        self.view.setUniformRowHeights(True)
        box = QtWidgets.QVBoxLayout()
        box.addWidget(self.view)
        self.load_trace = QtWidgets.QPushButton('&Trace')
        self.load_trace.setToolTip('Load into the Trace Window')
        self.load_trace.setEnabled(False)
        for activation_signal in [
                self.view.activated,
                self.view.entered,
                self.view.pressed]:
            activation_signal.connect(lambda _: self.update_controls_state())
        self.load_trace.clicked.connect(self.load_current_trace)
        self.view.doubleClicked.connect(self.jump_to_index)
        hbox = QtWidgets.QHBoxLayout()
        self.filter = QtWidgets.QLineEdit()
        self.filter.textChanged.connect(self.filter_model)
        filter_label = QtWidgets.QLabel('&Search')
        filter_label.setBuddy(self.filter)
        hbox.addWidget(filter_label)
        hbox.addWidget(self.filter)
        hbox.addWidget(self.load_trace)
        box.addLayout(hbox)
        self.setLayout(box)
        self.model = None
        self.proxy = None

    def display(self, incidents, locations):
        self.model = IncidentModel(incidents, locations, self)
        self.proxy = QSortFilterProxyModel(self)
        self.proxy.setSourceModel(self.model)
        self.proxy.setFilterRole(self.model.filter_role)
        self.proxy.setFilterRegExp(QRegExp(self.filter.text()))
        self.view.setModel(self.proxy)

    def filter_model(self, txt):
        if self.proxy:
            self.proxy.setFilterRegExp(QRegExp(txt))

    def update_controls_state(self):
        curr = self.view.currentIndex()
        self.load_trace.setEnabled(curr.isValid() and
                                   curr.parent().isValid())

    def load_current_trace(self):
        idx = self.proxy.mapToSource(self.view.currentIndex())
        if not idx.isValid() or index_level(idx) not in (1, 2):
            raise ValueError('load_current_trace: invalid index')

        if index_level(idx) == 2:
            idx = idx.parent()

        incident = self.model.incidents[idx.parent().row()]
        location = incident.locations[idx.row()]
        backtrace = self.model.locations[location]

        for p in reversed(backtrace):
            self.load_trace_point(p)

    def jump_to_index(self, idx):
        idx = self.proxy.mapToSource(idx)
        if index_level(idx) != 2:
            # don't mess with parents, they are used to create children
            return
        grandpa = idx.parent().parent()
        incident = self.model.incidents[grandpa.row()]
        location = incident.locations[idx.parent().row()]
        trace = self.model.locations[location]
        point = trace[idx.row()]
        self.show_trace_point(point)

    def load_trace_point(self, p):
        add_insn_to_trace_view(p.addr)

    def show_trace_point(self, p):
        idaapi.jumpto(p.addr)


class TraceLoaderController(QtWidgets.QWidget):
    finished = pyqtSignal()

    def __init__(self, parent=None):
        super(TraceLoaderController, self).__init__(parent)
        self.loader = None
        box = QtWidgets.QVBoxLayout(self)
        self.location = TraceFileSelector(self)
        self.handlers = HandlerSelector(self)
        self.machines = MachineSelector(self)
        box.addWidget(self.location)
        box.addWidget(self.handlers)
        box.addWidget(self.machines)
        self.load = QtWidgets.QPushButton('&Load')
        self.load.setDefault(True)
        self.load.setEnabled(self.location.is_ready())
        self.cancel = QtWidgets.QPushButton('&Stop')
        self.cancel.setVisible(False)
        hor = QtWidgets.QHBoxLayout()
        hor.addWidget(self.load)
        hor.addWidget(self.cancel)
        self.progress = QtWidgets.QProgressBar()
        self.progress.setVisible(False)
        hor.addWidget(self.progress)
        hor.addStretch(2)
        box.addLayout(hor)

        def enable_load():
            self.load.setEnabled(self.location.is_ready() and
                                 self.machines.is_ready())
        self.location.updated.connect(enable_load)
        self.machines.updated.connect(enable_load)
        enable_load()
        self.processor = QTimer()
        self.processor.timeout.connect(self.process)
        self.load.clicked.connect(self.processor.start)
        self.cancel.clicked.connect(self.stop)
        self.setLayout(box)

    def start(self):
        self.cancel.setVisible(True)
        self.load.setVisible(False)
        filename = self.location.text()
        self.loader = trace.Loader(file(filename))
        self.progress.setVisible(True)
        stat = os.stat(filename)
        self.progress.setRange(0, stat.st_size)
        machines = self.machines.selected()
        if machines is not None:
            self.loader.enable_filter('filter-machine', id=machines)

        for name in self.handlers.options:
            if self.handlers.options[name].isChecked():
                self.loader.enable_handlers([name])

    def stop(self):
        self.processor.stop()
        self.progress.setVisible(False)
        self.cancel.setVisible(False)
        self.load.setVisible(True)
        self.loader = None
        self.finished.emit()

    def process(self):
        if not self.loader:
            self.start()
        try:
            self.loader.next()
            self.progress.setValue(self.loader.parser.lexer.instream.tell())
        except StopIteration:
            self.stop()


def index_level(idx):
    if idx.isValid():
        return 1 + index_level(idx.parent())
    else:
        return -1


def index_up(idx, level=0):
    if level == 0:
        return idx
    else:
        return index_up(idx.parent(), level=level-1)


class IncidentIndex(object):
    def __init__(self, model, index):
        self.model = model
        self.index = index

    @property
    def incidents(self):
        return self.model.incidents

    @property
    def level(self):
        return index_level(self.index)

    @property
    def column(self):
        return self.index.column()

    @property
    def row(self):
        return self.index.row()

    @property
    def incident(self):
        top = index_up(self.index, self.level)
        return self.incidents[top.row()]

    @property
    def location(self):
        if self.level in (1, 2):
            top = self.index
            if self.level == 2:
                top = index_up(self.index, 1)
            location_id = self.incident.locations[top.row()]
            if self.model.locations is None:
                return None
            else:
                return self.model.locations.get(location_id)

    @property
    def point(self):
        if self.level == 2:
            return self.location[self.index.row()]


class IncidentModel(QAbstractItemModel):
    filter_role = Qt.UserRole
    sort_role = Qt.UserRole + 1

    handlers = []

    def __init__(self, incidents, locations, parent=None):
        super(IncidentModel, self).__init__(parent)
        self.incidents = incidents
        self.locations = locations
        self.parents = {0: QModelIndex()}
        self.child_ids = 0

    def dispatch(self, role, index):
        for handler in self.handlers:
            def sat(c, v):
                if c == 'roles':
                    return role in v
                if c == 'level':
                    return index.level == v
                if c == 'column':
                    return index.column == v

            for (c, v) in handler['constraints'].items():
                if not sat(c, v):
                    break
            else:
                return handler['accept'](index)

    def index(self, row, col, parent):
        if parent.isValid():
            self.child_ids += 1
            index = self.createIndex(row, col, self.child_ids)
            self.parents[self.child_ids] = parent
            return index
        else:
            return self.createIndex(row, col, 0)

    def parent(self, child):
        return self.parents[child.internalId()]

    def rowCount(self, parent):
        n = self.dispatch('row-count', IncidentIndex(self, parent))
        return 0 if n is None else n

    def columnCount(self, parent):
        return 2 if not parent.isValid() or parent.column() == 0 else 0

    def data(self, index, role):
        role = {
            Qt.DisplayRole: 'display',
            self.sort_role: 'sort',
            self.filter_role: 'filter'
        }.get(role)

        if role:
            return QVariant(self.dispatch(role, IncidentIndex(self, index)))
        else:
            return QVariant()


def defmethod(*args, **kwargs):
    def register(method):
        kwargs['roles'] = args
        IncidentModel.handlers.append({
            'name': method.__name__,
            'constraints': kwargs,
            'accept': method})
    return register


@defmethod('display', level=2, column=0)
def display_point(msg):
    return '{:x}'.format(msg.point.addr)


@defmethod('display', level=2, column=1)
def display_point_machine(msg):
    return msg.point.machine


@defmethod('display', level=1, column=0)
def display_incident_location(msg):
    return 'location-{}'.format(msg.row)


@defmethod('display', level=0, column=0)
def display_incident_name(msg):
    return msg.incident.name


@defmethod('display', level=0, column=1)
def display_incident_id(msg):
    return msg.row


@defmethod('sort', 'filter', column=0)
def incident_name(msg):
    return msg.incident.name


@defmethod('row-count', level=-1)
def number_of_incidents(msg):
    return len(msg.incidents)


@defmethod('row-count', level=0, column=0)
def number_of_locations(msg):
    return len(msg.incident.locations)


@defmethod('row-count', level=1, column=0)
def backtrace_length(msg):
    return 0 if msg.location is None else len(msg.location)


class Incident(object):
    __slots__ = ['name', 'locations']

    def __init__(self, name, locations):
        self.name = name
        self.locations = locations

    def __repr__(self):
        return 'Incident({}, {})'.format(repr(self.name),
                                         repr(self.locations))


class Point(object):
    __slots__ = ['addr', 'machine']

    def __init__(self, addr, machine=None):
        self.addr = addr
        self.machine = machine

    def __str__(self):
        if self.machine:
            return '{}:{}'.format(self.machine, self.addr)
        else:
            return str(self.addr)

    def __repr__(self):
        if self.machine:
            return 'Point({},{})'.format(self.machine, self.addr)
        else:
            return 'Point({})'.format(repr(self.addr))


def parse_point(data):
    parts = data.split(':')
    if len(parts) == 1:
        return Point(int(data, 16))
    else:
        return Point(int(parts[1], 16), int(parts[0]))


class BapTraceMain(idaapi.PluginForm):
    def OnCreate(self, form):
        form = self.FormToPyQtWidget(form)
        self.control = TraceLoaderController(form)
        self.incidents = IncidentView(form)

        def display():
            self.incidents.display(incidents, locations)
        self.control.finished.connect(display)
        box = QtWidgets.QHBoxLayout()
        split = QtWidgets.QSplitter()
        split.addWidget(self.control)
        split.addWidget(self.incidents)
        box.addWidget(split)
        form.setLayout(box)


class BapTracePlugin(idaapi.plugin_t):
    wanted_name = 'BAP: Load Observations'
    wanted_hotkey = ''
    flags =  idaapi.PLUGIN_FIX
    comment = 'Load Primus Observations'
    help = """
    Loads Primus Observations into IDA for further analysis
    """

    def __init__(self):
        self.form = None
        self.name = 'Primus Observations'

    def init(self):
        return idaapi.PLUGIN_KEEP

    def term(self):
        pass

    def run(self, arg):
        if not self.form:
            self.form = BapTraceMain()
        return self.form.Show(self.name, options=(
            self.form.FORM_PERSIST |
            self.form.FORM_SAVE))


def PLUGIN_ENTRY():
    return BapTracePlugin()


main = None


def bap_trace_test():
    global main
    main = BapTraceMain()
    main.Show('Primus Observations')


repository: mahmoudimus/.idapro
file: scripts/blzd/topaqueminator.py
created_at: 2025-02-16 21:33:20+00:00
pushed_at: 2025-04-09 18:15:30+00:00
url:
https://github.com/mahmoudimus/.idapro/blob/e0bf0d58fb6de4eb6af0b87ffbd869d8c7deef08/scripts/blzd/topaqueminator.py

content:
import idaapi, idautils, idc, ida_funcs, ida_auto
import sys
import os

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 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 = False
CONTINUOS_ANALYZING = False
TESTS_MODE = False


SINGLE_BLOCK_MODE = False
WHOLE_FUNCTION_MODE = True
FIX_DEFLOW_MODE = False


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):
    insn = idaapi.insn_t()
    idaapi.decode_insn(insn, ea)
    dtype = insn.ops[op_idx].dtype
    dtyp_size = idaapi.get_dtype_size(dtype)
    return dtype, dtyp_size


def set_comment(ea, comment):
    idc.set_cmt(ea, comment, 0)


def get_comment(ea):
    return idc.get_cmt(ea, 0)


def read_range_selection():
    return idaapi.read_range_selection(None)


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


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

        
        is_be = ida_ida.cvar.inf.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
        }
        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) if patch_data else 0

        (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
                fmted_asm = "\n  ".join(orig_asm if orig_asm else "")
                new_patch_comment = ""
                if orig_comment != "":
                    new_patch_comment = "\n"
                new_patch_comment += f"Keypatch modified this from:\n  {fmted_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)
        if not inst:
            return False
        inst_str = inst.get_canon_mnem()
        if "jmp" in inst_str:
            return True
        return False

    def check_conditional_x86_64(self, inst_ea):
        inst = idautils.DecodeInstruction(inst_ea)
        if not inst:
            return False
        inst_str = inst.get_canon_mnem()
        if "j" in inst_str and "jmp" not in inst_str:
            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)
            if not code_addr:
                return False
            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:
                jumploc = self.get_the_location_of_the_jump(opaque_addr)
                if not jumploc:
                    print("unknown jump location!")
                    return
                jmp_location_str = hex(jumploc)
            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)
            inst = idautils.DecodeInstruction(opaque_addr)
            if not inst:
                return False, None
            init_inst_string = inst.get_canon_mnem()
            inst = idautils.DecodeInstruction(next_jmp_loc)
            if not inst:
                return False, None
            jmp_inst_string = inst.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 ")
                inst = idautils.DecodeInstruction(next_jmp_loc)
                if not inst:
                    logging.debug(f"cannot decode instruction at {next_jmp_loc}")
                    continue
                jmp_inst_string = inst.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 and final_jmp_loc is not None:
            inst = idautils.DecodeInstruction(opaque_addr)
            if not inst:
                return False, None
            inst_cmd_str = inst.get_canon_mnem()
            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 Exception("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 Exception("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)
            if not encoding:
                continue
            opcodes = b""
            for i in encoding:
                opcodes += bytes([i])
            all_shell_code += opcodes
        return all_shell_code

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

if __name__ == "__main__":
    run_main_class()

# """------------------------------"""


# 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(f"Error: {err}\n{traceback.format_exc()}")
#         raise


# """------------------------------"""


repository: mak/ipad
file: ipad.py
description: Yet another IDA Pro real time syncing plugin
url: https://github.com/mak/ipad/blob/fed5df9b48bcb94b09b93fc0169448abaaad9dd4/ipad.py
wanted_name: Changer Plugin
comment: comment
created_at: 2018-03-27 23:23:53+00:00
pushed_at: 2018-03-27 23:24:12+00:00
content:
import sys,os
import idc,idaapi

## my imports
from ipad.config import Config
from ipad.changer import Changer
from ipad.ui import UI
from ipad.cmd import Commands
from ipad.controler import IdaAction
## other things
import os
        
class changer(idaapi.plugin_t):
    flags = idaapi.PLUGIN_FIX | idaapi.PLUGIN_DRAW
    
    comment = "comment"
    help = "help"
    wanted_name = "Changer Plugin"
    wanted_hotkey = "alt-shift-p"

    @property
    def have_idb(self):
        return bool(idc.GetIdbPath())
    
    def init(self):
        self.ui = UI(self)
        self.load_config()
        return idaapi.PLUGIN_KEEP

    def load_config(self):

        if hasattr(self,'cfg'):
            del self.cfg
                
        self.cfg = Config()
        if not os.path.exists(self.cfg.dirname()):
            os.mkdir(self.cfg.dirname())
        self.cfg.load_cfg()
        
    def open_ui(self):
        self.ui.Show('ipad')

    def store_idb(self):
        if self.have_idb:
            idc.SaveBase(idc.GetIdbPath())
        self.cc.store_idb()
        
    def load_idb(self,*a):
        self.cc.load_idb(*a)

    def dispatch(self,d):
        self.a.dispatch(d)
        
    def run(self,arg):
        
        print '[*] ipad started - uid: %s | db: %s' % (self.cfg.uid,self.cfg.db)
                    
        if not self.have_idb:
            self.cc = Commands(self,None)#self.t)
            self.ui.storage.setCtrl(self)
            self.ui.storage._populate(self.cc.get_storage())
            self.ui.Show('ipad',UI.FORM_CENTERED|UI.FORM_MENU)
        else:            
            self.a = IdaAction(self.cfg,self.ui.history)
            self.t = Changer(self.a)
            self.cc = Commands(self.a,self.t)
            
            if not self.ui.all_active():
                self.ui.Show('ipad')
                
            if self.cc.get_changes():
                self.ui.sessions._populate(self.cc.get_sessions())
                
            self.ui.sessions.setCtrl(self.a)
            self.t.start()
            self.a.start()
            
    def term(self):
        print '[+] going down'

        if hasattr(self,'a'):
            self.a.end()
        if hasattr(self,'t'):
            self.t.quit()
        
    
def PLUGIN_ENTRY():
  return changer()


# def loadx():
#     idaapi.load_and_run_plugin(__file__,0)
    
# if __name__ == '__main__':
#     print sys.argv
#     idaapi.autoWait()


repository: idapython/src
file: examples/ui/auto_instantiate_widget_plugin.py
description: IDAPython project for Hex-Ray's IDA Pro
wanted_name: autoinst
created_at: 2015-03-14 20:09:27+00:00
pushed_at: 2025-03-13 15:18:11+00:00
url:
https://github.com/idapython/src/blob/af978e79e68af701c45bab401afce5145aa4b80b/examples/ui/auto_instantiate_widget_plugin.py

content:
"""
summary: restore custom widgets across sessions

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

level: advanced
"""

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: Mondego/pyreco
file: repoData/devttys0-ida/allPythonContent.py
description: Type discovery for Python
wanted_name: Define all data and code
created_at: 2015-11-09 18:50:20+00:00
pushed_at: 2016-06-27 05:15:24+00:00
url:
https://github.com/Mondego/pyreco/blob/0ac6653219c2701c13c508c5c4fc9bc3437eea06/repoData/devttys0-ida/allPythonContent.py

content:
__FILENAME__ = codatify
# IDA plugin that converts all data in data segments to defined data types, and all data in code segments to code.
#
# Use by going to Options->Define data and code, or use the Alt+3 hotkey.
#
# Craig Heffner
# Tactical Network Solutions

import idc
import idaapi
import idautils

class Codatify(object):

    CODE = 2
    DATA = 3
    SEARCH_DEPTH = 25

    def __init__(self):
        pass

    # Get the start of the specified segment type (2 == code, 3 == data)
    def get_start_ea(self, attr):
        ea = idc.BADADDR
        seg = idc.FirstSeg()

        while seg != idc.BADADDR:
            if idc.GetSegmentAttr(seg, idc.SEGATTR_TYPE) == attr:
                ea = seg
                break
            else:
                seg = idc.NextSeg(seg)
    
        return ea

    # Creates ASCII strings
    def stringify(self):
        n = 0
        ea = self.get_start_ea(self.DATA)

        if ea == idc.BADADDR:
            ea = idc.FirstSeg()

        print ("Looking for possible strings starting at: %s:0x%X..." % (idc.SegName(ea), ea)),

        for s in idautils.Strings():
            if s.ea > ea:
                if not idc.isASCII(idc.GetFlags(s.ea)) and idc.MakeStr(s.ea, idc.BADADDR):
                    n += 1

        print "created %d new ASCII strings" % n

    # Converts remaining data into DWORDS.
    def datify(self):
        ea = self.get_start_ea(self.DATA)
        if ea == idc.BADADDR:
            ea = idc.FirstSeg()
        
        print "Converting remaining data to DWORDs...",
    
        while ea != idc.BADADDR:
            flags = idc.GetFlags(ea)
                    
            if idc.isUnknown(flags) or idc.isByte(flags):
                idc.MakeDword(ea)
                idc.OpOff(ea, 0, 0)

            ea = idc.NextAddr(ea)

        print "done."

    def pointify(self):
        counter = 0

        print "Renaming pointers...",

        for (name_ea, name) in idautils.Names():
            for xref in idautils.XrefsTo(name_ea):
                xref_name = idc.Name(xref.frm)
                if xref_name and xref_name.startswith("off_"):
                    i = 0
                    new_name = name + "_ptr"
                    while idc.LocByName(new_name) != idc.BADADDR:
                        new_name = name + "_ptr%d" % i
                        i += 1

                    if idc.MakeName(xref.frm, new_name):
                        counter += 1
                    #else:
                    #    print "Failed to create name '%s'!" % new_name

        print "renamed %d pointers" % counter

    # Creates functions and code blocks
    def codeify(self, ea=idc.BADADDR):
        func_count = 0
        code_count = 0

        if ea == idc.BADADDR:
            ea = self.get_start_ea(self.CODE)
            if ea == idc.BADADDR:
                ea = idc.FirstSeg()

        print "\nLooking for undefined code starting at: %s:0x%X" % (idc.SegName(ea), ea)

        if self.get_start_ea(self.DATA) == idc.BADADDR:
            print "WARNING: No data segments defined! I don't know where the code segment ends and the data segment begins."
    

        while ea != idc.BADADDR:
            try:
                if idc.GetSegmentAttr(ea, idc.SEGATTR_TYPE) == self.CODE:
                    if idc.GetFunctionName(ea) != '':
                        ea = idc.FindFuncEnd(ea)
                        continue
                    else:
                        if idc.MakeFunction(ea):
                            func_count += 1
                        elif idc.MakeCode(ea):
                            code_count += 1
            except:
                pass
            
            ea = idc.NextAddr(ea)
    
        print "Created %d new functions and %d new code blocks\n" % (func_count, code_count)



class codatify_t(idaapi.plugin_t):
    flags = 0
    comment = ""
    help = ""
    wanted_name = "Define all data and code"
    wanted_hotkey = ""

    def init(self):
        self.menu_context = idaapi.add_menu_item("Options/", "Fixup code", "Alt-3", 0, self.fix_code, (None,))
        self.menu_context = idaapi.add_menu_item("Options/", "Fixup data", "Alt-4", 0, self.fix_data, (None,))
        return idaapi.PLUGIN_KEEP

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

    def run(self, arg):
        pass

    def fix_code(self, arg):
        cd = Codatify()
        cd.codeify()

    def fix_data(self, arg):
        cd = Codatify()
        cd.stringify()
        cd.datify()
        cd.pointify()

def PLUGIN_ENTRY():
    return codatify_t()


########NEW FILE########
__FILENAME__ = fuzz
import idc

class _FuzzHelper(object):
    
    def __init__(self, idasim):
        self.idasim = idasim

    def sanitize(self, data):
        try:
            return data.replace('"', '\\"')
        except:
            return data

    def display(self, message):
        print "%-25s %s" % (idc.GetFuncOffset(self.idasim.cpu.ReturnAddress()), message)
    
class Fuzz(object):

    def __init__(self, idasim):
        self.helper = _FuzzHelper(idasim)
        self.idasim = idasim

    def strcpy(self, dst, src=''):
        self.helper.display('strcpy(0x%X, "%s")' % (dst, self.helper.sanitize(src)))
        return None

    def strcat(self, dst='', src=''):
        self.helper.display('strcat("%s", "%s")' % (self.helper.sanitize(dst), self.helper.sanitize(src)))
        return None

    def sprintf(self, dst, fmt=''):
        string = self.idasim.vsprintf(fmt, 2)
        self.helper.display('sprintf(0x%X, "%s")' % (dst, self.helper.sanitize(string)))
        return None

    def system(self, cmd=''):
        self.helper.display('system("%s")' % self.helper.sanitize(cmd))
        return None

    def popen(self, cmd='', attrib=''):
        self.helper.display('popen("%s", "%s")' % (self.helper.sanitize(cmd), self.helper.sanitize(attrib)))
        return None

    def strncpy(self, dst, src='', n=0):
        if len(src) >= n:
            self.helper.display('strncpy(0x%X, "%s", %d)' % (dst, self.helper.sanitize(src), n))
        return None

    def strncat(self, dst='', src='', n=0):
        self.helper.display('strncat("%s", "%s", %d)' % (self.helper.sanitize(dst), self.helper.sanitize(src), n))

    def snprintf(self, dst, size, fmt=''):
        string = self.idasim.vsprintf(fmt, 3)
        if len(string) >= size:
            self.helper.display('snprintf(0x%X, %d, "%s")' % (dst, size, self.helper.sanitize(string)))
        return None

    def printf(self, fmt=''):
        if '%' not in fmt:
            self.helper.display('printf("%s")' % self.helper.sanitize(fmt))
        return None

    def fprintf(self, fd, fmt=''):
        if '%' not in fmt:
            self.helper.display('fprintf(%d, "%s")' % (fd, self.helper.sanitize(fmt)))
        return None


########NEW FILE########
__FILENAME__ = libc
import os
import time
import idc
import idaapi
import idautils

class LibC(object):
    '''
    Class containing simulators for various common libc functions.
    '''

    __IDASIM_DEFAULT_HANDLER_CLASS__ = True

    def __init__(self, idasim=None):
        '''
        Class constructor.
        '''
        self.idasim = idasim

    def sleep(self, t):
        time.sleep(t)
        return 0

    def atoi(self, string=''):
        return int(string)

    def atol(self, string=''):
        return self.atoi(string)

    def atoll(self, string=''):
        return self.atoi(string)

    def atoq(self, string=''):
        return self.atoi(string)

    def strlen(self, string=''):
        return len(string)

    def getenv(self, envar=''):
        return os.getenv(envar) + "\x00"

    def malloc(self, n):
        return "\x00" * n

    def memset(self, buf, c, n):
        idc.DbgWrite(buf, (chr(c) * n))
        return buf

    def memcpy(self, dst, src, n):
        idc.DbgWrite(dst, idc.GetManyBytes(src, n, use_dbg=False))
        return dst

    def strcpy(self, dst, src=''):
        '''
        Monitors, reports and simulates the strcpy function.
        '''
        print 'strcpy(0x%X, "%s")' % (dst, src)
        idc.DbgWrite(dst, src + "\x00")
        return dst

    def strcat(self, dst, src=''):
        '''
        Monitors, reports and simulates the strcat function.
        '''
        print 'strcat(0x%X, "%s")' % (dst, src)
        addr = dst + len(idc.GetString(dst))
        idc.DbgWrite(addr, src + "\x00")
        return dst

    def strncpy(self, dst, src='', n=0):
        idc.DbgWrite(dst, src + "\x00", length=n)
        return dst

    def strncat(self, dst, src='', n=0):
        addr = dst + len(idc.GetString(dst))
        idc.DbgWrite(addr, src + "\x00", length=n)
        return dst

    def strdup(self, string=''):
        return string + "\x00"

    def strcmp(self, s1='', s2=''):
        if s1 == s2:
            return 0
        else:
            return 1

    def strncmp(self, s1='', s2='', n=0):
        if s1[:n] == s2[:n]:
            return 0
        else:
            return 1

    def memcmp(self, dp1, dp2, n):
        d1 = idc.DbgRead(dp1, n)
        d2 = idc.DbgRead(dp2, n)
    
        if d1 == d2:
            return 0
        else:
            return 1

    def memchr(self, dp, c, n):
        c = chr(c)
        data = idc.DbgRead(dp, n)
        
        offset = data.find(c)

        if offset == -1:
            return 0
        else:
            return dp + offset

    def system(self, command=''):
        '''
        Displays the system() command, does not execute.
        '''
        print '0x%X : system("%s");' % (self.idasim.cpu.ReturnAddress(), command)
        return 0

    def strstr(self, hayptr, needle=''):
        haystack = idc.GetString(hayptr)
        offset = haystack.find(needle)
        
        if offset == -1:
            return 0
        else:
            return hayptr + offset

    def strchr(self, hayptr, needle):
        haystack = idc.GetString(hayptr)
        needle = chr(needle)
        offset = haystack.find(needle)
        
        if offset == -1:
            return 0
        else:
            return hayptr + offset

    def daemon(self):
        '''
        Fakes a daemon(), returns 0.
        '''
        return 0

    def fork(self):
        '''
        Fakes a fork(), always returns 0.
        '''
        return 0
    
    def free(self, address):
        '''
        Frees heap data not allocated by IDASimulator.
        '''
        if self.idasim.mmu.allocated_addresses.has_key(address):
            return 0
        else:
            return None

    def strtol(self, string='', base=0):
        return int(string, base)

    def strtoul(self, string='', base=0):
        return self.strtol(string, base)
    
    def strtod(self, string='', base=0):
        return self.strtod(string, base)

    def strcasecmp(self, s1='', s2=''):
        if s1.lower() == s2.lower():
            return 0
        else:
            return 1

    def strncasecmp(self, s1='', s2='', n=0):
        if s1[:n].lower() == s2[:n].lower():
            return 0
        else:
            return 1
    
    def exit(self, code):
        '''
        Prints exit code and stops debugger.
        '''
        print "Exit code:", code
        idc.StopDebugger()

    def setgroups(self):
        '''
        Fakes setgroups(), returns 0.
        '''
        return 0


########NEW FILE########
__FILENAME__ = libcsman
import idc

class LibCSMAN(object):

    _CONFIG = "/tmp/nvram.cfg"

    __IDASIM_DEFAULT_HANDLER_CLASS__ = True

    def __init__(self, idasim):
        self.idasim = idasim

        self.config = {}

        try:
            for line in open(self._CONFIG_FILE).readlines():
                if '=' in line:
                    kv = line.strip().split('=')
                    name = kv[0]
                    key = int(kv[1], 16)
                    if len(kv) == 3:
                        value = kv[2].decode('string_escape')
                    else:
                        value = "\x00"

                    if not self.config.has_key(key):
                        self.config[key] = {
                                'name'    : name,
                                'value'    : value
                        }
        except:
            pass

    def open_csman(self):
        return 128

    def close_csman(self):
        return 0

    def write_csman(self, fd, key, buf, size, default):
        return 0

    def read_csman(self, fd, key, value, size, default):
        if self.config.has_key(key):
            print "read_csman(%s)" % self.config[key]['name']
            idc.DbgWrite(value, self.config[key]['value'])
        else:
            print "UNKNOWN CSID: 0x%.8X called from 0x%.8X" % (key, self.idasim.cpu.ReturnAddress())

        return 0


########NEW FILE########
__FILENAME__ = libnvram
import idc

class LibNVRAM(object):

    _CONFIG_FILE = "/tmp/nvram.cfg"

    __IDASIM_DEFAULT_HANDLER_CLASS__ = True

    def __init__(self, idasim):
        self.idasim = idasim

        self.config = {}

        try:
            for line in open(self._CONFIG_FILE).readlines():
                if '=' in line:
                    kv = line.strip().split('=')
                    if not self.config.has_key(kv[0]):
                        self.config[kv[0]] = kv[1]
        except:
            pass

    def nvram_init(self):
        return 0

    def nvram_get(self, zero, key=''):
        return self.nvram_bufget(zero, key)

    def nvram_bufget(self, zero, key=''):
        try:
            value = self.config[key]
        except:
            value = ''

        print "nvram_get: {'%s' : '%s'}" % (key, value)
        return value + "\x00"

    def nvram_bufset(self, zero, key='', value=''):
        self.config[key] = value
        print "nvram_set: {'%s' : '%s'}" % (key, value)
        return 0

    def nvram_get_ex(self, key='', dst=0, size=0):
        idc.DbgWrite(dst, self.nvram_bufget(0, key)[:size])
        return 0

    def nvram_match(self, key='', match=''):
        if self.nvram_bufget(0, key)[:-1] == match:
            return 1
        return 0

    def nvram_invmatch(self, key='', match=''):
        if self.nvram_match(key, match):
            return 0
        return 1


########NEW FILE########
__FILENAME__ = pthread

class PThread(object):

    __IDASIM_DEFAULT_HANDLER_CLASS__ = True

    def __init__(self, idasim):
        self.idasim = idasim

    def pthread_create(self, thread, attr, start_routine, arg):
        '''
        Calls the start_routine, does not create a thread.
        '''
        self.idasim.app.Call(start_routine, arguments=[arg], block_until_return=False)
        return None

########NEW FILE########
__FILENAME__ = stdio
import os
import sys
import idc

class Stdio(object):
    '''
    Class containing simulators for common stdio functions.
    '''

    __IDASIM_DEFAULT_HANDLER_CLASS__ = True

    def __init__(self, sim=None):
        '''
        Class constructor.
        '''
        if sim is not None:
            self.mmu = sim.mmu
            self.cpu = sim.cpu
            self.sim = sim

        self.file_descriptors = {
                0 : sys.stdin,
                1 : sys.stdout,
                2 : sys.stderr,
        }

    def _next_fd(self):
        i = 0
        while self.file_descriptors.has_key(i):
            i += 1
        return i

    def _add_fd(self, fp):
        fd = self._next_fd()
        self.file_descriptors[fd] = fp
        return fd

    def _close(self, fd):
        if self.file_descriptors.has_key(fd):
            self.file_descriptors[fd].close()
            del self.file_descriptors[fd]

    def _open(self, fname, mode="rwb"):
        try:
            fp = open(fname, mode)
            fd = self._add_fd(fp)
        except:
            fd = -1

        return fd

    def _read(self, fd, size):
        data = ""
        if self.file_descriptors.has_key(fd):
            data = self.file_descriptors[fd].read(size)
        return data

    def _write(self, fd, data):
        if self.file_descriptors.has_key(fd):
            self.file_descriptors[fd].write(data)
        return len(data)

    def puts(self, string=''):
        print string
        return 0

    def printf(self, fmt=''):
        print self.sim.vsprintf(fmt, 1),
        return 0

    def syslog(self, i, fmt=''):
        print self.sim.vsprintf(fmt, 2)
        return 0

    def fprintf(self, fd, fmt=''):
        formatted_string = self.sim.vsprintf(fmt, 2)

        if self.file_descriptors.has_key(fd) and fd != 0:
            self._write(fd, formatted_string)
        else:
            print formatted_string,
        return 0

    def sprintf(self, dst, fmt=''):
        '''
        Monitors, reports and simulates sprintf.
        '''
        data = self.sim.vsprintf(ftm, 2)
        print 'sprintf(0x%X, "%s")' % (dst, data)
        idc.DbgWrite(dst, data + "\x00")
        return len(data)

    def snprintf(self, dst, n, fmt=''):
        idc.DbgWrite(dst, self.sim.vsprintf(fmt, 3)[:n] + "\x00")
        return dst

    def popen(self, command='', mode=''):
        '''
        Displays the popen() command, does not execute.
        '''
        print '0x%X : popen("%s", "%s");' % (self.cpu.ReturnAddress(), command, mode)
        #fp = os.popen(command, mode)
        #return self._add_fd(fp)
        return 0

    def pclose(self, fd):
        self._close(fd)
        return 0

    def fopen(self, fname='', modes=''):
        fd = self._open(fname, modes)
        if fd > -1:
            return fd
        else:
            return 0
    
    def fclose(self, fd):
        self._close(fd)
        return 0

    def fread(self, ptr, size, nmemb, fd):
        data = self._read(fd, (size * nmemb))
        idc.DbgWrite(ptr, data)
        return len(data)

    def fwrite(self, ptr, size, nmemb, fd):
        data = idc.DbgRead(ptr, (size * nmemb))
        self._write(fd, data)
        return len(data)

    def fflush(self, fd):
        if self.file_descriptors.has_key(fd):
            self.file_descriptors[fd].flush()
        return 0

    def fgets(self, dst, size, fd):
        if self.file_descriptors.has_key(fd):
            while not data.endswith('\n') and len(data) < (size-1):
                data += self._read(fd, 1)

            data += "\x00"
            idc.DbgWrite(dst, data, len(data))

        return dst

    def fseek(self, fd, offset, whence):
        if self.file_descriptors.has_key(fd):
            self.file_descriptors[fd].seek(offset, whence)
            return self.file_descriptors[fd].tell()
        return -1

    def rewind(self, fd):
        if self.file_descriptors.has_key(fd):
            self.file_descriptors[fd].seek(0, 0)
        return 0

    def ftell(self, fd):
        if self.file_descriptors.has_key(fd):
            return self.file_descriptors[fd].tell()
        return -1



########NEW FILE########
__FILENAME__ = idasimulator
import os
import pickle
import inspect
import idaapi
import idautils
import idc
import idasim

IDASIM = None

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

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

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

        @sim - idasimulator_t class instance.

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

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

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

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

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

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

        return cfg

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

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

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

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

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

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

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

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

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

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

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

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

        self._save_config_file(fdata)

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

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

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


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

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

        self.PopulateItems()

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


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

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

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

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

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

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

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

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

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

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

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

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

        @name - Named location.

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

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

        @name - Handler name.

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

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

        @name - Named location.

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

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

        @name - Named location.

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

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

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

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

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

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

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

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

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

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

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

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

        entry = {}
        entry[handler_name] = {}

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

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

        return entry

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

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

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

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

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

def PLUGIN_ENTRY():
    return idasimulator_t()


########NEW FILE########
__FILENAME__ = application
__all__ = ['Application']

import idc
from architecture import Architecture

class Application(object):
    '''
    Class for invoking functions in the target process.
    '''

    def __init__(self):
        '''
        Class constructor.
        '''
        self.cpu = Architecture()

    def Call(self, function, arguments=[], retaddr=0, block_until_return=True):
        '''
        Call a given function. Arguments must already be configured.
        This should not be used to call functions hooked with IDASimulator or it likely won't work.

        @function           - The name or address of the function to call.
        @arguments          - A list of function arguments.
        @retaddr            - The address to return to.
        @block_until_return - If set to True, this method will not return until the function does.
                              If set to False, this method will return immediately after calling the function.

        Returns the return value of the function on success.
        Returns None on failure, or if block_until_return is False.
        '''
        retval = None

        # Process should already be paused, but just in case...
        idc.PauseProcess()

        # If a function name was specified, get its address
        if isinstance(function, type('')):
            function = idc.LocByName('.' + function)

            if function == idc.BADADDR:
                function = idc.LocByName(function)

        if function != idc.BADADDR:
            if not retaddr:
                retaddr = self.cpu.ProgramCounter()

            # Set the specified function arguments
            self.cpu.SetArguments(arguments)

            # Do any arch-specific initialization before the function call
            self.cpu.PreFunctionCall(function)
    
            # Set up the return address and point the program counter to the start of the target function    
            self.cpu.ReturnAddress(value=retaddr)
            self.cpu.ProgramCounter(value=function)
            idc.Jump(function)

            if block_until_return:
                # Resume process and wait for the target function to return
                idc.StepUntilRet()
                idc.GetDebuggerEvent(idc.WFNE_CONT|idc.WFNE_SUSP, -1)
                idc.Jump(retaddr)
                retval = self.cpu.ReturnValue()
            else:
                idc.ResumeProcess()

        return retval

########NEW FILE########
__FILENAME__ = architecture
__all__ = ['Architecture']

import idc

class Architecture(object):
    '''
    Abstraction class for accessing CPU-specific registers and data.
    '''

    BIG = 'big'
    LITTLE = 'little'

    ARCH = {
        'mips'    : {
                'spoffset'  : 0x10,
                'argreg'    : ['a0', 'a1', 'a2', 'a3'],
                'retval'    : ['v0', 'v1'],
                'sp'        : 'sp',
                'ra'        : 'ra',
                'pc'        : 'pc',
                # Common calling convention for MIPS GCC is to place the address of the function
                # into $t9 and then jalr $t9. The callee then expects $t9 to point to the beginning
                # of itself, so $t9 is used to calculate the relative offset to the global pointer.
                # If $t9 is not set appropriately, any data/code xrefs that rely on $gp will fail.
                'callreg'   : 't9'
        },
        'arm'    : {
                'spoffset'  : 0x10,
                'argreg'    : ['R0', 'R1', 'R2', 'R3'],
                'retval'    : ['R0', 'R1'],
                'sp'        : 'SP',    
                'ra'        : 'LR',
                'pc'        : 'PC',
        },
        'ppc'    : {
                'spoffset'  : 8,
                'argreg'    : ['R3', 'R4', 'R5', 'R6', 'R7', 'R8', 'R9', 'R10'],
                'retval'    : ['R3'],
                'sp'        : 'R1',
                'ra'        : 'LR',
                'pc'        : 'PC',
                # GDB stubs for PPC are special...
                'bpt_size'  : 1,
                'bpt_type'  : idc.BPT_EXEC
        },
        'ia32'    : {
                'spoffset'  : 4,
                'argreg'    : [],
                'retval'    : ['EAX'],
                'sp'        : 'ESP',
                'ra'        : '*ESP',
                'pc'        : 'EIP',
        },
        'ia64'    : {
                'spoffset'  : 8,
                'argreg'    : ['RDI', 'RSI', 'RDX', 'RCX', 'R8', 'R9'],
                'retval'    : ['RAX'],
                'sp'        : 'RSP',
                'ra'        : '*RSP',
                'pc'        : 'RIP',
        },
        'win64'    : {
                'spoffset'  : 8,
                'argreg'    : ['RDX', 'RCX', 'R8', 'R9'],
                'retval'    : ['RAX'],
                'sp'        : 'RSP',
                'ra'        : '*RSP',
                'pc'        : 'RIP',
        }
    }
    
    PROCESSORS = {
        'mipsl'    : [{
                'architecture'  : 'mips',
                'endianess'     : LITTLE,
                'bits'          : 32,
        }],
        'mipsb'    : [{
                'architecture'    : 'mips',
                'endianess'    : BIG,
                'bits'        : 32,
        }],
        'arm'    : [{
                'architecture'    : 'arm',
                'endianess'    : LITTLE,
                'bits'        : 32,
        }],
        'armb'    : [{
                'architecture'    : 'arm',
                'endianess'    : BIG,
                'bits'        : 32,
        }],
        'ppc'    : [{
                'architecture'    : 'ppc',
                'endianess'    : BIG,
                'bits'        : 32,
        }],
        'metapc': [{
                'architecture'    : 'ia32',
                'endianess'    : LITTLE,
                'bits'        : 32,
               },
               {
                'architecture'    : 'win64',
                'endianess'    : LITTLE,
                'bits'        : 64,
                # Windows passes args differently in x86_64
                'file_types'    : [idc.FT_PE, idc.FT_EXE]
               },
               {
                'architecture'    : 'ia64',
                'endianess'    : LITTLE,
                'bits'        : 64,
               }
        ],
                
    }

    def __init__(self):
        '''
        Class constructor.
        
        Returns None.
        '''
        self.cpu = None
        self.cpu_name = None
        self.architecture = None
        self.bits64 = False
        self.bits = 0
        self.bsize = 0    

        self.__cpu_id()

        if self.cpu == None:
            if self.bits64:
                bits = '64'
            else:
                # This is an assumption, but it's only for the error message.
                # TODO: How to determine if a target is 16/32 bit from IDA's API?
                bits = '32'

            raise Exception("Unsupported cpu type: %s.%s" % (self.cpu_name, bits))

    def __stack_dword(self, n, value=None):
        addr = self.StackPointer() + self.cpu['spoffset'] + (n * self.bsize)

        if value is not None:
            sval = self.ToString(value, size=self.bsize)
            idc.DbgWrite(addr, sval)

        return idc.DbgDword(addr)

    def __reg_value(self, reg, value=None):
        if value is not None:
            if reg.startswith('*'):
                idc.DbgWrite(idc.GetRegValue(reg[1:]), self.ToString(value))
            else:
                idc.SetRegValue(value, reg)
        
        if reg.startswith('*'):
            return idc.DbgDword(idc.GetRegValue(reg[1:]))
        else:
            return idc.GetRegValue(reg)

    def __cpu_id(self):
        self.cpu_name = idc.GetShortPrm(idc.INF_PROCNAME).lower()
        
        if (idc.GetShortPrm(idc.INF_LFLAGS) & idc.LFLG_64BIT) == idc.LFLG_64BIT:
            self.bits64 = True
        else:
            self.bits64 = False

        for (processor, architectures) in self.PROCESSORS.iteritems():
            if self.cpu_name == processor:
                for arch in architectures:
                    # Only use 64-bit processor modules for a 64 bit binary
                    if (self.bits64 and arch['bits'] != 64) or (not self.bits64 and arch['bits'] == 64):
                        continue

                    # If specific file types were specified for this processor module, make sure the target file is in that list
                    if arch.has_key('file_types') and idc.GetShortPrm(idc.INF_FILETYPE) not in arch['file_types']:
                        continue

                    self.cpu = self.ARCH[arch['architecture']]
                    self.architecture = arch['architecture']
                    self.endianess = arch['endianess']
                    self.bits = arch['bits']
                    self.bsize = self.bits / 8
                    break

                if self.cpu:
                    break
        return None

    def ToString(self, value, size=None):
        '''
        Converts an integer value of size bytes into a raw string of bytes.

        @value - Integer value to be represented as a raw string.
        @size  - Size of the integer value, in bytes.

        Returns a raw string containing the integer value in string form, and in the appropriate endianess.
        '''
        data = ""

        if size is None:
            size = self.bsize

        for i in range(0, size):
            data += chr((value >> (8*i)) & 0xFF)

        if self.endianess != self.LITTLE:
            data = data[::-1]

        return data

    def FromString(self, data, size=None):
        '''
        Converts raw string data into an integer value, with appropriate endianess.

        @data - Raw string data.
        @size - Number of bytes to convert.

        Returns an integer value.
        '''
        i = 0
        value = 0

        if size is None:
            size = len(data)

        if self.endianess != self.LITTLE:
            data = data[::-1]

        for c in data[:size]:
            value += (ord(c) << (8*i))
            i += 1

        return value

    def GetArguments(self, index, n):
        '''
        Get a list of function arguments. Any valid string pointers will be converted to strings.

        @index - First argument index, 0-indexed.
        @n     - The number of arguments to retrieve.

        Returns a list of n arguments.
        '''
        args = []

        for j in range(index, n+index):
            arg = self.Argument(j)
            try:
                sval = idc.GetString(arg)
            except:
                sval = None

            if sval is not None:
                args.append(sval)
            else:
                args.append(arg)

        return args

    def SetArguments(self, arguments):
        '''
        Sets a list of function arguments.

        @arguments - List of function arguments.

        Returns None.
        '''
        for i in range(0, len(arguments)):
            self.Argument(i, value=arguments[i])

    def Argument(self, n, value=None):
        '''
        Read/write function arguments.

        @n     - Argument index number, 0-indexed.
        @value - If specified, the argument will be set to this value.

        Returns the current argument value.
        '''
        regn = len(self.cpu['argreg'])

        if value is not None:
            if n < regn:
                self.__reg_value(self.cpu['argreg'][n], value)
            else:
                self.__stack_dword(n-regn, value)
            
        if n < regn:
            return self.__reg_value(self.cpu['argreg'][n])
        else:
            return self.__stack_dword(n-regn)

    def StackPointer(self, value=None):
        '''
        Read/write the stack pointer register.

        @value - If specified, the stack pointer register will be set to this value.

        Returns the current stack pointer register value.
        '''
        return self.__reg_value(self.cpu['sp'], value)

    def ReturnValue(self, value=None, n=0):
        '''
        Read/write the function return register value.

        @value - If specified, the return register will be set to this value.
        @n     - Return register index number, for those architectures with multiple return registers.

        Returns the current return register value.
        '''
        return self.__reg_value(self.cpu['retval'][n], value)

    def ProgramCounter(self, value=None):
        '''
        Read/write the program counter register.

        @value - If specified, the program counter register will be set to this value.

        Returns the current value of the program counter register.
        '''
        return self.__reg_value(self.cpu['pc'], value)

    def ReturnAddress(self, value=None):
        '''
        Read/write the return address.

        @value - If specified, the return address will be set to this value.

        Returns the current return address value.
        '''
        return self.__reg_value(self.cpu['ra'], value)

    def StackCleanup(self):
        '''
        Cleans up values automatically pushed onto the stack by some architectures (return address in x86 for example).
        '''
        if self.cpu['ra'].startswith('*') and self.cpu['ra'][1:] == self.cpu['sp']:
            self.StackPointer(self.StackPointer() + self.bsize)

    def SetBreakpoint(self, address):
        '''
        Some GDB stubs for various architectures require different breakpoint settings.
        This method sets the appropriate breakpoint for the selected architecture.

        @address - The breakpoint address.

        Returns True on success, False on failure.
        '''
        bpt_size = 0
        bpt_type = idc.BPT_SOFT

        if self.cpu.has_key('bpt_size'):
            bpt_size = self.cpu['bpt_size']
        if self.cpu.has_key('bpt_type'):
            bpt_type = self.cpu['bpt_type']

        return idc.AddBptEx(address, bpt_size, bpt_type)

    def PreFunctionCall(self, function):
        '''
        Configure architecture-specific pre-requisites before calling a function.
        Called internally by Application.Call.

        @function - The address of the function to call.

        Returns None.
        '''
        if self.cpu.has_key('callreg'):
            idc.SetRegValue(function, self.cpu['callreg'])

########NEW FILE########
__FILENAME__ = exceptions
class JumpTo(Exception):
    pass

class GoTo(Exception):
    pass

########NEW FILE########
__FILENAME__ = handler
__all__ = ['IDASimFunctionHandler']

import idc
import idaapi
import inspect
import traceback
from exceptions import *
from architecture import Architecture

class IDASimFunctionHandler(object):
    '''
    Registers and manages function simulators.
    '''
    FUNCTION_HANDLERS = {}
    DEFAULT_HANDLER = None
    BPT_CND = '%s.Handler()'
    STUB_NAMING_CONVENTION = '.%s'

    def __init__(self, idbm, name=None, verbose=False):
        '''
        Class constructor.

        @idbm    - Instance of IDASimMMU.
        @name    - Name that will be assigned to the class instance.
        @verbose - Enable verbose mode.

        Returns None.
        '''
        self.idbm = idbm
        self.name = name
        self.verbose = verbose
        self.cpu = Architecture()

        if self.name == None:
            self.name = self.__get_my_name()

        self.bpt_cnd = self.BPT_CND % self.name

        self.UnregisterHandlers()

        # Eval this IDC expression to ensure that Python is set as the
        # preferred external scripting language. This is necessary for the
        # Python function handler to operate correctly.
        idc.Eval('RunPlugin("python", 3)')

    def cleanup(self):
        idc.Eval('RunPlugin("python", 4)')

    def __del__(self):
        self.cleanup()

    def __get_my_name(self):
        '''
        This is a hack to get the name of the class instance variable. For internal use only.
        '''
        i = -3
        (filename, line_number, function_name, text) = traceback.extract_stack()[i]
        name = text[:text.find('=')].strip()
        while 'self' in name:
            i -= 1
            (filename, line_number, function_name, text) = traceback.extract_stack()[i]
            name = name.replace('self', text[:text.find('=')].strip())
        return name

    def SetHandlerBreakpoint(self, address):
        '''
        Sets a handler breakpoint on the specified address.

        @address - Address to set the breakpoint at.

        Returns True on success, False on failure.
        '''
        # Some remote debugger stubs have special needs for different architectures (e.g., gdb).
        # Thus, setting breakpoints should be done through the architecture abstraction class, 
        # rather than directly through AddBpt/AddBptEx.
        self.cpu.SetBreakpoint(address)

        # A bug in versions of IDAPython shipped with IDA prior to 6.4sp1 improperly interpreted 
        # the is_lowcnd value set via SetBptCnd/SetBptCndEx. Do this directly through idaapi
        # ourselves in order to support older versions.
        bpt = idaapi.bpt_t()
        idaapi.get_bpt(address, bpt)
        bpt.condition = self.bpt_cnd
        bpt.flags &= ~idc.BPT_LOWCND
        return idaapi.update_bpt(bpt)

    def __register_internal_handler(self, name, handler):
        '''
        Internal handler registration function. For internal use only.
        '''
        if type(name) == type(""):
            address = idc.LocByName(name)
        else:
            address = name

        if address != idc.BADADDR:
            bpt_result = self.SetHandlerBreakpoint(address)

            if bpt_result:
                self.FUNCTION_HANDLERS[name] = {}
                self.FUNCTION_HANDLERS[name]["handler"] = handler
                self.FUNCTION_HANDLERS[name]["address"] = address

            return bpt_result
        else:
            return False

    def Handler(self):
        '''
        Breakpoint condition handler, called by IDA to evaluate conditional brekpoints. It in turn calls the 
        appropriate function handler, populates the return value and puts execution back at the return address. 
    
        This is a (slight) abuse of IDA's conditional breakpoints; this function always returns 0, indicating that
        the breakpoint condition has not been met. However, it does ensure that every call to a given function
        can be intercepted and simulated, regardless of whether the process is running freely, or the function has 
        been stepped over, stepped into, etc.
        '''
        retval = 0
        retaddr = None

        if self.verbose:
            print self.FUNCTION_HANDLERS

        for (name, properties) in self.FUNCTION_HANDLERS.iteritems():
            if self.cpu.ProgramCounter() == properties["address"]:
                handler = properties["handler"]
                break

        # If no explicit handler was found, use the default handler
        if not handler and self.DEFAULT_HANDLER:
            handler = self.DEFAULT_HANDLER

        if handler:
            if self.verbose:
                print "Using function handler:", handler.__name__

            parameters = {}
        
            # Enumerate the arguments and default values for the handler    
            args, varargs, keywords, defaults = inspect.getargspec(handler)
            try:
                defaults = dict(zip(reversed(args), reversed(defaults)))
            except:
                defaults = {}

            # Build the handler parameters
            try:
                i = 0
                for arg in args:
                    if arg != 'self':
                        parameters[arg] = self.cpu.Argument(i)
                        
                        if defaults.has_key(arg):
                            # If default value is of type string, get the string automatically
                            if type(defaults[arg]) == type(''):
                                parameters[arg] = idc.GetString(parameters[arg])
                            # If default value is of type list, get an array of bytes
                            elif type(defaults[arg]) == type([]) and len(defaults[arg]) == 1:
                                parameters[arg] = [c for c in idc.DbgRead(parameters[arg], defaults[arg][0])]
                        i += 1
            except Exception, e:
                print "WARNING: Failed to parse handler parameters:", str(e)
                parameters = {}

            try:
                retval = handler(**parameters)
            except JumpTo, offset:
                retaddr = self.cpu.ReturnAddress() + offset.message
            except GoTo, addr:
                retaddr = addr.message
            except Exception, e:
                print "WARNING: Failed to simulate function '%s': %s" % (handler.__name__, str(e))
                retval = 0

            if retval is not None:
                if retaddr is None:
                    retaddr = self.cpu.ReturnAddress()

                # If a string type was returned by the handler, place the string in memory and return a pointer
                if type(retval) == type(""):
                    retval = self.idbm.malloc(retval)
                # Map python's True and False to 1 and 0 repsectively
                elif retval == True:
                    retval = 1
                elif retval == False:
                    retval = 0

                self.cpu.ReturnValue(retval)
                self.cpu.ProgramCounter(retaddr)
                self.cpu.StackCleanup()
    
                # Since the PC register is manually manipulated, a breakpoint set on the return
                # address won't be triggered. In this case, make sure we pause the process manually.
                if idc.CheckBpt(self.cpu.ProgramCounter()) > 0:
                    idc.PauseProcess()
            
        return 0

    def RegisterDefaultHandler(self, handler):
        '''
        Register a default "catch-all" handler.

        @handler - Method/function handler.

        Returns None.
        '''
        self.DEFAULT_HANDLER = handler

    def UnregisterDefaultHandler(self):
        '''
        Unregister a default "catch-all" handler.
        
        Returns None.
        '''
        self.DEFAULT_HANDLER = None

    def RegisterHandler(self, name, handler, stubs=True):
        '''
        Registers a given function handler for a given function name.
    
        @name    - Name of the function.
        @handler - The function handler to call.
        @stubs   - If True, handle calls to both extern and stub addresses.

        Returns True on success, False on failure.
        '''

        retval = self.__register_internal_handler(name, handler)

        if retval and stubs and type(name) == type(""):
            stub_name = self.STUB_NAMING_CONVENTION % name
            retval = self.__register_internal_handler(stub_name, handler)
            
        return retval

    def RegisterHandlers(self, handlers, stubs=True):
        '''
        Registers a set of function handlers.

        @handlers - A dictionary consisting of 'name':handler pairs.
        @stubs    - If True, handle calls to both extern and stub addresses.
        
        Returns the number of handlers successfully registered.
        '''
        count = 0

        for (name, handler) in handlers.iteritems():
            if self.RegisterHandler(name, handler, stubs):
                count += 1

        return count

    def UnregisterHandler(self, name, stubs=True):
        '''
        Removes a function handler by name.

        @name  - The name of the function handler to be removed.
        @stubs - If True, corresponding function stub handlers that were automatically created by RegisterHandler will also be removed.

        Returns None.
        '''
        addr = None
        stub_name = None
        stub_addr = None

        if name is not None:
            try:
                stub_name = self.STUB_NAMING_CONVENTION % name
            except:
                pass

            if self.FUNCTION_HANDLERS.has_key(name):
                addr = self.FUNCTION_HANDLERS[name]['address']

            if self.FUNCTION_HANDLERS.has_key(stub_name):
                stub_addr = self.FUNCTION_HANDLERS[stub_name]['address']

        if addr is not None and name is not None:
            idc.DelBpt(addr)
            del self.FUNCTION_HANDLERS[name]

        if stubs and stub_addr is not None and stub_name is not None:
            idc.DelBpt(stub_addr)
            del self.FUNCTION_HANDLERS[stub_name]

    def UnregisterHandlers(self, purge=False):
        '''
        Deletes breakpoints for all registered handlers.

        @purge - Removes all handlers for all instances of IDBFunctionHandler.

        Returns None.
        '''
        self.UnregisterDefaultHandler()
    
        if not purge:
            # Only remove this instance's handlers
            for (name, info) in self.FUNCTION_HANDLERS.iteritems():
                condition = idc.GetBptAttr(info['address'], idc.BPTATTR_COND)

                if condition == self.bpt_cnd:
                    idc.DelBpt(info['address'])
        else:
            # Try to remove ALL instance's handlers (this could remove other conditional breakpoints...)
            for i in range(0, idc.GetBptQty()):
                ea = idc.GetBptEA(i)
                condition = idc.GetBptAttr(ea, idc.BPTATTR_COND)
                if condition.endswith(self.BPT_CND % ''):
                    idc.DelBpt(ea)
        
        self.FUNCTION_HANDLERS = {}

    def GetHandler(self, name):
        '''
        Returns the current handler for the named location.

        @name - Function/location name.

        Returns the handler instance.
        '''
        if self.FUNCTION_HANDLERS.has_key(name):
            return self.FUNCTION_HANDLERS[name]["handler"]
        else:
            return None


########NEW FILE########
__FILENAME__ = idasim
__all__ = ['IDASim']

import idc
import idaapi
import idautils
from mmu import *
from handler import *
from exceptions import *
from application import *
from architecture import *

class IDASimDbgHook(idaapi.DBG_Hooks):
    '''
    Resets the IDASimMMU base address (MP) whenever the debugger is started/stopped.
    Executes startup code when the debugger is started, if specified.
    Only used internally by the IDASim class.
    '''
    def dbg_init(self, idasim):
        self.debugging = False
        self.sim = idasim
        self.hook()

    def dbg_process_start(self, pid, tid, ea, name, base, size):
        self.sim.mmu.reset()
        if not self.debugging:
            self.debugging = True
            self.sim.InitHandler(init=True)

    def dbg_process_exit(self, pid, tid, ea, code):
        self.debugging = False
        self.sim.mmu.reset()

    def dbg_process_attach(self, pid, tid, ea, name, base, size):
        self.sim.mmu.reset()
        if not self.debugging:
            self.debugging = True
            self.sim.InitHandler(init=True)

    def dbg_process_detatch(self, pid, tid, ea):
        self.debugging = False
        self.sim.mmu.reset()

class IDASim(object):
    '''
    Class for easily simulating library function calls and initializing memory/registers when debugging emulated code in IDA. 
    '''

    def __init__(self, handlers={}, debug=False, attach=False, membase=None):
        '''
        Class constructor.

        @handlers  - A dictionary of function names/addresses to simulate and their corresponding handlers.
        @debug     - Set to True to automatically start debugging.
        @attach    - Set to True to attach to a process, rather than directly running the debugger.
        @membase   - Specify the base address to start at when allocating memory.

        Returns None.
        '''
        self.user_handlers = handlers
        self.script = None
        self.script_name = None

        self.cpu = Architecture()
        self.mmu = IDASimMMU()
        self.app = Application()
        self.FunctionHandler = IDASimFunctionHandler(self.mmu)
        self.dbg_hook = IDASimDbgHook()
        self.dbg_hook.dbg_init(self)

        self.__register_handlers()
        
        if attach:
            self.AttachDebugger()
        elif debug:
            self.StartDebugger()

        if membase is not None:
            self.mmu.base(membase)

    def __register_handlers(self):
        '''
        Registers function names and handlers with the IDB function handler.
        For internal use only.
        '''
        for (name, handler) in self.user_handlers.iteritems():
            self.FunctionHandler.RegisterHandler(name, handler)

    def __get_instance_methods(self, instance):
        methods = {}

        for name in dir(instance):
            if not name.startswith('_'):
                obj = getattr(instance, name)
                if 'method' in type(obj).__name__:
                    methods[name] = obj

        return methods

    def InitHandler(self, init=False):
        if self.script is not None:
            if (self.script_name is not None and not init) or (self.script_name is None and init):
                script_globals = {
                        'IDASIM'     : self,
                        'idc'        : idc,
                        'idaapi'    : idaapi,
                        'idautils'    : idautils,
                }

                script_globals.update(self.__get_instance_methods(self))
                script_globals.update(self.__get_instance_methods(self.cpu))

                try:
                    exec(self.script, script_globals)
                except Exception, e:
                    print "Failed to exec startup script:", str(e)
                    print "################"
                    print self.script
                    print "################"
        return None

    def ExecuteOnStart(self, script=None, name=None, disable=False):
        '''
        Specify a Python string to be evaluated when the debugger is started/attahced.

        @script - Python string to be evaluated. If None, this feature will be disabled.

        Returns None.
        '''
        self.script = script

        if disable:
            self.FunctionHandler.UnregisterHandler(self.script_name)
            self.script_name = None
        elif name is not None and script is not None:
            self.FunctionHandler.RegisterHandler(name, self.InitHandler)
            
        self.script_name = name
        
    def vsprintf(self, fmt, index):
        '''
        Builds a string from a format string and format arguments.
                
        @fmt   - The format string.
        @index - The function argument number at which the format string arguments start (0-indexed).

        Returns a formatted string.
        '''
        n = 0
        for i in range(0, len(fmt)-1):
            if fmt[i] == '%' and fmt[i+1] != '%':
                n += 1

        return fmt % tuple(self.cpu.GetArguments(index, n))

    def WaitForDebugger(self):
        '''
        Waits for the debugger event (WFNE_CONT | WFNE_SUSP).
        Called internally by StartDebugger and AttachDebugger.

        Returns None.
        '''
        idc.GetDebuggerEvent(idc.WFNE_CONT | idc.WFNE_SUSP, -1)

    def StartDebugger(self):
        '''
        Starts the debugger (equivalent of pressing F9).

        Returns None.
        '''
        idc.StartDebugger('', '', '')
        self.WaitForDebugger()

    def AttachDebugger(self, pid=-1):
        '''
        Attaches the debugger to a running process.

        @pid - The PID of the process to attach to (user will be prompted if not specified).

        Returns None.
        '''
        idc.AttachProcess(pid, -1)
        self.WaitForDebugger()

    def Malloc(self, data=None, size=0):
        '''
        Allocates space in the debugger's memory.

        @data - Fill the allocated space with this data.
        @size - If data is None, allocate and zero out size bytes of memory.

        Returns the address of the allocated memory.
        '''
        return self.mmu.malloc(data, size)

    def String(self, string, raw=False):
        '''
        Creates a NULL-terminated string in the debugger's memory.

        @string - The string, or list of strings, to place into memory.
        @raw    - If set to True, the string will not be NULL terminated.

        Returns the address, or list of addresses, of the string(s) in memory.
        '''
        addrs = []

        if type(string) == type(""):
            array = [string]
        else:
            array = string

        for s in array:
            if not raw:
                s = s + "\x00"
            addrs.append(self.Malloc(s))

        if type(string) == type(""):
            addrs = addrs[0]

        return addrs

    def Int(self, value, size):
        '''
        Creates an integer value of size bytes in the debugger's memory.

        @value - The integer value, or list of values, to place into memory.
        @size  - The size of the interger value(s), in bytes.

        Returns the address, or a list of addresses, of the integer(s) in memory.
        '''
        data = []

        if type(value) != type([]):
            value = [value]

        for d in value:
            data.append(self.cpu.ToString(d, size))

        return self.String(data, raw=True)

    def DoubleWord(self, dword):
        '''
        Places a double word integer into the debugger's memory.
        
        @dword - The value, or list of values, to place into memory.

        Returns the address, or a list of addresses, of the dword(s) in memory.
        '''
        return self.Int(dword, self.cpu.bsize*2)

    def Word(self, word):
        '''
        Places a word-sized integer into the debugger's memory.

        @word - The four byte integer value, or list of values, to place into memory.

        Returns the address, or a list of addresses, of the word(s) in memory.
        '''
        return self.Int(word, self.cpu.bsize)

    def HalfWord(self, hword):
        '''
        Places a half-word sized integer into the debugger's memory.

        @hword - The two byte value, or list of values, to place into memory.

        Returns the address, or a list of addresses, of the half word(s) in memory.
        '''
        return self.Int(hword, self.cpu.bsize/2)

    def Byte(self, byte):
        '''
        Places one byte of data into the debugger's memory.
        
        @byte - The byte value, or list of values, to place into memory.
        
        Returns the address, or a list of addresses, of the byte(s) in memory.
        '''
        return self.Int(byte, 1)

    def ARGV(self, argv):
        '''
        Allocates space for an argv data structure.

        @argv - A list of argv strings.

        Returns the address of the argv array of pointers.
        '''
        return self.Word(self.String(argv))[0]

    def Cleanup(self):
        '''
        Removes all registered function simulation hooks.

        Returns None.
        '''
        self.FunctionHandler.UnregisterHandlers()


########NEW FILE########
__FILENAME__ = mmu
__all__  = ['IDASimMMU']

import idc
import idaapi
from application import Application
from architecture import Architecture

class IDASimMMU(object):
    '''
    Manages the allocation of memory while running in the debugger.
    The term 'manage' is used very loosely here; it really only allocates memory.
    '''
    ALIGN = 4
    DEFAULT_MP = 0x100000
    SEGNAME = 'MMU'
    LAST_SEGNAME = ['MEMORY', 'RAM']

    def __init__(self, base=None):
        '''
        Class constructor.
        '''
        # Disable this for now, it doesn't work.
        self.use_native_malloc = False
        self.allocated_addresses = {}

        self.app = Application()
        self.cpu = Architecture()

        if base is not None:
            self.MP = self.BASE_MP = base
        else:
            self.MP = self.BASE_MP = idc.BADADDR

    def _detect_membase(self):
        '''
        Attempts to locate a section of memory for IDBMMU's internal memory allocation.
        For internal use only.
        '''
        if self.BASE_MP == idc.BADADDR:

            # Look for the MMU segment
            ea = idc.SegByName(self.SEGNAME)

            # No MMU segment?
            if ea == idc.BADADDR:
                ea = 0

                # Find the very last defined segment
                while True:
                    segea = idc.NextSeg(ea)
                    
                    if segea == idc.BADADDR:
                        break
                    else:
                        ea = segea

                # Is it not a memory segment?
                if idc.SegName(ea) not in self.LAST_SEGNAME:
                    try:
                        # Find the start of the stack
                        ea = idc.SegStart(self.cpu.StackPointer())

                        # Still nothing? Use the default.
                        if ea == idc.BADADDR:
                            ea = self.DEFAULT_MP
                    except:
                        if not self.use_native_malloc:
                            raise Exception("No available segments for memory allocation! Try defining segment %s." % self.SEGNAME)
            self.BASE_MP = ea

        if self.MP == idc.BADADDR:
            self.MP = self.BASE_MP

        return self.BASE_MP

    def reset(self):
        '''
        Resets the current allocation address.
        '''
        self.MP = idc.BADADDR
        self.allocated_addresses = {}

    def base(self, base=None):
        '''
        Set the base address at which to start allocating memory. Default: 0x100000.

        @base - The base address. If specified BASE_MP will be set to this value.

        Returns the current BASE_MP value.
        '''
        if base is not None:
            self.MP = self.BASE_MP = base
        return self.BASE_MP

    def malloc(self, data=None, size=0):
        '''
        Allocates space for data in the debugger's memory and populates it.
    
        @data - Data to place into memory. If None, NULL bytes will be used.
        @size - Size of memory to allocate. If 0, len(data) bytes will be allocated.
    
        Returns the address of the allocated memory.
        '''
        if size == 0 and data is not None:
            size = len(data)
    
        if data is None:
            data = "\x00" * size

        if self.use_native_malloc:
            addr = self.app.Call('malloc', arguments=[size], retaddr=self.cpu.ReturnAddress())
        else:
            self._detect_membase()
    
            addr = self.MP
            self.MP += size
            # This ensures memory addresses are 4-byte aligned. This is important for some architectures.
            if (self.MP % self.ALIGN) > 0:
                self.MP += (self.ALIGN - (self.MP % self.ALIGN))

            # Keep a dictionary of allocated addresses and their sizes
            self.allocated_addresses[addr] = size
    
        idc.DbgWrite(addr, data)
    
        return addr


########NEW FILE########
__FILENAME__ = localxrefs
# IDA Plugin to search for cross references only within the current defined function.
#
# Useful, for example, to find instructions that use a particular register, or that reference a literal value.
#
# Invoke by highlighting the desired text in IDA, then going to Jump->List local xrefs, or by pressing Alt+8.
# Highlighting is also supported; once xrefs are found, type the following in the Python command window:
#
#	Python> localxrefs.highlight()       <-- Highlight all xrefs
#	Python> localxrefs.highlight(False)  <-- Un-highlight all xrefs
#
# Craig Heffner
# Tactical Network Solutions

import idc
import idaapi

localxrefs = None

class LocalXrefs(object):

	UP   = 'Up  '
	DOWN = 'Down'
	THIS = '-   '

	READ    = 'r'
	WRITE   = 'w'

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

	def __init__(self):
		self.xrefs = {}
		self.function = ''
		self._profile_function()

	def _profile_function(self):
		current_ea = ScreenEA()
		current_function = idc.GetFunctionName(current_ea)
		current_function_ea = idc.LocByName(current_function)

		if current_function:
			self.function = current_function

		ea = start_ea = idc.GetFunctionAttr(current_function_ea,  idc.FUNCATTR_START)
		end_ea = idc.GetFunctionAttr(current_function_ea, idc.FUNCATTR_END)

		self.highlighted = idaapi.get_highlighted_identifier()

		while ea < end_ea and ea != idc.BADADDR and self.highlighted:

			i = 0
			match = False
			optype = self.READ
			comment = None

			idaapi.decode_insn(ea)
			
			mnem = idc.GetMnem(ea)

			if self.highlighted in mnem:
				match = True
			elif idaapi.is_call_insn(ea):
				for xref in idautils.XrefsFrom(ea):
					if xref.type != 21:
						name = idc.Name(xref.to)
						if name and self.highlighted in name:
							match = True
							break
			else:	
				while True:
					opnd = idc.GetOpnd(ea, i)
					if opnd:
						if self.highlighted in opnd:
							match = True
							if (idaapi.insn_t_get_canon_feature(idaapi.cmd.itype) & self.OPND_WRITE_FLAGS[i]):
								optype = self.WRITE
						i += 1
					else:
						break

			if not match:
				comment = idc.GetCommentEx(ea, 0)
				if comment and self.highlighted in comment:
					match = True
				else:
					comment = idc.GetCommentEx(ea, 1)
					if comment and self.highlighted in comment:
						match = True
					else:
						comment = None

			if match:
				if ea > current_ea:
					direction = self.DOWN
				elif ea < current_ea:
					direction = self.UP
				else:
					direction = self.THIS

				self.xrefs[ea] = {
					'offset' 	: idc.GetFuncOffset(ea),
					'mnem'	 	: mnem,
					'type'		: optype,
					'direction'	: direction,
					'text'		: idc.GetDisasm(ea),
				}

			ea += idaapi.cmd.size

	def highlight(self, highlight=True, mnem=None, optype=None, direction=None, text=None):
		for (ea, info) in self.xrefs.iteritems():
			if mnem and info['mnem'] != mnem:
				highlight = False
			elif optype and info['optype'] != optype:
				highlight = False
			elif direction and info['direction'] != direction:
				highlight = False
			elif text and info['text'] != text:
				highlight = False

			if highlight:
				color = 0x00ff00
			else:
				color = idc.DEFCOLOR

			idc.SetColor(ea, idc.CIC_ITEM, color)

	def unhighlight(self):
		self.highlight(False)
		
	
class localizedxrefs_t(idaapi.plugin_t):
	flags = 0
	comment = "IDA Localized Xrefs"
	help = ""
	wanted_name = "Localized Xrefs"
	wanted_hotkey = ""

	DELIM = '-' * 86
	HEADER = '\nXrefs to %s from %s:'

	def init(self):
		self.menu_context = idaapi.add_menu_item("Jump/", "List local xrefs", "Alt-8", 0, self.run, (None,))
		return idaapi.PLUGIN_KEEP

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

	def run(self, arg):
		global localxrefs
		fmt = ''

		r = LocalXrefs()
		localxrefs = r

		offsets = r.xrefs.keys()
		offsets.sort()

		if r.highlighted:
			print self.HEADER % (r.highlighted, r.function)
			print self.DELIM
			
			for ea in offsets:
				info = r.xrefs[ea]
	
				if not fmt:
					fmt = "%%s   %%s   %%-%ds   %%s" % (len(info['offset']) + 15)

				print fmt % (info['direction'], info['type'], info['offset'], info['text'])
	
			print self.DELIM

def PLUGIN_ENTRY():
	return localizedxrefs_t()


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

import idc
import idaapi
import idautils

class NameMIPSSavedRegisters(object):

	INSIZE = 4
	SEARCH_DEPTH = 25

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

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

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

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

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

				mea += self.INSIZE

		print "done."


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

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

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

	def run(self, arg):
		pass

	def name_saved_registers(self, arg):
		NameMIPSSavedRegisters()

def PLUGIN_ENTRY():
	return mips_saved_registers_t()


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

import re
import idc
import idaapi
import idautils

# Global instance of MIPSROPFinder
mipsrop = None

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

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

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

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

		return string[:-1]

class ROPGadget(object):
	'''
	Class for storing information about a specific ROP gadget.
	'''
	
	def __init__(self, control, jump, operation=None, description="ROP gaget"):
		self.h = '-' * 112
		self.control = control
		self.exit = jump
		self.operation = operation
		self.description = description
		
		if self.control.opnd1:
			self.control.register = self.control.opnd1
		else:
			self.control.register = self.control.opnd0

		if self.exit.opnd1:
			self.exit.register = self.exit.opnd1
		else:
			self.exit.register = self.exit.opnd0
		
		if self.operation:
			if self.operation.ea < self.control.ea:
				self.entry = self.operation
			else:
				self.entry = self.control
		else:
			self.operation = self.control
			self.entry = self.control


	def header(self):
		return self.h + "\n|  Address     |  Action                                              |  Control Jump                          |\n" + self.h

	def footer(self):
		return self.h

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

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

	INSIZE = 4
	SEARCH_DEPTH = 25

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

	def build_code(self):
		keys = self.gadgets.keys()
		keys.sort()
	
		for key in keys[::-1]:
			last_instruction = False
			ea = self.gadgets[key]
			end_ea = ea + self.SEARCH_DEPTH

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

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

class MIPSROPFinder(object):
	'''
	Primary ROP finder class.
	'''

	CODE = 2
        DATA = 3
	INSIZE = 4
        SEARCH_DEPTH = 25
	
	def __init__(self):
		self.start = idc.BADADDR
		self.end = idc.BADADDR
		self.system_calls = []
		self.double_jumps = []
		self.controllable_jumps = []
		start = 0
		end = 0

		for (start, end) in self._get_segments(self.CODE):
			self.controllable_jumps += self._find_controllable_jumps(start, end)
			self.system_calls += self._find_system_calls(start, end)
			self.double_jumps += self._find_double_jumps(start, end)
			if self.start == idc.BADADDR:
				self.start = start
		self.end = end
	
		if self.controllable_jumps or self.system_calls:
			print "MIPS ROP Finder activated, found %d controllable jumps between 0x%.8X and 0x%.8X" % (len(self.controllable_jumps), self.start, self.end)
		
	def _get_segments(self, attr):
		segments = []
		start = idc.BADADDR
		end = idc.BADADDR
		seg = idc.FirstSeg()

		while seg != idc.BADADDR:
			if idc.GetSegmentAttr(seg, idc.SEGATTR_TYPE) == attr:
				start = idc.SegStart(seg)
				end = idc.SegEnd(seg)
				segments.append((start, end))
			seg = idc.NextSeg(seg)

		return segments

	def _get_instruction(self, ea):
		return MIPSInstruction(idc.GetMnem(ea), idc.GetOpnd(ea, 0), idc.GetOpnd(ea, 1), idc.GetOpnd(ea, 2), ea)

	def _does_instruction_match(self, ea, instruction, regex=False):
		i = 0
		op_cnt = 0
		op_ok_cnt = 0
		match = False
		ins_size = idaapi.decode_insn(ea)
		mnem = GetMnem(ea)

		if (not instruction.mnem) or (instruction.mnem == mnem) or (regex and re.match(instruction.mnem, mnem)):
			for operand in instruction.operands:
				if operand:
					op_cnt += 1
					op = idc.GetOpnd(ea, i)

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

			if op_cnt == op_ok_cnt:
				match = True

		return match

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

		if mnem and mnem[0] in bad_instructions:
			bad = True
		else:
			for register in no_clobber:
				if (idaapi.insn_t_get_canon_feature(idaapi.cmd.itype) & idaapi.CF_CHG1) == idaapi.CF_CHG1:
					if idc.GetOpnd(ea, 0) == register:
						bad = True

		return bad

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

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

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

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

			ea -= self.INSIZE

		return instruction_ea

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

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

			ea += self.INSIZE

		return instruction_ea

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

		for possible_control_instruction in t9_controls+ra_controls:
			ea = start_ea
			found = 0

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

			while ea <= end_ea:

				ea = self._find_next_instruction_ea(ea, possible_control_instruction, end_ea, regex=True)
				if ea != idc.BADADDR:
					ins_size = idaapi.decode_insn(ea)

					control_instruction = self._get_instruction(ea)
					control_register = control_instruction.operands[1]
					
					if control_register:
						for jump in jumps:
							jump_ea = self._find_next_instruction_ea(ea+ins_size, jump, end_ea, no_baddies=True, regex=True, dont_overwrite=musnt_clobber)
							if jump_ea != idc.BADADDR:
								jump_instruction = self._get_instruction(jump_ea)
								controllable_jumps.append(ROPGadget(control_instruction, jump_instruction, description="Controllable Jump"))
								ea = jump_ea
					
					ea += ins_size

		return controllable_jumps

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

		for xref in idautils.XrefsTo(idc.LocByName('system')):
			ea = xref.frm
			if ea >= start_ea and ea <= end_ea and idc.GetMnem(ea)[0] in ['j', 'b']:
				a0_ea = self._find_next_instruction_ea(ea+self.INSIZE, stack_arg_zero, ea+self.INSIZE)
				if a0_ea == idc.BADADDR:
					a0_ea = self._find_prev_instruction_ea(ea, stack_arg_zero, ea-(self.SEARCH_DEPTH*self.INSIZE))
				
				if a0_ea != idc.BADADDR:
					control_ea = self._find_prev_instruction_ea(ea-self.INSIZE, system_load, ea-(self.SEARCH_DEPTH*self.INSIZE))
					if control_ea != idc.BADADDR:
						system_calls.append(ROPGadget(self._get_instruction(control_ea), self._get_instruction(ea), self._get_instruction(a0_ea), description="System call"))

				ea += self.INSIZE
			else:
				break

		return system_calls

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

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

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

		return double_jumps

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

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

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

		return gadget_list

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

		for gadget in gadgets:
			print str(gadget)

		if gadgets:
			print gadgets[0].footer()
		
		print "Found %d matching gadgets" % (len(gadgets))

	def _get_marked_gadgets(self):
		rop_gadgets = {}

		for i in range(1, 1024):
			marked_pos = idc.GetMarkedPos(i)
			if marked_pos != idc.BADADDR:
				marked_comment = idc.GetMarkComment(i)
				if marked_comment and marked_comment.lower().startswith("rop"):
					rop_gadgets[marked_comment] = marked_pos
			else:
				break

		return rop_gadgets

	def double(self):
		self.doubles()

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

	def stackfinder(self):
		self.stackfinders()

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

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

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

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

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

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

		self._print_gadgets(tail_gadgets)

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

	def find(self, instruction_string=""):
		'''
		Locates all potential ROP gadgets that contain the specified instruction.

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

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

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

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

		instruction = MIPSInstruction(instruction_parts[0], instruction_parts[1], instruction_parts[2], instruction_parts[3])
		gadgets = self._find_rop_gadgets(instruction)
		if gadgets:
			self._print_gadgets(gadgets)
		else:
			print "No ROP gadgets found!"

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

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

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

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

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

					ea += self.INSIZE

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

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

				summaries.append(summary)

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

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

			print delim
			print line_fmt % (headings['name'], headings['offset'], headings['summary'])
			print delim
			
			for i in range(0, len(gadget_keys)):
				line_count = 0
				marked_comment = gadget_keys[i]
				offset = "0x%.8X" % rop_gadgets[marked_comment]
				summary = summaries[i]
				
				for line in summary:
					if line_count == 0:
						print line_fmt % (marked_comment, offset, line)
					else:
						print line_fmt % ('', '', line)

					line_count += 1

				print delim

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

	def help(self):
		'''
		Show help info.
		'''
		delim = "---------------------------------------------------------------------" * 2
		
		print ""
		print "mipsrop.find(instruction_string)"
		print delim
		print self.find.__doc__

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

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

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

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

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


class mipsropfinder_t(idaapi.plugin_t):
	flags = 0
	comment = "MIPS ROP Finder"
	help = ""
	wanted_name = "MIPS ROP Finder"
	wanted_hotkey = ""

	def init(self):
		self.menu_context = idaapi.add_menu_item("Search/", "mips rop gadgets", "Alt-1", 0, self.run, (None,))
		return idaapi.PLUGIN_KEEP

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

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

# DEBUG
#if __name__ == '__main__':
#	mipsrop = MIPSROPFinder()


########NEW FILE########
__FILENAME__ = idapathfinder
import idc
import idaapi
import idautils
import pathfinder
import time

class idapathfinder_t(idaapi.plugin_t):

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

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

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

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

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

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

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

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

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

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

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

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

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

            if not many:
                break

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

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

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

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

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

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

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

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

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

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

def PLUGIN_ENTRY():
    return idapathfinder_t()

########NEW FILE########
__FILENAME__ = install
#!/usr/bin/env python
# Simple installer script to drop the files where they need to be.

import sys
import shutil
import os.path

try:
	ida_dir = sys.argv[1]
except:
	print "Usage: %s <path to IDA install directory>" % sys.argv[0]
	sys.exit(1)

if os.path.exists(ida_dir):
	shutil.copyfile('pathfinder.py', os.path.join(ida_dir, 'python', 'pathfinder.py'))
	shutil.copyfile('idapathfinder.py', os.path.join(ida_dir, 'plugins', 'idapathfinder.py'))
	print "PathFinder installed to '%s'." % ida_dir
else:
	print "Install failed, '%s' does not exist!" % ida_dir
	sys.exit(1)

########NEW FILE########
__FILENAME__ = pathfinder
import idc
import idaapi
import idautils
import time

class History(object):
    '''
    Manages include/exclude graph history.
    '''

    INCLUDE_ACTION = 0
    EXCLUDE_ACTION = 1

    def __init__(self):
        self.reset()

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

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

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

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

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

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

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

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

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

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

class PathFinderGraph(idaapi.GraphViewer):
    '''
    Displays the graph and manages graph actions.
    '''

    def __init__(self, results, title="PathFinder Graph"):
        idaapi.GraphViewer.__init__(self, title)
        self.results = results

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

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

    def Show(self):
        '''
        Display the graph.

        Returns True on success, False on failure.
        '''
        if not idaapi.GraphViewer.Show(self):
            return False
        else:
            self.cmd_undo = self.AddCommand("Undo", "U")
            self.cmd_redo = self.AddCommand("Redo", "R")
            self.cmd_reset = self.AddCommand("Reset graph", "G")
            self.cmd_exclude = self.AddCommand("Exclude node", "X")
            self.cmd_include = self.AddCommand("Include node", "I")
            return True

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

        includes = self.history.get_includes()
        excludes = self.history.get_excludes()
        
        for path in self.results:
            parent_node = None

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

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

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

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

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

        return True

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

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

        return (self[node_id], color)

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

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

        return hint

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

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

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

        if self.edges.has_key(node_id):
            for edge_node_id in self.edges[node_id]:

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

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

        if xref_locations:
            xref_locations.sort()

            print ""
            print "Path Xrefs from %s:" % self[node_id]
            print "-" * 100
            for (xref_ea, dst_ea) in xref_locations:
                print "%-50s  =>  %s" % (self.get_name_by_ea(xref_ea), self.get_name_by_ea(dst_ea))
            print "-" * 100
            print ""
            
            idc.Jump(xref_locations[0][0])
        else:
            idc.Jump(node_ea)

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

    def get_ea_by_name(self, name):
        '''
        Get the address of a location by name.

        @name - Location name

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

        return ea

    def get_name_by_ea(self, ea):
        '''
        Get the name of the specified address.

        @ea - Address.

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

class PathFinder(object):
    '''
    Base class for finding the path between two addresses.
    '''

    # Limit the max recursion depth
    MAX_DEPTH = 500

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

        @destination - The end node ea.

        Returns None.
        '''
        self.tree = {}
        self.nodes = {}
        self.depth = 0
        self.last_depth = 0
        self.full_paths = []
        self.current_path = []
        self.destination = self._name2ea(destination)
        self.build_call_tree(self.destination)

    def __enter__(self):
        return self
        
    def __exit__(self, t, v, traceback):
        return

    def _name2ea(self, nea):
        if isinstance(nea, type('')):
            return idc.LocByName(nea)
        return nea

    def paths_from(self, source, exclude=[], include=[], xrefs=[], noxrefs=[]):
        '''
        Find paths from a source node to a destination node.

        @source  - The source node ea to start the search from.
        @exclude - A list of ea's to exclude from paths.
        @include - A list of ea's to include in paths.
        @xrefs   - A list of ea's that must be referenced from at least one of the path nodes.
        @noxrefs - A list of ea's that must not be referenced from any of the path nodes.

        Returns a list of path lists.
        '''
        paths = []
        good_xrefs = []
        bad_xrefs = []

        source = self._name2ea(source)

        # If all the paths from the destination node have not already
        # been calculated, find them first before doing anything else.
        if not self.full_paths:
            s = time.time()
            self.find_paths(self.destination, source)
            e = time.time()

        for xref in xrefs:
            xref = self._name2ea(xref)

            for x in idautils.XrefsTo(xref):
                f = idaapi.get_func(x.frm)
                if f:
                    good_xrefs.append(f.startEA)

        for xref in noxrefs:
            bad_xrefs.append(self._name2ea(xref))
            xref = self._name2ea(xref)

            for x in idautils.XrefsTo(xref):
                f = idaapi.get_func(x.frm)
                if f:
                    bad_xrefs.append(f.startEA)

        for p in self.full_paths:
            try:
                index = p.index(source)

                if exclude:
                    for ex in excludes:
                        if ex in p:
                            index = -1
                            break
                
                if include:
                    orig_index = index
                    index = -1

                    for inc in include:
                        if inc in p:
                            index = orig_index
                            break

                if good_xrefs:
                    orig_index = index
                    index = -1
                    
                    for xref in good_xrefs:
                        if xref in p:
                            index = orig_index

                    if index == -1:
                        print "Sorry, couldn't find", good_xrefs, "in", p

                if bad_xrefs:
                    for xref in bad_xrefs:
                        if xref in p:
                            index = -1
                            break

                # Be sure to include the destinatin and source nodes in the final path
                p = [self.destination] + p[:index+1]
                # The path is in reverse order (destination -> source), so flip it
                p = p[::-1]
                # Ignore any potential duplicate paths
                if p not in paths:
                    paths.append(p)
            except:
                pass

        return paths

    def find_paths(self, ea, source=None, i=0):
        '''
        Performs a depth-first (aka, recursive) search to determine all possible call paths originating from the specified location.
        Called internally by self.paths_from.

        @ea - The start node to find a path from.
        @i  - Used to specify the recursion depth; for internal use only.

        Returns None.
        '''
        # Increment recursion depth counter by 1
        i += 1
        # Get the current call graph depth
        this_depth = self.depth

        # If this is the first level of recursion and the call
        # tree has not been built, then build it.
        if i == 1 and not self.tree:
            self.build_call_tree(ea)

        # Don't recurse past MAX_DEPTH    
        if i >= self.MAX_DEPTH:
            return

        # Loop through all the nodes in the call tree, starting at the specified location
        for (reference, children) in self.nodes[ea].iteritems():
            # Does this node have a reference that isn't already listed in our current call path?
            if reference and reference not in self.current_path:
                    # Increase the call depth by 1
                    self.depth += 1
                    # Add the reference to the current path
                    self.current_path.append(reference)
                    # Find all paths from this new reference
                    self.find_paths(reference, source, i)

        # If we didn't find any additional references to append to the current call path (i.e., this_depth == call depth)
        # then we have reached the limit of this call path.
        if self.depth == this_depth:
            # If the current call depth is not the same as the last recursive call, and if our list of paths
            # does not already contain the current path, then append a copy of the current path to the list of paths
            if self.last_depth != self.depth and self.current_path and self.current_path not in self.full_paths:
                self.full_paths.append(list(self.current_path))
            # Decrement the call path depth by 1 and pop the latest node out of the current call path
            self.depth -= 1
            if self.current_path:
                self.current_path.pop(-1)

        # Track the last call depth
        self.last_depth = self.depth    

    def build_call_tree(self, ea):
        '''
        Performs a breadth first (aka, iterative) search to build a call tree to the specified address.

        @ea - The node to generate a tree for.

        Returns None.
        '''
        self.tree[ea] = {}
        self.nodes[ea] = self.tree[ea]
        nodes = [ea]

        while nodes:
            new_nodes = []

            for node in nodes:
                if node and node != idc.BADADDR:
                    node_ptr = self.nodes[node]

                    for reference in self.node_xrefs(node):
                        if reference not in self.nodes:
                            node_ptr[reference] = {}
                            self.nodes[reference] = node_ptr[reference]
                            new_nodes.append(reference)
                        elif not node_ptr.has_key(reference):
                            node_ptr[reference] = self.nodes[reference]
            
            nodes = new_nodes

    def node_xrefs(self, node):
        '''
        This must be overidden by a subclass to provide a list of xrefs.

        @node - The EA of the node that we need xrefs for.

        Returns a list of xrefs to the specified node.
        '''
        return []

class FunctionPathFinder(PathFinder):
    '''
    Subclass to generate paths between functions.
    '''

    def __init__(self, destination):
        # IDA 6.4 needs the extra import here, else idaapi is type None
        import idaapi
        func = idaapi.get_func(self._name2ea(destination))
        super(FunctionPathFinder, self).__init__(func.startEA)

    def node_xrefs(self, node):
        '''
        Return a list of function EA's that reference the given node.
        '''
        xrefs = []

        for x in idautils.XrefsTo(node):
            if x.type != idaapi.fl_F:
                f = idaapi.get_func(x.frm)
                if f and f.startEA not in xrefs:
                    xrefs.append(f.startEA)
        return xrefs
    
class BlockPathFinder(PathFinder):
    '''
    Subclass to generate paths between code blocks inside a function.
    '''

    def __init__(self, destination):
        func = idaapi.get_func(destination)
        self.blocks = idaapi.FlowChart(f=func)
        self.block_table = {}

        for block in self.blocks:
            self.block_table[block.startEA] = block
            self.block_table[block.endEA] = block

        self.source_ea = func.startEA
        dst_block = self.LookupBlock(destination)

        if dst_block:
            super(BlockPathFinder, self).__init__(dst_block.startEA)

    def LookupBlock(self, ea):
        try:
            return self.block_table[ea]
        except:
            for block in self.blocks:
                if ea >= block.startEA and ea < block.endEA:
                    return block
        return None
        
    def node_xrefs(self, node):
        '''
        Return a list of blocks that reference the provided block.
        '''
        xrefs = []

        block = self.LookupBlock(node)
        if block:
            for xref in idautils.XrefsTo(block.startEA):
                xref_block = self.LookupBlock(xref.frm)
                if xref_block and xref_block.startEA not in xrefs:
                    xrefs.append(xref_block.startEA)

        return xrefs

class Find(object):

    def __init__(self, start=[], end=[], include=[], exclude=[], xrefs=[], noxrefs=[]):
        self.start = self._obj2list(start)
        self.end = self._obj2list(end)
        self.include = self._obj2list(include)
        self.exclude = self._obj2list(exclude)
        self.xrefs = self._obj2list(xrefs)
        self.noxrefs = self._obj2list(noxrefs)

        if len(self.start) > 0:
            first_ea = self._obj2ea(self.start[0])
            func = idaapi.get_func(self.start[0])
            if func:
                results = []

                end_func = idaapi.get_func(self.end[0])
                if end_func and end_func.startEA == self.end[0]:
                    pfclass = FunctionPathFinder
                else:
                    pfclass = BlockPathFinder
                print pfclass
                
                for destination in self.end:
                    pf = pfclass(destination)
                    for source in self.start:
                        results += pf.paths_from(source, exclude=self.exclude, include=self.include, xrefs=self.xrefs, noxrefs=self.noxrefs)
                    del pf

                print "RESULTS:", results
                if results:
                    pg = PathFinderGraph(results)
                    pg.Show()
                    del pg

    def _obj2list(self, obj):
        '''
        Converts the supplied object to a list, if it is not already a list.

        @obj - The object.
    
        Returns a list.
        '''
        l = []

        if not isinstance(obj, type([])):
            l.append(self._obj2ea(obj))
        else:
            for o in obj:
                l.append(self._obj2ea(o))
        return l

    def _obj2ea(self, ea):
        if isinstance(ea, type('')):
            return idc.LocByName(ea)
        return ea

#if __name__ == "__main__":
    #Find(['main'], ['strcpy'])
    #Find('execute_other_requests', 'loc_408E80')


########NEW FILE########


repository: herosi/PyClassInformer
file: pyclassinformer_plugin.py
wanted_name: PyClassInformer
comment: Yet Another RTTI Parser
created_at: 2025-02-27 05:43:08+00:00
pushed_at: 2025-04-04 09:20:19+00:00
url:
https://github.com/herosi/PyClassInformer/blob/38d236ff1bb4c8acbaaf0c67439fe4247ac4c78d/pyclassinformer_plugin.py

content:
import ida_idaapi
import ida_kernwin

import os
import sys

dirpath = os.path.dirname(os.path.abspath(__file__))
script_dir = os.path.join(dirpath, "pyclassinformer")
if not os.path.isdir(script_dir):
    script_dir = os.path.join(dirpath, "..", "pyclassinformer")

ida_idaapi.require("pyclassinformer")
ida_idaapi.require("pyclassinformer.qtutils")

# for IDA 7.4 or earlier
try:
    g_flags = ida_idaapi.PLUGIN_MULTI
except AttributeError:
    g_flags = ida_idaapi.PLUGIN_DRAW

# for IDA 7.4 or earlier
try:
    g_obj = ida_idaapi.plugmod_t
except AttributeError:
    g_obj = object

g_plugmod_flag = False
if g_flags != ida_idaapi.PLUGIN_DRAW and g_obj != object:
    g_plugmod_flag = True

class pci_plugin_t(ida_idaapi.plugin_t):
    flags = g_flags
    comment = "Yet Another RTTI Parser"
    wanted_name = "PyClassInformer"
    wanted_hotkey = "Alt-Shift-L"
    help = "Press '" + wanted_hotkey + "' to display the " + wanted_name + " result."
    
    toolbar_displayed_name = wanted_name
    toolbar_name = toolbar_displayed_name + 'Toolbar'
    
    action_name = "pyclassinformer:execute"
    menu_path = "Edit/Plugins/"
    
    icon_data = open(os.path.join(script_dir, "pci_icon.png"), "rb").read()
    act_icon = ida_kernwin.load_custom_icon(data=icon_data, format="png")
    
    class exec_pci_action(ida_kernwin.action_handler_t):
        def __init__(self, plugin):
            ida_kernwin.action_handler_t.__init__(self)
            import weakref
            self.v = weakref.ref(plugin)
        
        def activate(self, ctx):
            self.v().plugin_mod.run(None)
            
        def update(self, ctx):
            return ida_kernwin.AST_ENABLE_ALWAYS
        
    def attach_to_menu_and_toolbar(self):
        # insert the action
        ida_kernwin.register_action(
            ida_kernwin.action_desc_t(
            self.action_name,
            self.comment,
            self.exec_pci_action(self),
            None,
            self.wanted_name,
            self.act_icon))
        
        # attach the action to menu
        ida_kernwin.attach_action_to_menu(
            self.menu_path,
            self.action_name,
            ida_kernwin.SETMENU_APP)

        # create the toolbar and attach the action to the toolbar
        ida_kernwin.create_toolbar(self.toolbar_name, self.toolbar_displayed_name)
        ida_kernwin.attach_action_to_toolbar(self.toolbar_name, self.action_name)
        
        # install ui hook to enable toolbar later
        self.ph = pyclassinformer.qtutils.enable_toolbar_t(self.toolbar_name)
        
    @staticmethod
    class register_icon(ida_kernwin.UI_Hooks):
        def updated_actions(self):
            if ida_kernwin.update_action_icon(pci_plugin_t.menu_path + pci_plugin_t.wanted_name, pci_plugin_t.act_icon):
                # unhook this if it's successful
                self.unhook()

    def init(self):
        ida_kernwin.msg("############### %s (%s) ###############%s" % (self.wanted_name, self.comment, os.linesep))
        ida_kernwin.msg("%s%s" % (self.help, os.linesep))

        # attach action to menu and toolbar
        self.attach_to_menu_and_toolbar()
        
        r = self.flags
        self.plugin_mod = pci_plugmod_t()
        if g_plugmod_flag:
            r = self.plugin_mod
        return r
    
    # for old IDA til 7.6
    def run(self, arg):
        self.plugin_mod.run(arg)
        
    # for old IDA til 7.6
    def term(self):
        self.plugin_mod.term()
        

class pci_plugmod_t(g_obj):
    toolbar_name = pci_plugin_t.toolbar_name
    menu_path = pci_plugin_t.menu_path
    action_name = pci_plugin_t.action_name
    act_icon = pci_plugin_t.act_icon
    
    def __del__(self):
        self.term()
        
    def run(self, arg):
        pci_plugmod_t.run_pci(icon=self.act_icon)
        
    def term(self):
        self.detatch_from_menu_and_toolbar()
        
    def detatch_from_menu_and_toolbar(self):
        ida_kernwin.detach_action_from_toolbar(self.toolbar_name, self.action_name)
        ida_kernwin.delete_toolbar(self.toolbar_name)
        ida_kernwin.detach_action_from_menu(self.menu_path, self.action_name)
        ida_kernwin.free_custom_icon(self.act_icon)
        ida_kernwin.unregister_action(self.action_name)
        
    @staticmethod
    def run_pci(icon=-1):
        ida_idaapi.require("pyclassinformer.pci_config_form")
        config = pyclassinformer.pci_config_form.pci_form_t.show()
        if config is not None:
            ida_idaapi.require("pyclassinformer.pyclassinformer")
            pyclassinformer.pyclassinformer.run_pci(config=config, icon=icon)
        else:
            print("PyClassInformer: Canceled")


def PLUGIN_ENTRY():
    return pci_plugin_t()


# install a UI hook to add icon in the plugin menu
ri = pci_plugin_t.register_icon()
ri.hook()


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

content:
"""
summary: show ctree graph

description:
  Registers an action that can be used to show the graph of the ctree.
  The current item will be highlighted in the graph.

  The command shortcut is `Ctrl+Shift+G`, and is also added
  to the context menu.

  To display the graph, we produce a .gdl file, and
  request that ida displays that using `ida_gdl.display_gdl`.
"""

from __future__ import print_function

import ida_pro
import ida_hexrays
import ida_kernwin
import ida_gdl
import ida_lines
import ida_idaapi

ACTION_NAME = "vds5.py:displaygraph"
ACTION_SHORTCUT = "Ctrl+Shift+G"

CL_WHITE            = ((255)+  (255<<8)+  (255<<16)) #   0
CL_BLUE             = ((0  )+  (0  <<8)+  (255<<16)) #   1
CL_RED              = ((255)+  (0  <<8)+  (0  <<16)) #   2
CL_GREEN            = ((0  )+  (255<<8)+  (0  <<16)) #   3
CL_YELLOW           = ((255)+  (255<<8)+  (0  <<16)) #   4
CL_MAGENTA          = ((255)+  (0  <<8)+  (255<<16)) #   5
CL_CYAN             = ((0  )+  (255<<8)+  (255<<16)) #   6
CL_DARKGREY         = ((85 )+  (85 <<8)+  (85 <<16)) #   7
CL_DARKBLUE         = ((0  )+  (0  <<8)+  (128<<16)) #   8
CL_DARKRED          = ((128)+  (0  <<8)+  (0  <<16)) #   9
CL_DARKGREEN        = ((0  )+  (128<<8)+  (0  <<16)) #  10
CL_DARKYELLOW       = ((128)+  (128<<8)+  (0  <<16)) #  11
CL_DARKMAGENTA      = ((128)+  (0  <<8)+  (128<<16)) #  12
CL_DARKCYAN         = ((0  )+  (128<<8)+  (128<<16)) #  13
CL_GOLD             = ((255)+  (215<<8)+  (0  <<16)) #  14
CL_LIGHTGREY        = ((170)+  (170<<8)+  (170<<16)) #  15
CL_LIGHTBLUE        = ((128)+  (128<<8)+  (255<<16)) #  16
CL_LIGHTRED         = ((255)+  (128<<8)+  (128<<16)) #  17
CL_LIGHTGREEN       = ((128)+  (255<<8)+  (128<<16)) #  18
CL_LIGHTYELLOW      = ((255)+  (255<<8)+  (128<<16)) #  19
CL_LIGHTMAGENTA     = ((255)+  (128<<8)+  (255<<16)) #  20
CL_LIGHTCYAN        = ((128)+  (255<<8)+  (255<<16)) #  21
CL_LILAC            = ((238)+  (130<<8)+  (238<<16)) #  22
CL_TURQUOISE        = ((64 )+  (224<<8)+  (208<<16)) #  23
CL_AQUAMARINE       = ((127)+  (255<<8)+  (212<<16)) #  24
CL_KHAKI            = ((240)+  (230<<8)+  (140<<16)) #  25
CL_PURPLE           = ((160)+  (32 <<8)+  (240<<16)) #  26
CL_YELLOWGREEN      = ((154)+  (205<<8)+  (50 <<16)) #  27
CL_PINK             = ((255)+  (192<<8)+  (203<<16)) #  28
CL_ORANGE           = ((255)+  (165<<8)+  (0  <<16)) #  29
CL_ORCHID           = ((218)+  (112<<8)+  (214<<16)) #  30
CL_BLACK            = ((0  )+  (0  <<8)+  (0  <<16)) #  31

COLORS_LUT = {
    CL_WHITE             : "white",
    CL_BLUE              : "blue",
    CL_RED               : "red",
    CL_GREEN             : "green",
    CL_YELLOW            : "yellow",
    CL_MAGENTA           : "magenta",
    CL_CYAN              : "cyan",
    CL_DARKGREY          : "darkgrey",
    CL_DARKBLUE          : "darkblue",
    CL_DARKRED           : "darkred",
    CL_DARKGREEN         : "darkgreen",
    CL_DARKYELLOW        : "darkyellow",
    CL_DARKMAGENTA       : "darkmagenta",
    CL_DARKCYAN          : "darkcyan",
    CL_GOLD              : "gold",
    CL_LIGHTGREY         : "lightgrey",
    CL_LIGHTBLUE         : "lightblue",
    CL_LIGHTRED          : "lightred",
    CL_LIGHTGREEN        : "lightgreen",
    CL_LIGHTYELLOW       : "lightyellow",
    CL_LIGHTMAGENTA      : "lightmagenta",
    CL_LIGHTCYAN         : "lightcyan",
    CL_LILAC             : "lilac",
    CL_TURQUOISE         : "turquoise",
    CL_AQUAMARINE        : "aquamarine",
    CL_KHAKI             : "khaki",
    CL_PURPLE            : "purple",
    CL_YELLOWGREEN       : "yellowgreen",
    CL_PINK              : "pink",
    CL_ORANGE            : "orange",
    CL_ORCHID            : "orchid",
    CL_BLACK             : "black",
}

def get_color_name(c):
    return COLORS_LUT[c] if c in COLORS_LUT.keys() else "?"

class cfunc_graph_t: # alas we can't inherit gdl_graph_t
    def __init__(self, highlight):
        self.items = [] # list of citem_t
        self.highlight = highlight
        self.succs = [] # list of lists of next nodes
        self.preds = [] # list of lists of previous nodes

    def nsucc(self, n):
        return len(self.succs[n]) if self.size() else 0

    def npred(self, n):
        return len(self.preds[n]) if self.size() else 0

    def succ(self, n, i):
        return self.succs[n][i]

    def pred(self, n, i):
        return self.preds[n][i]

    def size(self):
        return len(self.preds)

    def add_node(self):
        n = self.size()

        def resize(array, new_size):
            if new_size > len(array):
                while len(array) < new_size:
                    array.append([])
            else:
                array = array[:new_size]
            return array

        self.preds = resize(self.preds, n+1)
        self.succs = resize(self.succs, n+1)
        return n

    def add_edge(self, x, y):
        self.preds[y].append(x)
        self.succs[x].append(y)

    def get_expr_name(self, expr):
        name = expr.print1(None)
        name = ida_lines.tag_remove(name)
        name = ida_pro.str2user(name)
        return name

    def get_node_label(self, n):
        item = self.items[n]
        op = item.op
        insn = item.cinsn
        expr = item.cexpr
        parts = [ida_hexrays.get_ctype_name(op)]
        if op == ida_hexrays.cot_ptr:
            parts.append(".%d" % expr.ptrsize)
        elif op == ida_hexrays.cot_memptr:
            parts.append(".%d (m=%d)" % (expr.ptrsize, expr.m))
        elif op == ida_hexrays.cot_memref:
            parts.append(" (m=%d)" % (expr.m,))
        elif op in [
                ida_hexrays.cot_obj,
                ida_hexrays.cot_var]:
            name = self.get_expr_name(expr)
            parts.append(".%d %s" % (expr.refwidth, name))
        elif op in [
                ida_hexrays.cot_num,
                ida_hexrays.cot_helper,
                ida_hexrays.cot_str]:
            name = self.get_expr_name(expr)
            parts.append(" %s" % (name,))
        elif op == ida_hexrays.cit_goto:
            parts.append(" LABEL_%d" % insn.cgoto.label_num)
        elif op == ida_hexrays.cit_asm:
            parts.append("<asm statements; unsupported ATM>")
            # parts.append(" %a.%d" % ())
        parts.append("\\n")
        parts.append("ea: %08X" % item.ea)
        if item.is_expr() and not expr.type.empty():
            parts.append("\\n")
            tstr = expr.type._print()
            parts.append(tstr if tstr else "?")
        return "".join(parts)

    def get_node_color(self, n):
        item = self.items[n]
        if self.highlight is not None and item.obj_id == self.highlight.obj_id:
            return CL_GREEN
        return None

    def gen_gdl(self, fname):
        with open(fname, "wb") as out:
            def write_out(s):
                out.write(s.encode("UTF-8"))
            write_out("graph: {\n")

            write_out("// *** nodes\n")
            for n in range(len(self.items)):
                item = self.items[n]
                node_label = self.get_node_label(n)
                node_props = [""]
                if n == 0:
                    node_props.append("vertical_order: 0")
                color = self.get_node_color(n)
                if color is not None:
                    node_props.append("color: %s" % get_color_name(color))
                write_out("""node: { title: "%d" label: "%d: %s" %s}\n""" % (
                    n,
                    n,
                    node_label,
                    " ".join(node_props)))

            write_out("// *** edges\n")
            for n in range(len(self.items)):
                item = self.items[n]
                write_out("// edges %d -> ?\n" % n)
                for i in range(self.nsucc(n)):
                    t = self.succ(n, i)
                    label = ""
                    if item.is_expr():
                        target = self.items[t]
                        if item.x is not None and item.x == target:
                            label = "x"
                        elif item.y is not None and item.y == target:
                            label = "y"
                        elif item.z is not None and item.z == target:
                            label = "z"
                        if label:
                            label = """ label: "%s" """ % label
                    write_out("""edge: { sourcename: "%s" targetname: "%s"%s}\n""" % (
                        str(n), str(t), label))

            write_out("}\n")

    def dump(self):
        print("%d items:" % len(self.items))
        for i in self.items:
            print("\t%s (%08x)" % (i, i.ea))

        print("succs:")
        for s in self.succs:
            print("\t%s" % s)

        print("preds:")
        for p in self.preds:
            print("\t%s" % p)


class graph_builder_t(ida_hexrays.ctree_parentee_t):

    def __init__(self, cg):
        ida_hexrays.ctree_parentee_t.__init__(self)
        self.cg = cg
        self.reverse = [] # (citem_t, node#) tuples

    def add_node(self, i):
        for k, _ in self.reverse:
            if i.obj_id == k.obj_id:
                ida_kernwin.warning("bad ctree - duplicate nodes! (i.ea=%x)" % i.ea)
                self.cg.dump()
                return -1

        n = self.cg.add_node()
        if n <= len(self.cg.items):
            self.cg.items.append(i)
        self.cg.items[n] = i
        self.reverse.append((i, n))
        return n

    def process(self, i):
        n = self.add_node(i)
        if n < 0:
            return n
        if len(self.parents) > 1:
            lp = self.parents.back().obj_id
            for k, v in self.reverse:
                if k.obj_id == lp:
                    p = v
                    break
            self.cg.add_edge(p, n)
        return 0

    def visit_insn(self, i):
        return self.process(i)

    def visit_expr(self, e):
        return self.process(e)


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

    def activate(self, ctx):
        vu = ida_hexrays.get_widget_vdui(ctx.widget)
        vu.get_current_item(ida_hexrays.USE_KEYBOARD)
        highlight = vu.item.e if vu.item.is_citem() else None

        cg = cfunc_graph_t(highlight)
        gb = graph_builder_t(cg)
        gb.apply_to(vu.cfunc.body, None)

        import tempfile
        fname = tempfile.mktemp(suffix=".gdl")
        cg.gen_gdl(fname)
        ida_gdl.display_gdl(fname)
        return 1

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


class vds5_hooks_t(ida_hexrays.Hexrays_Hooks):
    def populating_popup(self, widget, handle, vu):
        ida_kernwin.attach_action_to_popup(vu.ct, None, ACTION_NAME)
        return 0

# a plugin interface, boilerplate code
class my_plugin_t(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_HIDE
    wanted_name = "Hex-Rays show C graph (IDAPython)"
    wanted_hotkey = ""
    comment = "Sample plugin5 for Hex-Rays decompiler"
    help = ""
    def init(self):
        if ida_hexrays.init_hexrays_plugin():
            ida_kernwin.register_action(
                ida_kernwin.action_desc_t(
                    ACTION_NAME,
                    "Hex-Rays show C graph (IDAPython)",
                    display_graph_ah_t(),
                    ACTION_SHORTCUT))
            self.vds5_hooks = vds5_hooks_t()
            self.vds5_hooks.hook()
            return ida_idaapi.PLUGIN_KEEP # keep us in the memory
    def term(self):
        self.vds5_hooks.unhook()
    def run(self, arg):
        pass

def PLUGIN_ENTRY():
    return my_plugin_t()


repository: techbliss/ida_pro_http_ip_geolocator
file: 7.0+/plugins/Geolocateloader.py
description: Google maps http and ip lookup for ida pro
wanted_name: Geo locate
comment: Run me
created_at: 2017-08-05 19:28:41+00:00
pushed_at: 2019-03-10 17:07:43+00:00
url:
https://github.com/techbliss/ida_pro_http_ip_geolocator/blob/56bb56374ac88b9b344bfcfba8c274ad7cdb4173/7.0%2B/plugins/Geolocateloader.py

content:
# Created by: Storm Shadow http://www.techbliss.org
import os
import ida_idaapi, ida_kernwin
import idc
from idc import *
from idaapi import *
import sys
sys.path.insert(0 , idaapi.idadir("plugins\\geo\\icons")) #change path to the icon
sys.path.insert(0 , idaapi.idadir("plugins"))
sys.path.insert(0, idaapi.idadir("plugins\\geo\\"))
import ico
from ico import *
import subprocess

PLUGIN_VERSION = "1.2" #change Plugin Version
IDAVERISONS    = "IDA PRO 7.0+" #change Ida version
AUTHORS        = "Storm Shadow" #change author
DATE           = "2019" #change date
TWITTER        = "Twitter @zadow28" #change social media

def banner():
    banner_options = (PLUGIN_VERSION, AUTHORS, DATE, TWITTER, IDAVERISONS)
    banner_titles = "Geo Locator v%s - (c) %s - %s - %s - %s" % banner_options #change Python editor

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

banner()

# 1) Create the handler class
class MyEditorHandler(idaapi.action_handler_t): # the class name must be unique for each plugin nb same as line (53)
    def __init__(self):
        idaapi.action_handler_t.__init__(self)

    # Run editor when invoked.
    def activate(self, ctx):
		g = globals()
		idahome = idaapi.idadir("plugins\\geo") #change to set plguin path where the main plguin script is
		process = subprocess.Popen(['python', idahome + "\\geo.py"])
		#print os.getcwd()		#change for the main plugin script

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

class geoeye(idaapi.plugin_t): #change class to unique name for each plugin , also change line (103)
    flags = idaapi.PLUGIN_FIX #different flags this is for plugin visible at startup
    comment = "Run me" #help me text
    help = "Geo helper" #help text
    wanted_name = "Geo locate" #change the plugins name
    wanted_hotkey = "" #the tooltip goes away in menu when setting it here ,DONT DO it! and only is shown in File/Plugins menu


    def locate_menuaction(self): #change for something unique.also change line (90)
        action_desc = idaapi.action_desc_t(
            'my:locateaction',  # The action name. This acts like an ID and must be unique same as line (64), (68)
            'Geo locate!',  # The action text.
            MyEditorHandler(),  # The action handler must be unique , also change line (99)
            'Ctrl+G',  # Optional: the action shortcut DO IT  HERE!
            'Geo locate',  # Optional: the action tooltip (available in menus/toolbar)
            idaapi.load_custom_icon(":/ico/python.png")  # hackish load action icon , if no custom icon use number from 1-150 from internal ida
        )

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

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

        form = idaapi.get_current_tform()
        idaapi.attach_action_to_popup(form, None, "my:locateaction", None) # The action ID (see line(51)

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

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

        return PLUGIN_KEEP

    def _install_plugin(self):
        """
        Initialize & integrate the plugin into IDA.
        """
        self.locate_menuaction() #same as line (49)
        self._init()

    def term(self):
        pass

    def run(self, arg = 0):
        #we need the calls again if we wanna load it via File/Plugins/editor
        idaapi.msg("Python Geo loacte Loaded to menu \n use Alt+G hot key to quick load ")
        hackish = MyEditorHandler() #must be the same as line (53) change also hackish =  must be unique also line (100)
        hackish.activate(self)  #hackish must the same as line hackish = (99)

def PLUGIN_ENTRY():
    return geoeye() #must be unique for each plugin and same as line 43


repository: techbliss/ADB_Helper_For_ida_Pro
file: ADB.py
description: Helps you debugging Android in Ida pro
wanted_name: ADB Helper
comment: This is a comment
created_at: 2014-12-30 12:18:08+00:00
pushed_at: 2015-01-26 21:25:21+00:00
url:
https://github.com/techbliss/ADB_Helper_For_ida_Pro/blob/6d36953ae60f55849ffb8d41ff1aa1be5b15b772/ADB.py

content:
__author__ = 'zadow'

import subprocess
import idaapi
import idc
from idc import *
from idaapi import *
import os
import sys


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

    help = "ADB"
    wanted_name = "ADB Helper"
    wanted_hotkey = "Alt-F7"

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

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

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

    def AddMenuElements(self):
        idaapi.add_menu_item("Debugger/", "1: adb Restart as Root", "", 0, self.night, ())
        idaapi.add_menu_item("Debugger/", "2: push android server to rooted phone", "", 0, self.sun, ())
        idaapi.add_menu_item("Debugger/", "3: chmod android server 755 ", "", 0, self.rain, ())
        idaapi.add_menu_item("Debugger/", "4: connect to android server ", "", 0, self.shinny, ())
        idaapi.add_menu_item("Debugger/", "5: adb Forward ports for debugging", "", 0, self.day, ())
        idaapi.add_menu_item("Debugger/", "6: adb Logcat", "", 0, self.snows, ())
        idaapi.add_menu_item("Debugger/", "7: adb Logcat to File", "", 0, self.thunder, ())
        idaapi.add_menu_item("Debugger/", "8: adb shell >>> :", "", 0, self.cloudy, ())



    def run(self, arg=0):
        idaapi.msg("")

        self.AddMenuElements()


    def night(self):
        subprocess.Popen('adb.exe root')
        print 'Restart abd connection as root'

    def sun(self):
        subprocess.Popen('adb.exe push android_server /data/local/android_server')
        print 'Pushed'

    def rain(self):
        subprocess.Popen('adb.exe shell chmod 755 /data/local/android_server')
        print 'Permissions changed'


    def shinny(self):
        subprocess.Popen('adb.exe shell ./data/local/android_server')
        print 'Should be connected'


    def day(self):
        subprocess.Popen('adb forward tcp:23946 tcp:23946')
        print 'You can Attach or start debugging'

    def snows(self):
        subprocess.Popen('adb logcat')
        print 'Logging'

    def thunder(self):
        file = get_root_filename()
        subprocess.Popen('adb logcat > logcat.txt')

    def cloudy(self):
        subprocess.Popen('adb shell')

def PLUGIN_ENTRY():
    return gonzola()


repository: wallds/NoVmpy
file: novmpy.py
url: https://github.com/wallds/NoVmpy/blob/ffdbfcf4b9f39ac598263a8f3468ff79617f09d1/novmpy.py
wanted_name: NoVmpy
created_at: 2022-02-02 00:41:42+00:00
pushed_at: 2025-01-01 07:13:21+00:00
content:
# idapython fix <= 7.6
# DO NOT REMOVE ME
import sys
sys.stdout.encoding = 'utf-8'


# for test
# ida_loader.load_plugin('F:/NoVmpy/novmpy.py')
import idaapi

NOVMPY_VERSION = "0.1"


class NoVmpyPlugin(idaapi.plugin_t):
    flags = 0
    comment = ""
    help = ""
    wanted_name = "NoVmpy"
    wanted_hotkey = ""

    def __init__(self):
        super(NoVmpyPlugin, self).__init__()

    def init(self):
        from novmpy.ui import UIManager
        self.ui = UIManager()

        return idaapi.PLUGIN_KEEP

    def run(self, args):
        pass

    def term(self):
        from novmpy.handler import vm_handlers
        vm_handlers.clear()


def PLUGIN_ENTRY():
    return NoVmpyPlugin()


repository: arizvisa/ida-minsc
file: plugins/minsc.py
wanted_name: About Minsc
comment: Makes IDAPython Not Suck Completely.
created_at: 2014-09-02 08:25:13+00:00
pushed_at: 2025-04-25 19:39:39+00:00
description:
A plugin based on IDAPython for a functional DWIM interface. Current development against most recent IDA is in the "persistence-refactor" branch, ancient (but stable) work is in "master", so... create an issue if you want/need something backported. Use "Wiki" or "Discussions" for examples, and smash that "Star" button if you like this.

url:
https://github.com/arizvisa/ida-minsc/blob/fb90d75a9ce5d649323d3eb2f27301960ddd4cf2/plugins/minsc.py

content:
"""
IDA-Minsc plugin -- https://arizvisa.github.io/ida-minsc

This file contains the entirety of the logic that is used to load the
plugin. The plugin is (mostly) a library that aims to simplify IDAPython.
However, it utilizes hooks and keyboard shortcuts in a variety of ways in
order to keep track of the changes that the user may make within their
database. Essentially the goal of this plugin is to make absolutely
_everything_ that's useful in the database serializeable and queryable
so that it can be exchanged with things outside of the database.

If you wish to change the directory that the plugin is loaded from, specify the
location of the plugin's git repository in the variable that is marked below.
"""
import sys, os, logging
import builtins, six, imp, fnmatch, ctypes, types
import idaapi

# :: Point this variable at the directory containing the repository of the plugin ::
root = idaapi.get_user_idadir()

# The following classes contain the pretty much all of the loader logic used by the
# plugin. This is being defined here so that this file can also be used as a module.
class internal_api(object):
    """
    Loader base-class for any api that's based on files contained within a directory.
    """
    os, imp, fnmatch = os, imp, fnmatch
    def __init__(self, directory, **attributes):
        '''Initialize the api using the contents within the specified `directory`.'''
        self.path = self.os.path.realpath(directory)
        [ setattr(self, name, attribute) for name, attribute in attributes.items() ]

    ### Api operations
    def load_api(self, path):
        '''Load the specified `path` into a module that can be used.'''
        path, filename = self.os.path.split(path)
        name, _ = self.os.path.splitext(filename)
        return self.imp.find_module(name, [path])

    def iterate_api(self, include='*.py', exclude=None, **attributes):
        """Iterate through all of the files in the directory specified when initializing the loader.

        The `include` string is a glob that specifies which files are part of the loader's api.
        If the `exclude` glob is specified, then exclude files that match it from the loader api.
        """
        result = []
        for filename in self.fnmatch.filter(self.os.listdir(self.path), include):
            if exclude and self.fnmatch.fnmatch(filename, exclude):
                continue

            path = self.os.path.join(self.path, filename)
            _, ext = self.os.path.splitext(filename)

            left, right = (None, None) if include == '*' else (include.index('*'), len(include) - include.rindex('*'))
            modulename = filename[left : -right + 1]
            yield modulename, path
        return

    def new_api(self, modulename, path):
        '''Load the file found at `path` into the specified `modulename`.'''
        file, path, description = self.load_api(path)
        try:
            return self.imp.load_module(modulename, file, path, description)
        finally: file.close()

    ### Module operations
    def new_module(self, fullname, documentation=None):
        '''Create a new module (empty) with the specified `fullname` and the provided `documentation`.'''
        res = self.imp.new_module(fullname)
        res.__package__ = fullname
        res.__doc__ = documentation or ''
        return res

    def find_module(self, fullname, path=None):
        raise NotImplementedError

    def load_module(self, fullname):
        raise NotImplementedError

    def __iter__(self):
        '''Yield the full path of each module that is provided by this class.'''
        return
        yield

class internal_path(internal_api):
    """
    Loader class which provides all api composed of all of the files within a directory
    as modules that can always be imported from anywhere.
    """
    def __init__(self, path, **attrs):
        '''Initialize the loader using the files from the directory specified by `path`.'''
        super(internal_path, self).__init__(path)
        attrs.setdefault('include', '*.py')
        self.attrs, self.cache = attrs, { name : path for name, path in self.iterate_api(**attrs) }

    def find_module(self, fullname, path=None):
        '''If the module with the name `fullname` matches one of the files handled by our api, then act as their loader.'''
        return self if path is None and fullname in self.cache else None

    def load_module(self, fullname):
        '''Iterate through all of the modules that we can handle, and then load it if we've been asked.'''
        self.cache = { name : path for name, path in self.iterate_api(**self.attrs) }
        if fullname not in self.cache:
            raise ImportError("path-loader ({:s}) was not able to find a module named `{:s}`".format(self.path, fullname))
        return self.new_api(fullname, self.cache[fullname])

    def __iter__(self):
        '''Yield each of the available modules.'''
        for name, _ in self.iterate_api(**self.attrs):
            yield name
        return

class internal_submodule(internal_api):
    """
    Loader class which provides an api composed of all of the files within a
    directory, and binds them to a module which is used to access them.
    """
    sys = sys
    def __init__(self, __name__, path, **attrs):
        '''Initialize the loader using `__name__` as the name of the submodule using the files underneath the directory `path`.'''
        super(internal_submodule, self).__init__(path)
        attrs.setdefault('include', '*.py')
        self.__name__, self.attrs = __name__, attrs

    def find_module(self, fullname, path=None):
        '''If the module with the name `fullname` matches our submodule name, then act as its loader.'''
        return self if path is None and fullname == self.__name__ else None

    def new_api(self, modulename, path):
        '''Load the file found at the specified `path` as a submodule with the specified `modulename`.'''
        cls, fullname = self.__class__, '.'.join([self.__name__, modulename])
        res = super(cls, self).new_api(fullname, path)
        res.__package__ = self.__name__
        return res

    def load_module(self, fullname):
        '''Iterate through all of the modules that we can handle, load the submodule with them, and return it.'''
        module = self.sys.modules[fullname] = self.new_module(fullname)
        # FIXME: make module a lazy-loaded object for fetching module-code on-demand

        # Build a temporary cache for the module names and paths to load the api,
        # and use them to build their documentation.
        cache = { name : path for name, path in self.iterate_api(**self.attrs) }
        maximum = max(map(len, cache)) if cache else 0
        documentation = '\n'.join("{:<{:d}s} : {:s}".format(name, maximum, path) for name, path in sorted(cache.items()))
        documentation = '\n\n'.join([self.attrs['__doc__'], documentation]) if '__doc__' in self.attrs else documentation
        module.__doc__ = documentation

        # Load each submodule that composes the api, and attach it to the returned submodule.
        stack, count, result = [item for item in cache.items()], len(cache), {}
        while stack and count > 0:
            name, path = stack.pop(0)

            # Take the submodule name we popped off of the cache, and try and load it.
            # If we were able to load it successfully, then we just need to attach the
            # loaded code as a submodule of the object we're going to return.
            try:
                res = self.new_api(name, path)
                modulename = '.'.join([res.__package__, name])

            # If an exception was raised, then remember it so that we can let the user
            # know after we've completely loaded the module.
            except Exception as E:
                import logging
                logging.info("{:s} : Error trying to import module `{:s}` from {!s}. Queuing it until later.".format(self.__name__, name, path), exc_info=True)

                # If we caught an exception while trying to import the module, then stash
                # our exception info state into a dictionary and decrease a counter. This
                # is strictly to deal with module recursion issues in Python3.
                result[name], count = sys.exc_info(), count - 1
                stack.append((name, path))

            # Add the submodule that we loaded into the module that we're going to return.
            else:
                setattr(module, name, res)
            continue

        # If we weren't able to load one of the submodules that should've been in our cache,
        # then go through all of our backtraces and log the exception that was raised.
        if stack:
            import logging, traceback
            for name, path in stack:
                logging.fatal("{:s} : Error trying to import module `{:s}` from {!s}.".format(self.__name__, name, path), exc_info=result[name])
            return module

        # If we caught an exception despite our stack being empty, then this is because of a
        # recursion issue. In case someone wants to track these situations down, we go through
        # our caught exceptions and create some logging events with the backtrace. These errors
        # are non-fatal because importing another sub-module helped resolve it.
        import logging
        for name, exc_info in result.items():
            logging.info("{!s} : Encountered a non-fatal exception while trying to import recursive module `{:s}` from {!s}".format(self.__name__, name, cache[name]), exc_info=result[name])

        # Return the module that we just created.
        return module

    def __iter__(self):
        '''Yield each of the available modules.'''
        for name, _ in self.iterate_api(**self.attrs):
            yield name
        return

class internal_object(object):
    """
    Loader class which will simply expose an object instance as the module.
    """
    sys = sys
    def __init__(self, __name__, object, **attributes):
        '''Initialize the loader with the specified `__name__` and returning the provided `object` as its module.'''
        self.__name__, self.object = __name__, object
        for name, value in attributes.items():
            try: setattr(object, name, value)
            except: continue

    def find_module(self, fullname, path=None):
        '''If the module being searched for matches our `fullname`, then act as its loader.'''
        return self if path is None and fullname == self.__name__ else None

    def load_module(self, fullname):
        '''Return the specific object for the module specified by `fullname`.'''
        if fullname != self.__name__:
            raise ImportError("object-loader ({:s}) was not able to find a module named `{:s}`".format(self.__name__, fullname))
        module = self.sys.modules[fullname] = self.object
        return module

class plugin_module(object):
    """
    Loader class which iterates through all of the files in a directory, and
    manually initializes each plugin similar to the way `idaapi.plugin_t` is
    supposed to be initialized.
    """
    def __init__(self, path, **attrs):
        # FIXME: go through all files in plugin/ and call PLUGIN_ENTRY() on each module
        #        this should return an idaapi.plugin_t.

        # idaapi.plugin_t will contain an init, run, and term method.
        # also, are some attributes to process:
        # 'wanted_name' which is for idc.
        # 'wanted_hotkey', which should be mapped to a keypress.
        # 'comment' self-explanatory
        # 'help' self-explanatory

        # hotkey can be done by:
        # idaapi.CompileLine('static myname() { RunPythonStateMent("CallSomePython()") }')
        # idc.AddHotKey(module.wanted_hotkey, "myname")

        # idaapi.require
        pass

# The following logic is responsible for replacing a namespace with
# the contents namespace that represents the entirety of the plugin.
def load(namespace, preserve=()):
    module = imp.load_source('__root__', os.path.join(root, '__root__.py'))

    # save certain things within the namespace
    preserved = {symbol : value for symbol, value in namespace.items() if symbol in preserve}

    # empty out the entire namespace so that we can replace it
    [namespace.pop(symbol) for symbol in namespace.copy() if not symbol.startswith('__')]

    # re-populate with the root namespace while restoring any symbols
    # that needed to be preserved.
    namespace.update({symbol : value for symbol, value in module.__dict__.items() if not symbol.startswith('__')})
    namespace.update({symbol : value for symbol, value in preserved.items()})

# Just a ctypes wrapper so that we can access the internal IDA api.
library = ctypes.WinDLL if os.name == 'nt' else ctypes.CDLL

# The following code is responsible for seeding Python's loader by
# yielding each of the objects that need to be added to its meta_path.
def finders():
    '''Yield each finder that will be used by the plugin to locate its modules.'''

    documentation = 'This is a ctypes-library to the shared object that is exposed to the IDA SDK.'

    # IDA's native lower-level api
    if sys.platform in {'darwin'}:
        yield internal_object('ida', library(idaapi.idadir("libida{:s}.dylib".format('' if idaapi.BADADDR < 0x100000000 else '64'))), __doc__=documentation)

    elif sys.platform in {'linux', 'linux2'}:
        yield internal_object('ida', library(idaapi.idadir("libida{:s}.so".format('' if idaapi.BADADDR < 0x100000000 else '64'))), __doc__=documentation)

    elif sys.platform in {'win32'}:
        if os.path.exists(idaapi.idadir('ida.wll')):
            yield internal_object('ida', library(idaapi.idadir('ida.wll')))
        else:
            yield internal_object('ida', library(idaapi.idadir("ida{:s}.dll".format('' if idaapi.BADADDR < 0x100000000 else '64'))), __doc__=documentation)

    else:
        logging.warning("{:s} : Unable to load IDA's native api via ctypes. Ignoring...".format(__name__))

    # private (internal) api
    documentation = 'This virtual module contains a number of internal submodules.'
    yield internal_submodule('internal', os.path.join(root, 'base'), include='_*.py', __doc__=documentation)

    # public api
    yield internal_path(os.path.join(root, 'base'), exclude='_*.py')
    yield internal_path(os.path.join(root, 'misc'))

    # tools and application api
    documentation = 'This virtual module contains a number of different files as submodules.'
    for directory in ['tools', 'application']:
        yield internal_submodule(directory, os.path.join(root, directory), __doc__=documentation)
    return

# The following logic is simply for detecting the version of IDA and
# for stashing it directly into the "idaapi" module.

# needed because IDA 6.95 is fucking stupid and sets the result of idaapi.get_kernel_version() to a string
def host_version():
    '''Return the version of the host application as the major, minor, and a floating-point variation.'''

    # if the api doesn't exist, then go back to some crazy version.
    if not hasattr(idaapi, 'get_kernel_version'):
        return 6, 0, 6.0

    import math
    res = str(idaapi.get_kernel_version())      # force it to a str because IDA 7.0 "fixed" it
    major, minor = map(int, res.split('.', 2))
    minor = int("{:<02d}".format(minor))
    if minor > 0:
        count = 1 + math.floor(math.log10(minor))
        return major, minor, float(major) + minor * pow(10, -count)
    return major, minor, float(major)

def patch_version(module):
    '''Patch the version of the host application into a given module.'''
    version = host_version()
    module.__version_major__, module.__version_minor__, module.__version__ = version
    return version

# The following logic is the display hook that we install in order to
# control how all of our output renders in the REPL.
class DisplayHook(object):
    """
    Re-implementation of IDAPython's displayhook that doesn't tamper with
    classes that inherit from base classes
    """
    def __init__(self, output, displayhook):
        self.orig_displayhook = displayhook or sys.displayhook

        # Save our output callable so we can use it to write raw
        # and unprocessed information to it.
        self.output = output

    def format_seq(self, num_printer, storage, item, open, close):
        storage.append(open)
        for idx, el in enumerate(item):
            if idx > 0:
                storage.append(', ')
            self.format_item(num_printer, storage, el)
        storage.append(close)

    def format_basestring(self, string):
        # FIXME: rather than automatically evaluating the string as we're
        #        currently doing, it'd be much cleaner if we just format the
        #        result from a function with some sort of wrapper object. This
        #        way we can check its type, and then choose whether to unwrap it
        #        or not. This can be done with a decorator of some sort that
        #        communicates to this implementation that it will need to
        #        distinguish between printable strings that we can output and
        #        strings that should be processed by the user.
        # XXX: maybe we can even use this wrapper object to allow this class to
        #      handle aligning columns in a table automatically such as when
        #      more than one element in a row is being returned.
        try:
            result = u"{!r}".format(string)
        except UnicodeDecodeError:
            import codecs
            encoded, _ = codecs.escape_encode(string)
            result = u"'{!s}'".format(encoded)
        return result

    def format_ctypes(self, num_printer, storage, item):
        cls, size = item.__class__, ctypes.sizeof(item)
        if isinstance(item, ctypes._SimpleCData):
            contents = "{:#0{:d}x}".format(item.value, 2 + 2 * size) if isinstance(item.value, six.integer_types) else "{!s}".format(item.value)
            storage.append("{:s}({:s})".format(cls.__name__, contents))

        # if it's anything else (or an unknown), then use the default formatter.
        else:
            storage.append("{!r}".format(item))
        return

    def format_item(self, num_printer, storage, item):
        if item is None or isinstance(item, bool):
            storage.append("{!s}".format(item))
        elif isinstance(item, six.string_types):
            storage.append(self.format_basestring(item))
        elif isinstance(item, six.integer_types):
            storage.append(num_printer(item))
        elif isinstance(item, idaapi.tinfo_t):
            storage.append("{!s}".format(item.dstr()))
        elif isinstance(item, (ctypes._SimpleCData, ctypes._Pointer, ctypes._CFuncPtr, ctypes.Array, ctypes.Structure)):
            self.format_ctypes(num_printer, storage, item)
        elif item.__class__ is list:
            self.format_seq(num_printer, storage, item, '[', ']')
        elif item.__class__ is tuple:
            self.format_seq(num_printer, storage, item, '(', ')')
        elif item.__class__ is set:
            self.format_seq(num_printer, storage, item, 'set([', '])')
        elif item.__class__ is dict:
            storage.append('{')
            for idx, pair in enumerate(item.items()):
                if idx > 0:
                    storage.append(', ')
                self.format_item(num_printer, storage, pair[0])
                storage.append(": ")
                self.format_item(num_printer, storage, pair[1])
            storage.append('}')
        else:
            storage.append("{!r}".format(item))

    def _print_hex(self, x):
        return "{:#x}".format(x)

    def displayhook(self, item):
        builtins._ = item
        if item is None or not hasattr(item, '__class__') or item.__class__ is bool:
            self.orig_displayhook(item)
            return
        try:
            storage = []
            if idaapi.__version__ < 7.0:
                import idaapi as ida_idp
            else:
                import ida_idp
            num_printer = self._print_hex
            dn = ida_idp.ph_get_flag() & ida_idp.PR_DEFNUM
            if dn == ida_idp.PRN_OCT:
                num_printer = oct
            elif dn == ida_idp.PRN_DEC:
                num_printer = str
            elif dn == ida_idp.PRN_BIN:
                num_printer = bin
            self.format_item(num_printer, storage, item)
            self.output("%s\n" % "".join(storage))
        except Exception:
            import traceback
            traceback.print_exc()
            self.orig_displayhook(item)

# Locate the user's dotfile and execute it within the specified namespace.
def dotfile(namespace, filename=u'.idapythonrc.py'):
    '''Execute the user's dotfile within the specified namespace.'''
    path = None

    # A closure that just consumes paths until it finds one that
    # it can read and execute it within the specified namespace.
    def read_and_execute(filename, namespace):
        path = (yield)
        while True:
            fp = os.path.join(path, filename)
            try:
                with open(fp, 'r') as infile:
                    content = infile.read()
            except IOError:
                logging.debug("{:s} : Error reading the dotfile at `{!s}`.".format(__name__, fp), exc_info=True)
            else:
                break
            path = (yield)

        exec(content, namespace)

    # Another closure that just tries to determine the user's home
    # directory from the environment variables it receives.
    def find_home(*variables):
        iterable = (var for var in variables if os.getenv(var, default=None) is not None)
        path = next(iterable, None)
        if path is None:
            raise OSError("{:s} : Unable to determine the user's home directory from the environment.".format(__name__))
        return path, os.getenv(path)

    # Create our coroutine, initialize it, and then feed it some paths.
    tribulations = read_and_execute(filename, namespace); next(tribulations)

    try:
        path = os.path.expanduser('~')
        tribulations.send(path)

        var, path = find_home('HOME', 'USERPROFILE')
        tribulations.send(path)

    # If we stopped, then it was read and executed successfully.
    except StopIteration:
        logging.debug("{:s} : Successfully read and executed the dotfile at `{!s}`.".format(__name__, os.path.join(path, filename)))

    # If we received an OSError, then this likely happened while we
    # were trying to find the home directory. Pass it to the user.
    except OSError as E:
        print(E)

    # Any other exception is because of an issue in the user's script,
    # so we'll do our best to log the backtrace for them to debug.
    except Exception:
        logging.warning("{:s} : Unexpected exception raised while trying to execute the dotfile at `{!s}`.".format(__name__, os.path.join(path, filename)), exc_info=True)

    # If we didn't get an exception, then literally we couldn't find
    # any file that we were supposed to execute. Log it and move on.
    else:
        vowels, alpha = tuple('aeiou'), next((filename[index:] for index, item in enumerate(filename.lower()) if item in 'abcdefghijklmnopqrstuvwxyz'), filename)
        logging.warning("{:s} : Unable to locate a{:s} {:s} dotfile in the user's {:s} directory ({!s}).".format(__name__, 'n' if alpha.startswith(vowels) else '', filename, var, path))
    finally:
        tribulations.close()

    # Verify that executing the dotfile did not add the root directory
    # to the system path because the contents of the root directory is
    # explicitly handled by our loaders.
    busted = [(index, item) for index, item in enumerate(sys.path) if os.path.realpath(item) in {os.path.realpath(root)}]
    if busted:
        logging.warning("{:s} : Execution of `{!s}` has resulted in a conflict between the repository path ({!r}) and the system path.".format(__name__, filename, root))
        [ logging.warning("{:s} : The system path at index {:d} ({!r}) resolves to {!r} which conflicts with the plugin and may interfere with imports.".format(__name__, index, item, os.path.realpath(item))) for index, item in busted ]
    return

# The following logic is actually responsible for starting up the whole
# plugin. This is done with by trying with a notification, falling back
# to a timer, and then straight-up executing things if all else fails.
def startup(namespace=None):
    '''Patch in a notification hander and start up everything left in the plugin.'''

    # First check that we've installed the version patch and it's the
    # right type as everything literally revolves around that critical step.
    if not hasattr(idaapi, '__version__'):
        raise SystemError("{:s} : Unable to start up plugin due to the \"{:s}\" attribute not having been assigned.".format(__name__, '.'.join([idaapi.__module__, '__version_'])))

    if not isinstance(idaapi.__version__, float):
        raise SystemError("{:s} : Unable to start up plugin due to the \"{:s}\" attribute not being a valid type ({!r})".format(__name__, '.'.join([idaapi.__module__, '__version_']), idaapi.__version__))

    # Now we need to make sure we have access to our internal module.
    # We can simply trap for ImportError to ensure this works.
    try:
        import internal
        internal.interface

    except ImportError:
        logging.critical("{:s} : An error occured while trying to access the \"{:s}\" module.".format(__name__, 'internal'), exc_info=True)
        raise SystemError("{:s} : Unable to start up plugin without being able to access its \"{:s}\" modules.".format(__name__, 'internal'))

    except AttributeError:
        logging.critical("{:s} : An error occured while trying to access the \"{:s}\" module.".format(__name__, '.'.join(['internal', 'interface'])), exc_info=True)
        raise SystemError("{:s} : Unable to start up plugin due to an error while loading its \"{:s}\" modules.".format(__name__, 'internal'))

    # The next module we need to make sure we have access to is our
    # hooks module which contains all of our startup logic.
    try:
        import hooks

    except ImportError:
        logging.critical("{:s} : An error occured while trying to access the \"{:s}\" module.".format(__name__, 'hooks'), exc_info=True)
        raise SystemError("{:s} : Unable to start up plugin without being able to access its \"{:s}\" module.".format(__name__, 'hooks'))

    # Finally we can construct our priority notification class and inject it into
    # the IDAPython module. This object needs to exist in order for everything to
    # initialize and deinitialize properly.
    idaapi.__notification__ = notification = internal.interface.prioritynotification()

    # Now we can proceed to install our hooks that actually initialize and uninitialize
    # MINSC. We define two closures here (with documentation) because these are actually
    # accessible by the user if they navigate the hook interface that we expose to them.
    def execute_user_dotfile(*args, **kwargs):
        '''This function is responsible for executing the dotfile in the home directory of the user.'''
        return dotfile(namespace)

    # This second closure is meant to be called from a timer. To avoid having to register
    # multiple timers, we pack up our logic that loads the plugin and the user's dotfile
    # within a single function to register.
    def load_plugin_and_execute_user_dotfile(*args, **kwargs):
        '''This function is responsible for loading the plugin and executing the dotfile in the home directory of the user.'''
        result = hooks.ida_is_busy_sucking_cocks(*args, **kwargs)
        dotfile(namespace)
        return result

    # Finally we can register the functions that will actually be responsible for
    # initializing the plugin. If we were given a namespace, then we can register
    # our closure that loads the user's dotfile too.
    try:
        notification.add(idaapi.NW_INITIDA, hooks.make_ida_not_suck_cocks, -1000)
        namespace and notification.add(idaapi.NW_INITIDA, execute_user_dotfile, 0)

    # If installing that hook failed, then check if we're running in batch mode. If
    # we are, then just immediately register things and load the user dotfile.
    except Exception:
        TIMEOUT = 5
        if idaapi.cvar.batch:
            hooks.ida_is_busy_sucking_cocks()
            namespace and dotfile(namespace)

        # Otherwise warn the user about this and register our hook with a timer.
        elif namespace is None:
            logging.warning("{:s} : Unable to add a notification via `{:s}` for {:s}({:d}).".format(__name__, '.'.join(['idaapi', 'notify_when']), '.'.join(['idaapi', 'NW_INITIDA']), idaapi.NW_INITIDA))
            logging.warning("{:s} : Registering {:.1f} second timer with `{:s}` in an attempt to load plugin...".format(__name__, TIMEOUT, '.'.join(['idaapi', 'register_timer'])))
            idaapi.register_timer(TIMEOUT, hooks.ida_is_busy_sucking_cocks)
            six.print_('=' * 86)

        # If we were given a namespace to load into, then we register the closure
        # that we defined into the timer so that the user's dotfile gets executed.
        else:
            logging.warning("{:s} : Unable to add a notification via `{:s}` for {:s}({:d}).".format(__name__, '.'.join(['idaapi', 'notify_when']), '.'.join(['idaapi', 'NW_INITIDA']), idaapi.NW_INITIDA))
            logging.warning("{:s} : Registering {:.1f} second timer with `{:s}` in an attempt to load plugin...".format(__name__, TIMEOUT, '.'.join(['idaapi', 'register_timer'])))
            idaapi.register_timer(TIMEOUT, load_plugin_and_execute_user_dotfile)
            six.print_('=' * 86)
        del(TIMEOUT)

    # If we were able to hook NW_INITIDA, then the NW_TERMIDA hook should also work.
    else:
        try:
            notification.add(idaapi.NW_TERMIDA, hooks.make_ida_suck_cocks, +1000)

        # Installing the termination hook failed, but it's not really too important...
        except Exception:
            logging.warning("{:s} : Unable to add a notification for idaapi.NW_TERMIDA({:d}).".format(__name__, idaapi.NW_TERMIDA))
    return

# Now we can define our plugin_t that literally does nothing if we've already been
# loaded via the idapythonrc.py file and control the entire default namespace.

class MINSC(idaapi.plugin_t):
    wanted_name = 'About Minsc'
    comment = 'Makes IDAPython Not Suck Completely.'
    wanted_hotkey = ''

    flags = idaapi.PLUGIN_FIX
    state = None

    help = 'You should totally check out the `dill` Python module so you can save your game.'

    def get_loader(self):
        '''Return the loader containing all the components needed for loading and initializing the plugin'''
        import imp

        # We explicitly create our own version of the loader from the current
        # file. The functionality we need is actually within our current module,
        # but IDA was responsible for loading it. Most importantly, though, is
        # that the loader is intended to be completely thrown away after usage.
        try:
            filename = __file__ if os.path.exists(__file__) else os.path.join(root, 'plugins', 'minsc.py')
            module = imp.load_source("{:s}__loader__".format(__name__), filename)

        except IOError:
            logging.critical("{:s} : A critical error occurred while trying to read the plugin loader from the file: {:s}".format(__name__, filename), exc_info=True)

        except ImportError:
            logging.critical("{:s} : A critical error occurred while initializing the plugin loader in \"{:s}\"".format(__name__, filename), exc_info=True)

        except Exception:
            logging.critical("{:s} : A critical error occurred while initializing the plugin loader".format(__name__, filename), exc_info=True)

        return module

    def init(self):
        version = getattr(idaapi, '__version__', None)

        # Check our version.. but not really. We're only checking it to see
        # whether the plugin has been loaded yet. If our version if a float,
        # then our module finders have already been loaded and we just need
        # to persist ourselves.
        if isinstance(version, float):
            self.state = self.__class__.state = 'persistent'
            return idaapi.PLUGIN_KEEP

        # If our state is already initialized, then we've done this before.
        elif self.state:
            logging.critical("{:s} : Loading plugin again despite it already being initialized ({:s}).".format(__name__, self.state))

        # Now the version hasn't been assigned yet, then the user didn't
        # install this globally. This means that we don't control the primary
        # namespace. So we'll need to load ourselves still and then afterwards
        # we can uninstall ourselves whenever our plugin is asked to terminate.
        loader = self.get_loader()
        if not loader:
            raise SystemError("{:s} : Unable to get the loader required by the plugin.".format(__name__))

        # Iterate through all of the items in our system path in order to remove
        # any that reference our plugin root. This is because we're using our
        # own loaders to find modules instead of Python's file loaders.
        sys.path[:] = [item for item in sys.path if os.path.realpath(item) not in {os.path.realpath(root)}]

        # Seed the metapath, then patch the version into the idaapi module.
        sys.meta_path.extend(loader.finders())
        _, _, version = loader.patch_version(idaapi)

        # Check if IDAPython (6.95) has replaced the display hook with their
        # own version. We're going to undo exactly what they did, because
        # we're going to replace it with our own anyways.
        ida_idaapi = __import__('ida_idaapi') if version >= 6.95 else idaapi
        if hasattr(ida_idaapi, '_IDAPython_displayhook') and hasattr(ida_idaapi._IDAPython_displayhook, 'orig_displayhook'):
            orig_displayhook = ida_idaapi._IDAPython_displayhook.orig_displayhook
            del(ida_idaapi._IDAPython_displayhook.orig_displayhook)

            sys.displayhook = loader.DisplayHook(sys.stdout.write, orig_displayhook).displayhook

        # If it's the builtin displayhook then we can use it as-is.
        elif getattr(sys.displayhook, '__module__', '') == 'sys':
            sys.displayhook = loader.DisplayHook(sys.stdout.write, sys.displayhook).displayhook

        # Anything else means that some plugin or somebody else did something
        # crazy, and we have no idea how to recover from this.
        else:
            logging.warning("{:s} : Skipping installation of the display hook at \"{:s}\" due to a lack of awareness about the current one ({!r}).".format(__name__, '.'.join(['sys', 'displayhook']), sys.displayhook))

        # Now we'll try and tamper with the user's namespace. We'll search through
        # Python's module list, and if we find it we'll just swap it out for root.
        if '__main__' in sys.modules:
            ns, banner_required = sys.modules['__main__'], {'print_banner', 'IDAPYTHON_VERSION', 'sys'}
            loader.load(ns.__dict__, preserve={'_orig_stdout', '_orig_stderr'} | banner_required)
            hasattr(ns, 'print_banner') and ns.print_banner()
            [ns.__dict__.pop(item, None) for item in banner_required]

        else:
            logging.warning("{:s} : Skipping the reset of the primary namespace as \"{:s}\" was not found in Python's module list.".format(__name__, '__main__'))

        # We don't bother tampering with the user's namespace, since technically
        # we don't have access to it.. However, we'll still try to install the
        # necessary hooks or other features depending what we found available.
        ok = True
        try:
            import internal, hooks

        except (ImportError, Exception):
            logging.critical("{:s} : An error occurred while trying to import the necessary modules \"{:s}\", and \"{:s}\".".format(__name__, 'internal', 'hooks'), exc_info=True)
            ok = False

        try:
            ok and internal.interface

        except AttributeError:
            logging.critical("{:s} : One of the internal modules, \"{:s}\", is critical but was not properly loaded.".format(__name__, '.'.join(['internal', 'interface'])))

        # Check to see if our notification instance was assigned into idaapi. If
        # it wasn't then try to construct one and assign it for usage.
        try:
            if ok and not hasattr(idaapi, '__notification__'):
                idaapi.__notification__ = notification = internal.interface.prioritynotification()

        except Exception:
            logging.warning("{:s} : An error occurred while trying to instantiate the notifications interface. Notifications will be left as disabled.".format(__name__))

        # Check to see if all is well, and if it is then we can proceed to install
        # the necessary hooks to kick everything off.
        if ok:
            logging.info("{:s} : Plugin has been successfully initialized and will now start attaching to the necessary handlers.".format(__name__))
            self.state = self.__class__.state = 'local'
            hooks.make_ida_not_suck_cocks(idaapi.NW_INITIDA)

            # If there's an accessible "__main__" namespace, then dump the dotfile into it.
            '__main__' in sys.modules and dotfile(sys.modules['__main__'].__dict__)

        else:
            logging.warning("{:s} : Due to previous errors the plugin was not properly attached. Modules may still be imported, but a number of features will not be available.".format(__name__))
            self.state = self.__class__.state = 'disabled'

        return idaapi.PLUGIN_KEEP

    def term(self):
        if self.state is None:
            logging.warning("{:s} : Ignoring the host application request to terminate as the plugin has not yet been initialized.".format(__name__))
            return

        # Figure out how we were started so that we can slowly tear things down.
        if self.state in {'disabled', 'persistent'}:
            logging.debug("{:s} : Host application requested termination of {:s} plugin.".format(__name__, self.state))
            return

        # We were run locally, so we're only allowed to interact with the current
        # database. This means that we now will need to shut everything down.
        try:
            import internal, hooks

        except ImportError:
            logging.critical("{:s} : An error occurred while trying to import the necessary modules \"{:s}\", and \"{:s}\" during plugin termination.".format(__name__, 'internal', 'hooks'), exc_info=True)
            return

        # Now we can just remove our hooks and all should be well.
        try:
            logging.debug("{:s} : Detaching from the host application as requested.".format(__name__))
            hooks.make_ida_suck_cocks(idaapi.NW_TERMIDA)

        except Exception:
            logging.critical("{:s} : An error occurred while trying to detach from the host application during plugin termination. Application may become unstable.".format(__name__), exc_info=True)
        return

    def run(self, args):
        import ui

        # Shove some help down the user's throat.
        print("Python>{:<{:d}s} # Use `help({:s})` for usage".format('ui.keyboard.list()', 40, 'ui.keyboard'))
        try:
            ui.keyboard.list()
        except Exception as E:
            print(E)
        print('')

        # Have some more...
        hooks = [name for name in dir(ui.hook) if not name.startswith('__')]
        print('The following hook types are locked and loaded:' if hooks else 'Currently no hooks have been initialized.')
        for name in hooks:
            item = getattr(ui.hook, name)
            fullname = '.'.join(['ui.hook', name])
            print("Python>{:<{:d}s} # Use `help({:s})` for usage and `{:s}.list()` to see availability".format(fullname, 40, fullname, fullname))
            print(item)
            print('')

        # Dead leaves on the dirty ground...when I know you're not around. Shiny
        # tops and soda pops, when I hear you make a sound.

        noise = '''Welcome to the IDA-minsc plugin. My arrow keys are broken.

        This plugin is (mostly) a library that aims to simplify IDAPython. However,
        it utilizes hooks and keyboard shortcuts in a variety of ways in order to
        keep track of the changes that the user may make within their database.

        Essentially the goal of this plugin is to make absolutely _everything_
        that a user may notate in their database serializeable (into a python type)
        and queryable so that things can be exchanged with other Python interpreters.

        Use "." to jump to the command-line and Shift+F2 if you need multi-line.
        Don't forget `dir()` to look around, and `help(thing)` to inquire.
        '''

        # If you can hear a piano fall, you can hear me coming down the hall. If
        # I can just hear your pretty voice, I don't think I need to see at all.

        home = os.path.expanduser('~')
        dotfile = "On startup, the {:s} file will be executed within the primary namespace.".format(os.path.join(home, '.idapythonrc.py'))

        # Every breath that is in your lungs is a tiny little gift to me.

        import database
        path = os.path.join(database.config.path() or '$IDB_DIRECTORY', 'idapythonrc.py')
        rcfile = "Upon {:s} database, the {:s} file will be loaded.".format('opening up the current' if database.config.path() else 'opening up a', os.path.abspath(path))
        ui.message('\n'.join([noise, '\n'.join([dotfile, rcfile])]))
        return

def PLUGIN_ENTRY():
    return MINSC()


repository: binsys/MKBundleManager
file: MKBundleManager/MKBundleManager.py
description: mono mkbundle unpacker and file replace tool.
wanted_name: MK Bundle Manager
comment: MK Bundle Manager.
created_at: 2014-10-06 13:22:16+00:00
pushed_at: 2015-07-07 06:34:45+00:00
url:
https://github.com/binsys/MKBundleManager/blob/3efb9c7b894036b30476055f4eb5b246b8511e7a/MKBundleManager/MKBundleManager.py

content:
﻿#!/usr/bin/env python
# coding=utf-8
# by BinSys <binsys@163.com>

# 支持 mtouch mtouch-64 mtouch.exe mandroid.exe 解包

# Readme

# 1. 将插件文件 MKBundleManager.py 放入 IDA Pro 的 plugins 目录
# 2. 用IDA打开待分析文件，等待分析完毕(左下角状态栏的 AU: idel)
# 3. IDA 菜单栏 点击 View -> Open subviews -> Bundled Assembly Manager
# 4. 在 Bundled Assembly Manager 窗口中可见程序集列表
# 5. 选择要修改的文件用右键菜单内 导出全部文件 或者 导出文件 命令导出到指定位置
# 6. 文件修改完毕后用右键菜单内 替换文件 命令 替换修改后的文件
# 7. 会在位于原程序所在目录内用原文件名+日期时间命名生成替换后的打包文件

# Note:
# 可能会存在问题，请看IDA的输出窗口获取详细出错信息
# .Net 程序集的修改可用 替换文件Radate .NET Reflector + Reflexil 插件
# 当修改后的文件被压缩后大于原始文件的压缩数据大小时无法替换，这时，请用Reflexil删除修改后的程序集的冗余IL指令，减少程序集大小




MKBundleManager_VERSION = "1.2"

# IDA libraries
import idaapi
import idautils
import idc
from idaapi import Form, Choose2, plugin_t

# Python modules
import io
import sys
import os
import shutil
import struct
import binascii
from datetime import datetime, date, time
import urllib2
import httplib
import zlib
import StringIO
import gzip
import traceback

from struct import *



InputFileType_EXE = 11
InputFileType_MachO = 25
InputFileType = -1

Is64Bit = False


string_type_map = {
	0 : "ASCSTR_C",       #              C-string, zero terminated
	1 : "ASCSTR_PASCAL",  #              Pascal-style ASCII string (length byte)
	2 : "ASCSTR_LEN2",    #              Pascal-style, length is 2 bytes
	3 : "ASCSTR_UNICODE", #              Unicode string
	4 : "ASCSTR_LEN4",    #              Delphi string, length is 4 bytes
	5 : "ASCSTR_ULEN2",   #              Pascal-style Unicode, length is 2 bytes
	6 : "ASCSTR_ULEN4",   #              Pascal-style Unicode, length is 4 bytes
}




filetype_t_map = {
	 0 : "f_EXE_old",            # MS DOS EXE File
	 1 : "f_COM_old",            # MS DOS COM File
	 2 : "f_BIN",                # Binary File
	 3 : "f_DRV",                # MS DOS Driver
	 4 : "f_WIN",                # New Executable (NE)
	 5 : "f_HEX",                # Intel Hex Object File
	 6 : "f_MEX",                # MOS Technology Hex Object File
	 7 : "f_LX",                 # Linear Executable (LX)
	 8 : "f_LE",                 # Linear Executable (LE)
	 9 : "f_NLM",                # Netware Loadable Module (NLM)
	10 : "f_COFF",               # Common Object File Format (COFF)
	11 : "f_PE",                 # Portable Executable (PE)
	12 : "f_OMF",                # Object Module Format
	13 : "f_SREC",               # R-records
	14 : "f_ZIP",                # ZIP file (this file is never loaded to IDA database)
	15 : "f_OMFLIB",             # Library of OMF Modules
	16 : "f_AR",                 # ar library
	17 : "f_LOADER",             # file is loaded using LOADER DLL
	18 : "f_ELF",                # Executable and Linkable Format (ELF)
	19 : "f_W32RUN",             # Watcom DOS32 Extender (W32RUN)
	20 : "f_AOUT",               # Linux a.out (AOUT)
	21 : "f_PRC",                # PalmPilot program file
	22 : "f_EXE",                # MS DOS EXE File
	23 : "f_COM",                # MS DOS COM File
	24 : "f_AIXAR",              # AIX ar library
	25 : "f_MACHO",              # Max OS X
}

class BundledAssembly():
	def __init__(self):
		self.Index = 0
		self.FileItemStructOffset = 0
		self.FileNameOffset = 0
		self.FileName = ""
		self.FileDataOffset = 0
		self.FileSize = 0
		self.FileSizeOffset = 0
		self.FileCompressedSizeOffset = 0
		self.FileCompressedSize = 0
		self.IsGZip = ""
		self.FileDataCompressed = ""
		self.IsCompressed = True
		self.IsME = False

		pass

class MKBundleTool():
	def __init__(self):
		self.Is64Bit = False
		self.InputFileType = -1

		print("Input File:{}".format(GetInputFile()))
		print("Input File Path:{}".format(GetInputFilePath()))
		print("Idb File Path:{}".format(GetIdbPath()))
		print("cpu_name:{}".format(idc.GetShortPrm(idc.INF_PROCNAME).lower()))
	
		self.InputFileType = idc.GetShortPrm(idc.INF_FILETYPE)
		#ida.hpp filetype_t f_PE=11 f_MACHO=25
	
		print("InputFileType:{}".format(filetype_t_map.get(self.InputFileType, None)))
	
	
		if self.InputFileType != InputFileType_EXE and self.InputFileType != InputFileType_MachO:
			print "Error,Input file type must is PE or MachO!"
			return
		
	
	
	
		if (idc.GetShortPrm(idc.INF_LFLAGS) & idc.LFLG_64BIT) == idc.LFLG_64BIT:
			self.Is64Bit = True
		else:
			self.Is64Bit = False
		
		print("Is64Bit:{}".format(self.Is64Bit))


	def GetBundledAssemblyList(self,UseScreenEA = False):

		if not UseScreenEA:


			StringEA = self.FindStringEA()
			if StringEA == -1:
				print "Can't find StringEA!"
				return
	
			Func = self.FindUnFunction(StringEA)
	
			if not Func:
				print "Can't find Func!"
				return

			FuncName = idc.GetFunctionName(Func.startEA)
	
			print "Found Data Function:" + FuncName


		
			BundledAssemblyListOffsetsVA = self.FindBundledAssemblyListOffsetsVA(Func.startEA)
			if not BundledAssemblyListOffsetsVA:
				print "Can't find BundledAssemblyListOffsetsVA!"
				return
		else:
			BundledAssemblyListOffsetsVA = ScreenEA()


	
		print("BundledAssemblyListOffsetsVA:0x{:016X}".format(BundledAssemblyListOffsetsVA))

		#StructOffsetList = self.GetStructOffsetList(BundledAssemblyListOffsetsVA)
		BundledAssemblyListOffsetList = self.GetBundledAssemblyListOffsetList(BundledAssemblyListOffsetsVA)
	
		if len(BundledAssemblyListOffsetList) == 0:
			print "Can't find BundledAssemblyListOffsetList!"
			return

		if len(BundledAssemblyListOffsetList) > 2:
			difflen = BundledAssemblyListOffsetList[1] - BundledAssemblyListOffsetList[0]
		else:
			difflen = 0
			return None

		if difflen == 16 or difflen == 32 or difflen == -16 or difflen == -32:
			IsCompressed = True
		else:
			IsCompressed = False

		print "IsCompressed:{}".format(IsCompressed)

		#return None

		#print BundledAssemblyListOffsetList

		BundledAssemblys = []
		BundledAssemblyItemIndex = 0
		for BundledAssemblyListOffset in BundledAssemblyListOffsetList:
			BundledAssemblyItem = self.MakeBundledAssemblyStruct(BundledAssemblyListOffset, IsCompressed)
			BundledAssemblyItem.Index = BundledAssemblyItemIndex
			BundledAssemblys.append(BundledAssemblyItem)
			BundledAssemblyItemIndex += 1

		#print BundledAssemblys

		return BundledAssemblys

	def FindStringEA(self):
		searchstr = str("mkbundle: Error %d decompressing data for %s\n")
		searchstr2 = str("Error %d decompresing data for %s\n")
	
		#Do not use default set up, we'll call setup().
		s = idautils.Strings(default_setup = False)
		# we want C & Unicode strings, and *only* existing strings.
		s.setup(strtypes=Strings.STR_C | Strings.STR_UNICODE, ignore_instructions = True, display_only_existing_strings = True)

		#loop through strings
		for i, v in enumerate(s):                
			if not v:
				#print("Failed to retrieve string at index {}".format(i))
				return -1
			else:
				#print("[{}] ea: {:#x} ; length: {}; type: {}; '{}'".format(i, v.ea,
				#v.length, string_type_map.get(v.type, None), str(v)))
				if str(v) == searchstr or str(v) == searchstr2:
					return v.ea
	
		return -1


	def FindUnFunction(self, StringEA):
		for ref in DataRefsTo(StringEA):
			f = idaapi.get_func(ref)

			if f:
				return f
		return None
	
	def FindBundledAssemblyListOffsetsVA(self, FuncEA):
		for funcitem in FuncItems(FuncEA):
			#print hex(funcitem)
			for dataref in DataRefsFrom(funcitem):
				return dataref
				#print " " + hex(dataref)
		return None
	
	def GetBundledAssemblyListOffsetList(self, DataOffset):



		if self.Is64Bit == True:
			addv = 8
			mf = MakeQword
			vf = Qword
		else:
			mf = MakeDword
			addv = 4
			vf = Dword
	
		AsmListStructListOffset = DataOffset
	
		currentoffset = AsmListStructListOffset



		mf(currentoffset)
		currentvalue = vf(currentoffset)
		currentoffset+=addv


		AsmListStructListOffsetList = []
		AsmListStructListOffsetList.append(currentvalue)
	
		while currentvalue != 0:


			mf(currentoffset)
			currentvalue = vf(currentoffset)
	
			if currentvalue != 0:
				AsmListStructListOffsetList.append(currentvalue)
			currentoffset+=addv
			
		#print len(AsmListStructListOffsetList)
		#for vv in AsmListStructListOffsetList:
			#print hex(vv)
			
		return AsmListStructListOffsetList
	
		
	

	
	def MakeBundledAssemblyStruct(self, FileItemStructOffset, IsCompressed = True):


		if self.Is64Bit == True:
			addv = 8
			mf = MakeQword
			vf = Qword
		else:
			mf = MakeDword
			addv = 4
			vf = Dword
	
		offset = FileItemStructOffset


	
		mf(offset)
		FileNameOffset = vf(offset)
		FileName = idc.GetString(FileNameOffset)
		offset+=addv
	
		mf(offset)
		FileDataOffset = vf(offset)
		offset+=addv
	
		mf(offset)
		FileSize = vf(offset)
		FileSizeOffset = offset
		offset+=addv
	
		if IsCompressed:
			IsME = "N/A"
			mf(offset)
			FileCompressedSize = vf(offset)
			FileCompressedSizeOffset = offset
			offset+=addv
	
			IsGZip = ""
	
			FileDataCompressed = idc.GetManyBytes(FileDataOffset,3)
	
			b1,b2,b3 = struct.unpack('ccc', FileDataCompressed[0:3])
			if b1 == '\x1f' and b2 == '\x8b' and b3 == '\x08':
				IsGZip = "Y"
			else:
				IsGZip = "N"
		else:
			IsME = ""
			IsGZip = "N/A"
			FileDataCompressed = idc.GetManyBytes(FileDataOffset,2)
	
			b1,b2 = struct.unpack('cc', FileDataCompressed[0:2])
			if b1 == '\x4d' and b2 == '\x45':
				IsME = "Y"
			else:
				IsME = "N"
	
		ba = BundledAssembly()
		ba.FileItemStructOffset = FileItemStructOffset
		ba.FileNameOffset = FileNameOffset
		ba.FileName = FileName
		ba.FileDataOffset = FileDataOffset
		ba.FileSize = FileSize
		ba.FileSizeOffset = FileSizeOffset
		if IsCompressed:
			ba.FileCompressedSizeOffset = FileCompressedSizeOffset
			ba.FileCompressedSize = FileCompressedSize
		ba.IsGZip = IsGZip
		ba.IsME = IsME
		if IsCompressed:
			ba.IsCompressed = "Y"
		else:
			ba.IsCompressed = "N"
		#ba.FileDataCompressed = FileDataCompressed

		return ba
		#return {\
		#	"FileItemStructOffset":FileItemStructOffset, \
		#	"FileNameOffset":FileNameOffset,\
		#	"FileName":FileName,\
		#	"FileDataOffset":FileDataOffset,\
		#	"FileSize":FileSize,\
		#	"FileSizeOffset":FileSizeOffset,\
		#	"FileCompressedSizeOffset":FileCompressedSizeOffset,\
		#	"FileCompressedSize":FileCompressedSize,\
		#	"IsGZip":IsGZip,\
		#	"FileDataCompressed":FileDataCompressed\
		#	 }




	#Python语言: Python Cookbook: 比系统自带的更加友好的makedir函数
	#from: http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/82465
	def _mkdir(self, newdir):
		"""works the way a good mkdir should :)
			- already exists, silently complete
			- regular file in the way, raise an exception
			- parent directory(ies) does not exist, make them as well
		"""
		if os.path.isdir(newdir):
			pass
		elif os.path.isfile(newdir):
			raise OSError("a file with the same name as the desired " \
						  "dir, '%s', already exists." % newdir)
		else:
			head, tail = os.path.split(newdir)
			if head and not os.path.isdir(head):
				_mkdir(head)
			#print "_mkdir %s" % repr(newdir)
			if tail:
				os.mkdir(newdir)

	def DecompressZLib(self, Data,Path):
		
		#compressedstream = StringIO.StringIO(Data)
		data2 = zlib.decompress(Data)
		f = open(Path, 'wb')
		f.write(data2)
		f.close()

	def DecompressGzipTo(self, Data,Path):

		compressedstream = StringIO.StringIO(Data)  
		gziper = gzip.GzipFile(fileobj=compressedstream)    
		data2 = gziper.read()   # 读取解压缩后数据

		f = open(Path, 'wb')
		f.write(data2)
		f.close()

	def DecompressFileTo(self, FileItem,OutputDir):
		
		if FileItem.IsME == "Y":
			extname = ".ME"
		else:
			extname = ""

		newpath = '{}\\{}{}'.format(OutputDir, FileItem.FileName, extname)

		if FileItem.IsCompressed == "Y":
			FileDataCompressed = idc.GetManyBytes(FileItem.FileDataOffset,FileItem.FileCompressedSize)
			if FileItem.IsGZip == "Y":
				self.DecompressGzipTo(FileDataCompressed,newpath)
			else:
				self.DecompressZLib(FileDataCompressed,newpath)
		else:
			FileData = idc.GetManyBytes(FileItem.FileDataOffset,FileItem.FileSize)
			f = open(newpath, 'wb')
			f.write(FileData)
			f.close()
			pass

	def CompressGzipToData(self, data):
		buf = StringIO.StringIO() 
		f = gzip.GzipFile(mode='wb',compresslevel=9,fileobj=buf)
		try:
			f.write(data)
		finally:
			f.close()
		compresseddata = buf.getvalue()

		return compresseddata

	def CompressZLibToData(self, data):
		compressobj = zlib.compressobj(zlib.Z_BEST_COMPRESSION, zlib.DEFLATED, zlib.MAX_WBITS, 8 ,zlib.Z_DEFAULT_STRATEGY)
		data2 = compressobj.compress(data)
		data2 += compressobj.flush()
		return data2

	def ReplaceFile(self, FileItem,BundleFilePath,NewFilePath):

		print "Start replace file: {}".format(NewFilePath)

		if not os.path.exists(BundleFilePath):
			print "BundleFilePath error!"
			return


		if not os.path.isfile(NewFilePath):
			print "NewFilePath error!"
			return

		f = open(NewFilePath, 'rb')
		NewFileData = f.read()
		f.close()


		if FileItem.IsCompressed == "Y":
			if FileItem.IsGZip == "Y":
				compresseddata = self.CompressGzipToData(NewFileData)
			else:
				compresseddata = self.CompressZLibToData(NewFileData)

			sizediff = FileItem.FileCompressedSize - len(compresseddata)
			print "FileCompressedSize - compresseddata = 0x{:016X} - 0x{:016X} = 0x{:016X}".format(FileItem.FileCompressedSize,len(compresseddata),sizediff)

			if sizediff < 0:
				print "FileCompressedSize < compresseddata,can't replace!"
				return
		else:
			compresseddata = NewFileData
			sizediff = FileItem.FileSize - len(compresseddata)
			print "FileSize - compresseddata = 0x{:016X} - 0x{:016X} = 0x{:016X}".format(FileItem.FileSize,len(compresseddata),sizediff)

			if sizediff < 0:
				print "FileSize < compresseddata,can't replace!"
				return




		FileSizeOffset = idaapi.get_fileregion_offset(FileItem.FileSizeOffset)

		if FileItem.IsCompressed == "Y":
			FileCompressedSizeOffset = idaapi.get_fileregion_offset(FileItem.FileCompressedSizeOffset)

		FileDataOffset = idaapi.get_fileregion_offset(FileItem.FileDataOffset)
		#ea = idaapi.get_fileregion_ea(offset)

		NewFileDataSize = len(NewFileData)
		if FileItem.IsCompressed == "Y":
			NewCompressedDataSize = len(compresseddata)
			print "FileSizeOffset = 0x{:016X},FileCompressedSizeOffset = 0x{:016X},FileDataOffset = 0x{:016X}".format(FileSizeOffset,FileCompressedSizeOffset,FileDataOffset)
		else:
			print "FileSizeOffset = 0x{:016X},FileDataOffset = 0x{:016X}".format(FileSizeOffset,FileDataOffset)
		

		input_file_dir = os.path.dirname(BundleFilePath)
		input_file_fullname = os.path.basename(BundleFilePath) 
		input_file_name,input_file_extname = os.path.splitext(input_file_fullname)
		#分离扩展名：os.path.splitext(r"c:\python\hello.py") --> ("c:\\python\\hello",
		#".py")


		output_file_fullname = '{}_{:%Y%m%d%H%M%S%f}{}'.format(input_file_name, datetime.now(),input_file_extname)
		output_file_fullpath = os.path.join(input_file_dir,output_file_fullname)


		#shutil.copy(BundleFilePath, output_file_fullpath)

		print "new BundleFilePath path:{}".format(output_file_fullpath)



		fp = open(BundleFilePath,"rb")
		data = fp.read() #读出文件内容
		
		fp.close()


		if self.Is64Bit == True:
			NewFileDataSizeData = struct.pack("q",NewFileDataSize)
			if FileItem.IsCompressed == "Y":
				NewCompressedDataSizeData = struct.pack("q",NewCompressedDataSize)
		else:
			NewFileDataSizeData = struct.pack("l",NewFileDataSize)
			if FileItem.IsCompressed == "Y":
				NewCompressedDataSizeData = struct.pack("l",NewCompressedDataSize)

		
		
		#data 不可更改元素，是固定的,必须转成可对元素操作的list
		data = list(data)


		#try:
		#	data[0] = 'a'
		#	data[1] = 'a'
		#	data[2] = 'a'
		#except Exception,e:
		#	print Exception,":",e


		#range(m, n)这里，range()函数产生的是一个从 m至n-1的整数列表


		#update FileSize
		for i in range(0, len(NewFileDataSizeData)):
			data[FileSizeOffset + i] = NewFileDataSizeData[i]

		##update FileCompressedSize
		if FileItem.IsCompressed == "Y":
			for i in range(0, len(NewCompressedDataSizeData)):
				data[FileCompressedSizeOffset + i] = NewCompressedDataSizeData[i]


		##clear FileData
		if FileItem.IsCompressed == "Y":
			for i in range(0, FileItem.FileCompressedSize):
				data[FileDataOffset + i] = chr(0x0)
				pass
		else:
			for i in range(0, FileItem.FileSize):
				data[FileDataOffset + i] = chr(0x0)
				pass

		##update FileData
		for i in range(0, len(compresseddata)):
			data[FileDataOffset + i] = compresseddata[i]
			pass

		fp2 = open(output_file_fullpath,"wb")
		#把data的list转为str并写入文件
		fp2.write(''.join(data))#重写
		fp2.close()


		print "replace ok!"

class ReplaceFileForm(Form):
	def __init__(self,bundleFile,impFile):
		Form.__init__(self,
r"""BUTTON YES* 替换
BUTTON CANCEL 取消
请选择文件（文件压缩后数据大小必须小于替换前压缩后数据大小）

{FormChangeCb}
<##选择被打包的文件:{bundleFile}>
<##选择修改后的文件:{impFile}>

""".decode('utf-8').encode(sys.getfilesystemencoding()), {  'bundleFile': Form.FileInput(open=True,value=bundleFile),
		'impFile': Form.FileInput(open=True,value=impFile),

		'FormChangeCb': Form.FormChangeCb(self.OnFormChange),
		})

		self.Compile()

	def OnFormChange(self, fid):
		# Form initialization
		if fid == -1:
			self.SetFocusedField(self.bundleFile)
			self.EnableField(self.bundleFile, True)
			self.EnableField(self.impFile, True)
		
		# Form OK pressed
		elif fid == -2:
		    pass
		elif fid == self.bundleFile.id:
			pass
		return 1

class SaveItemsToDirForm(Form):
	def __init__(self,defaultpath):
		Form.__init__(self,
r"""BUTTON YES* 保存
BUTTON CANCEL 取消
请选择输出目录

{FormChangeCb}

<##输出目录:{impFile}>
""".decode('utf-8').encode(sys.getfilesystemencoding()), {
		'impFile': Form.DirInput(value=defaultpath),

		'FormChangeCb': Form.FormChangeCb(self.OnFormChange),
		})

		self.Compile()

	def OnFormChange(self, fid):
		# Form initialization
		if fid == -1:
			self.SetFocusedField(self.impFile)
			self.EnableField(self.impFile, True)
		
		# Form OK pressed
		elif fid == -2:
			pass
		return 1

class BundledAssemblyManagerView(Choose2):
	def __init__(self):
		Choose2.__init__(self,
			"Bundled Assembly Manager",
			[
				["Index",                      6 | Choose2.CHCOL_DEC], 
				["FileItemStructOffset",      18 | Choose2.CHCOL_HEX], 
				["FileNameOffset",            18 | Choose2.CHCOL_HEX], 
				["FileDataOffset",            18 | Choose2.CHCOL_HEX], 
				["FileSize",                  18 | Choose2.CHCOL_HEX], 
				["FileSizeOffset",            18 | Choose2.CHCOL_HEX], 
				["FileCompressedSizeOffset",  18 | Choose2.CHCOL_HEX], 
				["FileCompressedSize",        18 | Choose2.CHCOL_HEX], 
				["IsCompressed",              4  | Choose2.CHCOL_PLAIN], 
				["IsGZip",                    4  | Choose2.CHCOL_PLAIN], 
				["IsME",                      4  | Choose2.CHCOL_PLAIN], 
				["FileName",                  18 | Choose2.CHCOL_PLAIN]
			])
		#self.popup_names = ["Insert", "Delete", "Edit", "Refresh"]

		self.icon = 47

		self.tool = None
		
		self.items = []
		self.items_data = []

		# Command callbacks
		self.cmd_Items_SaveAs = None
		self.cmd_Item_SaveAs = None
		self.cmd_Item_ReplaceBy = None

	def show(self):

		try:
			# Initialize/Refresh the view
			if self.refreshitems():
				if self.Show() < 0: return False
			#self.Refresh()
		except:
			traceback.print_exc()
		

		# Attempt to open the view
		



		if self.cmd_Items_SaveAs == None:
			self.cmd_Items_SaveAs = self.AddCommand("导出全部文件...".decode('utf-8').encode(sys.getfilesystemencoding()), flags = idaapi.CHOOSER_POPUP_MENU | idaapi.CHOOSER_NO_SELECTION, icon=139)
		
		if self.cmd_Item_SaveAs == None:
			self.cmd_Item_SaveAs = self.AddCommand("导出文件...".decode('utf-8').encode(sys.getfilesystemencoding()), flags = idaapi.CHOOSER_POPUP_MENU, icon=139)
		
		if self.cmd_Item_ReplaceBy == None:
			self.cmd_Item_ReplaceBy = self.AddCommand("替换文件...".decode('utf-8').encode(sys.getfilesystemencoding()), flags = idaapi.CHOOSER_POPUP_MENU, icon=139)

		return True

	def refreshitems(self):
		#print "refreshitems"
		self.items_data = []
		self.items = []
		try:

			#-1:cancel,0-no,1-ok
			UseScreenEAInt = idc.AskYN(1,"是否自动获取数据位置?(否则使用ScreenEA)".decode('utf-8').encode(sys.getfilesystemencoding()))

			if UseScreenEAInt == -1:
				UseScreenEAInt = 1

			if UseScreenEAInt == 1:
				UseScreenEA = False
			else:
				UseScreenEA = True

			print "UseScreenEA:{}".format(UseScreenEA)


			self.tool = MKBundleTool()
			asms = self.tool.GetBundledAssemblyList(UseScreenEA)
			if not asms:
				return False
			for BundledAssemblyItem in asms:
				#print BundledAssemblyItem
			
				#print("FileItemStructOffset:{:016X} FileNameOffset:{:016X}
				#FileDataOffset:{:016X} FileSize:{:016X} FileCompressedSize:{:016X}
				#IsGZip:{} FileName:{}"\
				#.format( \
				#BundledAssemblyItem.FileItemStructOffset , \
				#BundledAssemblyItem.FileNameOffset,\
				#BundledAssemblyItem.FileDataOffset,\
				#BundledAssemblyItem.FileSize,\
				#BundledAssemblyItem.FileCompressedSize,\
				#BundledAssemblyItem.IsGZip,\
				#BundledAssemblyItem.FileName))


				if  self.tool.Is64Bit:
					fstr = "0x%016X"
				else:
					fstr = "0x%08X"


				self.items_data.append(BundledAssemblyItem)
				self.items.append(["%d" % BundledAssemblyItem.Index,
									fstr % BundledAssemblyItem.FileItemStructOffset,
									fstr % BundledAssemblyItem.FileNameOffset,
									fstr % BundledAssemblyItem.FileDataOffset,
									fstr % BundledAssemblyItem.FileSize,
									fstr % BundledAssemblyItem.FileSizeOffset,
									fstr % BundledAssemblyItem.FileCompressedSizeOffset,
									fstr % BundledAssemblyItem.FileCompressedSize,
									BundledAssemblyItem.IsCompressed,
									BundledAssemblyItem.IsGZip,
									BundledAssemblyItem.IsME,
									BundledAssemblyItem.FileName])
			return True

		except:
			traceback.print_exc()
			return False

	def OnCommand(self, n, cmd_id):


		if self.tool == None:
			return 1;

		if cmd_id == self.cmd_Items_SaveAs:

			OutputDir = '{}_{:%Y%m%d%H%M%S%f}'.format(GetInputFilePath(), datetime.now())

			f = SaveItemsToDirForm(OutputDir)

			# Execute the form
			ok = f.Execute()
			if ok == 1:
				try:
					imp_file = f.impFile.value
					self.tool._mkdir(imp_file)

					for item in self.items_data:
						self.tool.DecompressFileTo(item,OutputDir)
				except:
					traceback.print_exc()

			# Dispose the form
			f.Free()
		elif cmd_id == self.cmd_Item_SaveAs:

			item = self.items_data[n]

			OutputDir = '{}_{:%Y%m%d%H%M%S%f}'.format(GetInputFilePath(), datetime.now())

			f = SaveItemsToDirForm(OutputDir)

			# Execute the form
			ok = f.Execute()
			if ok == 1:
				try:
					imp_file = f.impFile.value
					self.tool._mkdir(imp_file)
					self.tool.DecompressFileTo(item,OutputDir)
				except:   
					traceback.print_exc()
			# Dispose the form
			f.Free()
		elif cmd_id == self.cmd_Item_ReplaceBy:
			item = self.items_data[n]

			f = ReplaceFileForm(GetInputFilePath(),item.FileName)

			# Execute the form
			ok = f.Execute()
			if ok == 1:
				try:   
					self.tool.ReplaceFile(item,f.bundleFile.value,f.impFile.value)
				except:
					traceback.print_exc()
					
			f.Free()
			pass

		return 1

	def OnClose(self):
		self.cmd_Items_SaveAs = None
		self.cmd_Item_SaveAs = None
		self.cmd_Item_ReplaceBy = None

	def OnSelectLine(self, n):
		idaapi.jumpto(self.items_data[n].FileItemStructOffset)
		pass

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

	def OnGetIcon(self, n):
		# Empty list
		if not len(self.items) > 0:
		    return -1

		#return -1
		return 137


	def OnGetSize(self):
		return len(self.items)
	
	def OnRefresh(self, n):
		#print "OnRefresh"
		self.refreshitems()
		return n

	def OnActivate(self):
		#print "OnActivate"
		self.refreshitems()

class MKBundleManager():
	""" Class that manages GUI forms and MKBundleManager methods of the plugin. """
	def __init__(self): 
		
		self.addmenu_item_ctxs = list()
		self.bundledAssemblyManagerView = None
		#self.bundledAssemblyManagerView = BundledAssemblyManagerView()

	#--------------------------------------------------------------------------
	# Menu Items
	#--------------------------------------------------------------------------
	def add_menu_item_helper(self, menupath, name, hotkey, flags, pyfunc, args):
	
		# add menu item and report on errors
		addmenu_item_ctx = idaapi.add_menu_item(menupath, name, hotkey, flags, pyfunc, args)
		if addmenu_item_ctx is None:
			return 1
		else:
			self.addmenu_item_ctxs.append(addmenu_item_ctx)
			return 0
	
	def add_menu_items(self):
	
		if self.add_menu_item_helper("View/Open subviews/Problems", "Bundled Assembly Manager", "", 1, self.Show_BundledAssemblyManagerView, None): return 1
		return 0
	
	def del_menu_items(self):
		for addmenu_item_ctx in self.addmenu_item_ctxs:
			idaapi.del_menu_item(addmenu_item_ctx)

	#--------------------------------------------------------------------------
	# View Callbacks
	#--------------------------------------------------------------------------
	
	# BundledAssemblyManagerView View
	def Show_BundledAssemblyManagerView(self):

		try:
			if not self.bundledAssemblyManagerView == None:
				self.bundledAssemblyManagerView.Close()
				self.bundledAssemblyManagerView = None

			self.bundledAssemblyManagerView = BundledAssemblyManagerView()
			self.bundledAssemblyManagerView.show()
		except:
			traceback.print_exc()
#--------------------------------------------------------------------------
# Plugin
#--------------------------------------------------------------------------
class MKBundleManager_t(plugin_t):

	flags = idaapi.PLUGIN_UNL
	comment = "MK Bundle Manager."
	help = "MK Bundle Manager.."
	wanted_name = "MK Bundle Manager"
	wanted_hotkey = ""
	
	def init(self):  
		global MKBundleManagerInstance
		
		# Check if already initialized
		if not 'MKBundleManagerInstance' in globals():
		
			MKBundleManagerInstance = MKBundleManager()
			if MKBundleManagerInstance.add_menu_items():
				print "Failed to initialize MK Bundle Manager."
				MKBundleManagerInstance.del_menu_items()
				del MKBundleManagerInstance
				return idaapi.PLUGIN_SKIP
			else:  
				print("Initialized MKBundleManager v%s (c) BinSys <binsys@163.com>" % MKBundleManager_VERSION)
			
		return idaapi.PLUGIN_KEEP
	
	def run(self, arg):
		global MKBundleManagerInstance
		idc.Wait()
		MKBundleManagerInstance.Show_BundledAssemblyManagerView()
	
	def term(self):
		pass

def PLUGIN_ENTRY():
	return MKBundleManager_t()

#--------------------------------------------------------------------------
# Script / Testing
#--------------------------------------------------------------------------
def MKBundleManager_main():
	global MKBundleManagerInstance
	
	if 'MKBundleManagerInstance' in globals():
		MKBundleManagerInstance.del_menu_items()
		del MKBundleManagerInstance
	
	MKBundleManagerInstance = MKBundleManager()
	MKBundleManagerInstance.add_menu_items()
	MKBundleManagerInstance.Show_BundledAssemblyManagerView()

if __name__ == '__main__':
	try:
		#print("Initialized MKBundleManager  %s (c) BinSys <binsys@163.com>" % MKBundleManager_VERSION)
		#MKBundleManager_main() #for Developer only
		pass
	except:
		traceback.print_exc()
	pass


repository: eset/nimfilt
file: nimfilt_ida.py
description: A collection of modules and scripts to help with analyzing Nim binaries
url: https://github.com/eset/nimfilt/blob/5e125d37c00e278432650f434a4af56bed5e54d5/nimfilt_ida.py
created_at: 2024-04-03 12:45:35+00:00
pushed_at: 2024-10-07 18:26:02+00:00
content:
# -*- encoding: utf8 -*-
#
# Copyright (c) 2024 ESET
# Author: Alexandre Côté Cyr <alexandre.cote@eset.com>
# See LICENSE file for redistribution.

import nimfilt
import posixpath

import ida_auto
import ida_bytes
import ida_dirtree
import ida_funcs
import ida_nalt
import ida_name
import ida_segment
import ida_xref
import idaapi
import idautils
import idc

from ida_idp import ph_get_cnbits
from ida_ida import inf_get_app_bitness, inf_is_be
from collections import namedtuple

# From IDA 9.0 on, bin_search replaces bin_search3
try:
    from ida_bytes import bin_search
except ImportError:
    from ida_bytes import bin_search3 as bin_search

AUTO_RUN = False

PLUGIN_NAME = "Nimfilt"
VERSION = "1.1.0"

ST_NIMSTRING = 1
ST_NIMSTRING_PTR = 2
ST_NIMSTRING_PLD = 3

# These global variables can only be calculated after IDA has loaded the program
BITNESS = None
INT_BYTES = None
ENDIANNESS = None
# Flag that marks a string object as a string literal
NIM_STRLIT_FLAG = None

PROGRAM_END = None

def init_globals():
    global PROGRAM_END, BITNESS, INT_BYTES, ENDIANNESS, NIM_STRLIT_FLAG
    PROGRAM_END = ida_segment.get_last_seg().end_ea
    BITNESS = inf_get_app_bitness()
    INT_BYTES = BITNESS // ph_get_cnbits()
    ENDIANNESS = "big" if inf_is_be() else "little"
    # lib/nimbase.h: 1 << (sizeof(int) * CHAR_BITS - 2)
    # Nim's int type is 64bit on 64 bit architectures
    NIM_STRLIT_FLAG = 1 << (BITNESS - 2)

def get_uint(ea):
    bts = idaapi.get_bytes(ea, INT_BYTES, 0)
    return int.from_bytes(bts, ENDIANNESS, signed=False)

# bin_search is significantly faster than ida_search.find_text
# Unlike find_text, it also matches data that isn't typed as a string
def find_text(string: str):
    pattern = ida_bytes.compiled_binpat_vec_t()
    ida_bytes.parse_binpat_str(pattern, 0, '"{}"'.format(string), 0, ida_nalt.STRENC_DEFAULT)  # Strings must be inside double quotes
    return bin_search(0, ida_segment.get_last_seg().end_ea, pattern, idaapi.BIN_SEARCH_FORWARD)

class Nimfilt_plugin(idaapi.plugin_t):
    comment = ""
    flags = idaapi.PLUGIN_MOD | idaapi.PLUGIN_FIX
    help = "Helps with reversing Nim compiled executables"
    wanted_hotkey = ""
    wanted_name = PLUGIN_NAME

    def init(self):
        idaapi.notify_when(idaapi.NW_OPENIDB, self._idb_loaded_handler)
        return idaapi.PLUGIN_KEEP

    # Check if the database is Nim and run if AUTO_RUN is on
    def _idb_loaded_handler(self, _, is_old_database):
        if is_nim_idb():
            print("IDB identified as Nim.")
            if not is_old_database and AUTO_RUN:
                print("Running Nimfilt.")
                self.run()
        else:
            print("IDB could not be confirmed as Nim. You can still run the plugin manually")

    def run(self, arg):
        main()

    def term(self):
        pass

def iterate_segments():
    seg = ida_segment.get_first_seg()
    while seg is not None:
        yield seg
        seg = ida_segment.get_next_seg(seg.start_ea)

def make_nim_strings():
    for seg in iterate_segments():
        if ida_segment.get_segm_name(seg) in [".rdata", ".rodata"] or (seg.type == ida_segment.SEG_DATA and seg.perm == ida_segment.SEGPERM_READ):
            ea = seg.start_ea
            while ea < seg.end_ea:
                # Skip if the address is already typed as a struct
                if not ida_bytes.is_struct(ida_bytes.get_flags(ea)) and (is_str := is_nim_str(ea)):
                    if is_str[0] == ST_NIMSTRING:
                        ea += apply_Nim_string_struct(*is_str[1:])
                    elif is_str[0] == ST_NIMSTRING_PTR:
                        ea += apply_Nim_string_ptr_struct(*is_str[1:])
                    elif is_str[0] == ST_NIMSTRING_PLD:
                        ea += apply_Nim_string_payload_struct(*is_str[1:])
                    else:
                        raise Exception()
                else:
                    ea += 1

# String is NUL terminated and contains only valid ascii non-NUL characters
# TODO: check encodings
def _is_valid_C_str(s: bytes):
    return s[-1] == 0x00 and all([x in range(0x01, 0x80) for x in s[:-1]])

def is_nim_str_payload(ea, ln):
    reserved = get_uint(ea)
    if reserved ^ NIM_STRLIT_FLAG in [0, ln] and ea + ln <= PROGRAM_END:
        return _is_valid_C_str(ida_bytes.get_bytes(ea + INT_BYTES, ln + 1))
    return False

# lib/system/strs_v2.nim -> NimStringV2
def is_nim_str(ea):
    ln = get_uint(ea) & ~NIM_STRLIT_FLAG
    # Contiguous block string
    if ln > 0:
        if is_nim_str_payload(ea + INT_BYTES, ln):
            return (ST_NIMSTRING, ea, ln)
        elif (addr := ida_xref.get_first_dref_from(ea + INT_BYTES)) != idaapi.BADADDR and is_nim_str_payload(addr, ln):
            return (ST_NIMSTRING_PTR, ea, addr, ln)
        elif is_nim_str_payload(ea, ln):
            return (ST_NIMSTRING_PLD, ea, ln)
    return False

"""
Depending on the Nim version, strings can be represented one of a few ways:
StringV2, lib/system/strs_v2.nim
    NimStrPayload {
        int cap;
        char data[len+1];
    }
    NimString {
        int len;
        NimStrPayload* p;
    }

Strings, lib/system.nim
    TGenericSeq { //Generic sequence type
        int len;
        int reserved;
    }
    NimStringDesc {
        TGenericSeq Sup;
        char data[len+1];
    }

For string literals, the value of cap/reserved is ORed with NIM_STRLIT_FLAG
reserved appears to be the same as cap
"""
StructMember = namedtuple("StructMember", "name, flag, member_type, size")
def create_Nim_string_structs():
    INT_TYPES = {1: ida_bytes.FF_BYTE, 2: ida_bytes.FF_WORD, 4: ida_bytes.FF_DWORD, 8: ida_bytes.FF_QWORD, 16: ida_bytes.FF_OWORD}
    str_opinfo = ida_nalt.opinfo_t()
    str_opinfo.strtype = ida_nalt.STRTYPE_TERMCHR
    if (nsc_struct_id := idc.get_struc_id("NimStrPayload")) == idaapi.BADADDR:
        NimStrPayload = [
            StructMember("reserved", INT_TYPES[INT_BYTES] | ida_bytes.FF_DATA, -1, INT_BYTES),
            StructMember("str", ida_bytes.FF_STRLIT, str_opinfo.tid, 0)
        ]
        nsc_struct_id = create_IDA_struct("NimStrPayload", NimStrPayload)
    # For structs or pointers to structs, the mt argument must be a opinfo_t struct with the tid field set to the structure's id
    structs = {}
    if idc.get_struc_id("NimString") == idaapi.BADADDR:
        structs["NimString"] = [
            StructMember("length", INT_TYPES[INT_BYTES] | ida_bytes.FF_DATA, -1, INT_BYTES),
            StructMember("content", ida_bytes.FF_STRUCT | ida_bytes.FF_DATA, nsc_struct_id, INT_BYTES)  # Flags for structs
        ]
    if idc.get_struc_id("NimStringPtr") == idaapi.BADADDR:
        structs["NimStringPtr"] = [
            StructMember("length", INT_TYPES[INT_BYTES] | ida_bytes.FF_DATA, -1, 4),
            StructMember("content", INT_TYPES[INT_BYTES] | ida_bytes.FF_0OFF | ida_bytes.FF_1OFF | ida_bytes.FF_DATA, nsc_struct_id, 4)  # Flags for 32 bit pointers
        ]
    for name, members in structs.items():
        create_IDA_struct(name, members)

def create_IDA_struct(name: str, members: list):
    struct_id = idc.add_struc(-1, name, False)
    for field in members:
        field = field._asdict()
        idc.add_struc_member(struct_id, field["name"], -1, field["flag"], field["member_type"], field["size"])
    return struct_id

def apply_Nim_string_payload_struct(start_addr, length):
    struct_id = idc.get_struc_id("NimStrPayload")
    size = idc.get_struc_size(struct_id) + length
    ida_bytes.create_struct(start_addr, size, struct_id, True)
    content = ida_bytes.get_bytes(start_addr + INT_BYTES, length)
    ida_name.set_name(start_addr, str_to_name(content), ida_name.SN_AUTO | ida_name.SN_IDBENC | ida_name.SN_PUBLIC | ida_name.SN_FORCE)
    return size

def apply_Nim_string_struct(start_addr, length):
    struct_id = idc.get_struc_id("NimString")
    size = idc.get_struc_size(struct_id) + length
    ida_bytes.create_struct(start_addr, size, struct_id, True)
    content = ida_bytes.get_bytes(start_addr + 2 * INT_BYTES, length)
    ida_name.set_name(start_addr, str_to_name(content), ida_name.SN_AUTO | ida_name.SN_IDBENC | ida_name.SN_PUBLIC | ida_name.SN_FORCE)
    return size

def apply_Nim_string_ptr_struct(start_addr, content_addr, length):
    ptr_struct_id = idc.get_struc_id("NimStringPtr")
    content_struct_id = idc.get_struc_id("NimStrPayload")
    ida_bytes.create_struct(start_addr, idc.get_struc_size(ptr_struct_id), ptr_struct_id, True)
    size = idc.get_struc_size(content_struct_id) + length
    ida_bytes.create_struct(content_addr, size, content_struct_id, True)
    content = ida_bytes.get_bytes(content_addr + INT_BYTES, length)
    name = str_to_name(content)
    ida_name.set_name(content_addr, name, ida_name.SN_AUTO | ida_name.SN_IDBENC | ida_name.SN_PUBLIC | ida_name.SN_FORCE)
    name = ida_name.get_name(content_addr)  # Get final name in case IDA auto-added a suffix
    ida_name.set_name(start_addr, "ptr_{:s}".format(name), ida_name.SN_AUTO | ida_name.SN_IDBENC | ida_name.SN_PUBLIC | ida_name.SN_FORCE)
    return idc.get_struc_size(ptr_struct_id)

# Returns a name like IDA's default names for string using s as a prefix instead of a
# TODO: handle strings that contain only special characters
def str_to_name(string):
    try:
        default_encoding = ida_nalt.get_encoding_name(ida_nalt.get_default_encoding_idx(1))
        string = string.decode(default_encoding)
    except AttributeError:
        pass
    cleaned = ida_name.validate_name(string, ida_name.NT_LOCAL, ida_name.SN_IDBENC | ida_name.SN_NOCHECK).title().replace("_", "")
    return "s{:s}".format(cleaned[:0xE])

# Parse all functions in the current IDB for ones that have nim mangled names
def parse_nim_functions():
    for i in range(ida_funcs.get_func_qty()):
        func = ida_funcs.getn_func(i)
        name = ida_funcs.get_func_name(func.start_ea)
        try:
            niname = nimfilt.NimNameCreator(name)
            yield func.start_ea, niname
        except ValueError:
            pass

# Rename function. Use mangler generated suffix if there is a duplicate
def rename(ea, nname: nimfilt.NimName):
    name = nname.get_ida_name()
    if ida_name.get_name_ea(0, name) != idaapi.BADADDR:
        name = nname.get_ida_name(suffix=nimfilt.SUF_NIM)
        if ida_name.get_name_ea(0, name) != idaapi.BADADDR:
            name = nname.get_ida_name(suffix=nimfilt.SUF_NIM | nimfilt.SUF_IDA)
    ida_name.set_name(ea, name, ida_name.SN_FORCE)
    return name

# Recursively merge directories that only have a single child that's also a directory
def merge_dir(dirtree: ida_dirtree.dirtree_t, path=""):
    iterator = ida_dirtree.dirtree_iterator_t()
    ok = dirtree.findfirst(iterator, "{}/*".format(path))
    has_fn = False

    # Moving and deleting directories messes with the iterator so we list all the children first
    children = []
    while ok:
        child_path = dirtree.get_abspath(iterator.cursor)
        children.append(child_path)
        ok = dirtree.findnext(iterator)

    new_children = []
    for child_path in children:
        if child_path.endswith("/.."):  # Avoid infinite recursion
            continue
        if dirtree.isdir(child_path):
            child_path = merge_dir(dirtree, child_path)
        else:
            has_fn = True
        new_children.append(child_path)
    if len(children) == 1 and not has_fn:
        # Ida uses POSIX-like path so we use the \\ as an internal separator in directory names
        new_path = "{}\\{}".format(path, posixpath.basename(child_path))
        dirtree.rename(child_path, new_path)
        dirtree.rmdir(path)
        return new_path
    return path

# Returns True if the function at ea has a name that matches known Nim functions
def _nim_func_filter(ea):
    name = ida_funcs.get_name(ea)
    return any(filter(lambda known_name: known_name in name, nimfilt.NIM_FUNC_NAMES))

# Use simple heuristics to see if the currently open program is Nim
def is_nim_idb():
    # Based on our tests, these strings are present in all Nim binaries even if they are stripped. However, it would be trivial for a threat actor to remove them
    if find_text("fatal.nim") != idaapi.BADADDR and find_text("sysFatal") != idaapi.BADADDR:
        return True
    # Function names
    if any(filter(_nim_func_filter, idautils.Functions())):
        return True
    # Other strings
    # Must contain at least 2 of those
    ERR_MSG = ["@value out of range", "@division by zero", "@over- or underflow", "@index out of bounds"]
    matches = 0
    for s in ERR_MSG:
        if find_text(s) != idaapi.BADDADDR:
            matches += 1
            if matches >= 2:
                return True
    return False

# TODO: create separate root level directories for Stdlib, project and nimble packages
# Rename functions and move them to subdirectories based on the package path/name
def main():
    ida_auto.auto_wait()
    init_globals()
    func_dirtree = ida_dirtree.get_std_dirtree(ida_dirtree.DIRTREE_FUNCS)
    for ea, nname in parse_nim_functions():
        name = rename(ea, nname)
        func_dirtree.mkdir(nname.ida_dirname)
        func_dirtree.rename(name, "{}/{}".format(nname.ida_dirname, name))
    merge_dir(func_dirtree)
    create_Nim_string_structs()
    make_nim_strings()

if __name__ == "__main__":
    main()

def PLUGIN_ENTRY():
    return Nimfilt_plugin()


repository: dead-null/idascope
file: plugins/idascope_plugin.py
description: IDAScope is a Neovim plugin that integrates IDA Pro with Telescope
wanted_name: IDAScope
comment: IDA Pro Python Plugin for Telescope Integration
created_at: 2024-12-16 11:23:15+00:00
pushed_at: 2025-02-05 18:13:51+00:00
url:
https://github.com/dead-null/idascope/blob/d305eea0e3333bbd54d7a0d05a3411ee1df000d7/plugins/idascope_plugin.py

content:
from xmlrpc.server import SimpleXMLRPCServer
from enum import Enum
from pathlib import Path
from typing import Callable, Dict, Any
import threading
import idaapi
import idautils


class FunctionNames(Enum):
    LIST_FUNCTIONS = "list_functions"
    DECOMPILE_FUNCTION = "decompile_function"
    DISASSEMBLE_FUNCTION = "disassemble_function"
    GET_PATH = "get_path"


class BaseFunctionRequest:
    def __init__(self, event: threading.Event):
        self.event = event
        self.result: Any = None

    def __call__(self):
        try:
            self.execute()
        except Exception as e:
            self.result = f"Error: {e}"
        finally:
            self.event.set()

    def execute(self):
        """Subclasses must implement this method."""
        raise NotImplementedError("Execute method must be implemented in subclasses.")

class ListFunctionsRequest(BaseFunctionRequest):
    def execute(self):
        """
        Lists functions defined in the current IDA database, excluding imported functions.
        """
        function_names = self.list_all_functions()
        self.result = "\n".join(function_names)

    def list_all_functions(self) -> list:
        """
        Retrieve all functions from the current binary, based on its file type.
        """
        file_type = idaapi.get_file_type_name()

        if "ELF" in file_type:
            functions = self._get_functions_elf()
        elif "PE" in file_type:
            functions = self._get_functions_pe()
        else:
            idaapi.msg(f"[IDAScope] Unsupported file type: {file_type}\n")
            functions = set()

        return sorted(functions)

    def _get_functions_elf(self) -> set:
        """
        Retrieve all functions from an ELF binary, including resolving external functions.
        """
        functions = set()

        for func_ea in idautils.Functions():
            segment_name = idaapi.get_segm_name(idaapi.getseg(func_ea))

            if segment_name == 'extern':
                func_name = self._resolve_extern(func_ea)
            else:
                func_name = idaapi.get_func_name(func_ea)

            if func_name:
                functions.add(func_name)

        # Handle special ELF-specific sections
        extern_functions = {func for func in functions if "extern" in func}
        plt_functions = {func for func in functions if ".plt" in func}

        # Include or exclude based on XOR logic
        return functions.union(extern_functions ^ plt_functions)

    def _get_functions_pe(self) -> set:
        """
        Retrieve all functions from a PE binary, including imports and exports.
        """
        functions = {idaapi.get_func_name(func_ea) for func_ea in idautils.Functions() if idaapi.get_func_name(func_ea)}
        functions.update(self._get_imported_functions())
        functions.update(self._get_exported_functions())
        return functions

    def _get_imported_functions(self) -> list:
        """
        Enumerate imported functions from the Import Address Table (IAT) in PE binaries.
        """
        imported_functions = set()

        def import_enum_callback(ea, name, ordinal):
            if name:
                imported_functions.add(name)
            return True

        for i in range(idaapi.get_import_module_qty()):
            idaapi.enum_import_names(i, import_enum_callback)

        return sorted(imported_functions)

    def _get_exported_functions(self) -> list:
        """
        Enumerate exported functions from the Export Address Table (EAT) in PE binaries.
        """
        exports = []
        for entry in idautils.Entries():
            try:
                ordinal, ea, name = entry
                if name:
                    exports.append(name)
            except ValueError:
                idaapi.msg(f"[IDAScope] Unexpected entry format: {entry}\n")
        return sorted(set(exports))

    def _resolve_extern(self, ea) -> str:
        """
        Resolve external function references for both ELF and PE binaries.
        """
        func_name = idaapi.get_func_name(ea)
        if not func_name:
            return None

        extern_ea = idaapi.get_name_ea(idaapi.BADADDR, func_name)
        if extern_ea == idaapi.BADADDR:
            return None

        for xref_extern_ea in idautils.XrefsTo(extern_ea, idaapi.XREF_ALL):
            if not xref_extern_ea.iscode:
                continue

            xref_type = idaapi.xrefchar(xref_extern_ea.type).lower()
            if xref_type in {'j', 'i'}:  # 'j' for ELF, 'i' for PE
                func = idaapi.get_func(xref_extern_ea.frm)
                if func:
                    return idaapi.get_name(func.start_ea)

        return None

class GetPathRequest(BaseFunctionRequest):
    def __init__(self, func_name: str, event: threading.Event):
        super().__init__(event)
        self.func_name = func_name

    def execute(self):
        parent_path = Path(idaapi.get_input_file_path()).resolve().parent
        binary_name = Path(idaapi.get_input_file_path()).name
        file_path_dir = parent_path / f"{binary_name}_functions"
        file_path_dir.mkdir(exist_ok=True)
        self.result = str(file_path_dir / self.func_name)

class DecompileFunctionRequest(BaseFunctionRequest):
    def __init__(self, func_name: str, filename: str, event: threading.Event):
        super().__init__(event)
        self.func_name = func_name
        self.filename = filename

    def execute(self):
        """
        Attempts to decompile a function. If decompilation fails, fallback to disassembly.
        """
        self.filetype = idaapi.get_file_type_name()
        try:
            ea = idaapi.get_name_ea(idaapi.BADADDR, self.func_name)

            # Check for cached cfunc
            if idaapi.has_cached_cfunc(ea):
                idaapi.clear_cached_cfuncs()

            cfunc_t = idaapi.decompile(ea)
            if not cfunc_t:
                # Decompilation failed. Falling back to disassembly
                disassembler = DisassembleFunctionRequest(self.func_name, self.filename, self.event)
                disassembler.execute()
                self.result += "\n" + disassembler.result
            else:
                self.result = str(cfunc_t)
        except Exception as e:
            self.result = f"IMPORT {self.func_name}" if "PE" in self.filetype else "[ERROR] Disassembly failed."


class DisassembleFunctionRequest(BaseFunctionRequest):
    def __init__(self, func_name: str, filename: str, event: threading.Event):
        super().__init__(event)
        self.func_name = func_name
        self.filename = filename

    def execute(self):
        """
        Disassembles a function, including:
        - Prototype.
        - Repeatable comments.
        - Stack frames.
        - Detailed disassembly lines with segment names and addresses.
        """
        self.filetype = idaapi.get_file_type_name()
        try:
            ea = idaapi.get_name_ea(idaapi.BADADDR, self.func_name)
            func_t = idaapi.get_func(ea)
            if not func_t:
                self.result = f"IMPORT {self.func_name}" if "PE" in self.filetype else "[ERROR] Disassembly failed."
                return

            result = []

            # Add prototype
            prototype = self._get_proto(func_t, ea, self.func_name)
            if prototype:
                result.append(prototype)

            # Add repeatable comments
            func_header = idaapi.get_func_cmt(func_t, True)
            if func_header:
                result.append(f"// {func_header}")

            # Add stack frame details
            frame_details = self._get_stack_frame(func_t)
            if frame_details:
                result.extend(frame_details)

            # Add disassembly lines
            disassembly_lines = self._get_disassembly_lines(func_t)
            result.extend(disassembly_lines)

            self.result = "\n".join(result)
        except Exception as e:
            self.result = f"Error during disassembly: {e}"


    def _get_proto(self, func, ea, func_name) -> str:
        prototype = ''
        if func:
            tif = idaapi.tinfo_t()
            if idaapi.get_tinfo(tif, func.start_ea):
                func_data = idaapi.func_type_data_t()
                if tif.get_func_details(func_data):
                    prototype = idaapi.print_tinfo('', 0, 1, idaapi.PRTYPE_MULTI | idaapi.PRTYPE_PRAGMA, tif, func_name, f"{ea:#x}")
        return prototype



    def _get_stack_frame(self, func) -> list:
        """
        Extracts stack frame details for the given function.
        """
        stack = []
        if hasattr(idaapi, "get_frame"): #could do backward compatibility
            frame = idaapi.get_frame(func)
            if not frame:
                return []

            for member in idautils.StructMembers(frame.id):
                stack.append(f"// {member[1]} = 0x{member[0]:x}")
        else:
            frame_tif = idaapi.tinfo_t()
            frame_udt = idaapi.udt_type_data_t()
            if not idaapi.get_func_frame(frame_tif, func) or not frame_tif.get_udt_details(frame_udt):
                return []
            for udm in frame_udt:
                sval = idaapi.calc_frame_offset(func, udm.offset//udm.type.get_size(), None, None)-udm.offset//4
                stack.append(f"// {udm.name}\t =  {-1 * sval:#x} {udm.type.dstr()} {udm.cmt}")
        return stack

    def _get_disassembly_lines(self, func) -> list:
        """
        Generates disassembly lines with addresses and segment names.
        """
        lines = []
        for head in idautils.Heads(func.start_ea, func.end_ea):
            disasm_line = idaapi.generate_disasm_line(head, idaapi.GENDSM_FORCE_CODE)
            if disasm_line:
                segment_name = idaapi.get_segm_name(idaapi.getseg(head))
                lines.append(f"{segment_name}: {head:08X} {idaapi.tag_remove(disasm_line)}")
        return lines


class FunctionRequestFactory:
    @staticmethod
    def create_request(request_type: FunctionNames, **kwargs) -> BaseFunctionRequest:
        if request_type == FunctionNames.LIST_FUNCTIONS:
            return ListFunctionsRequest(kwargs['event'])
        elif request_type == FunctionNames.DECOMPILE_FUNCTION:
            return DecompileFunctionRequest(kwargs['func_name'], kwargs['filename'], kwargs['event'])
        elif request_type == FunctionNames.DISASSEMBLE_FUNCTION:
            return DisassembleFunctionRequest(kwargs['func_name'], kwargs['filename'], kwargs['event'])
        elif request_type == FunctionNames.GET_PATH:
            return GetPathRequest(kwargs['func_name'], kwargs['event'])
        else:
            raise ValueError(f"Unknown request type: {request_type}")


def handle_request(request_type: FunctionNames, **kwargs) -> str:
    """
    Centralized handler for scheduling and executing function requests.
    """
    event = threading.Event()
    request = FunctionRequestFactory.create_request(request_type, event=event, **kwargs)
    idaapi.execute_ui_requests([request])
    event.wait()
    return request.result


class XMLRPCServerManager:
    _instance = None  # Class-level variable to hold the singleton instance

    def __new__(cls, *args, **kwargs):
        if not cls._instance:
            cls._instance = super(XMLRPCServerManager, cls).__new__(cls, *args, **kwargs)
        return cls._instance

    def __init__(self, host: str = "localhost", port: int = 65432):
        # Prevent re-initialization of the singleton instance
        if hasattr(self, "_initialized") and self._initialized:
            return

        self.host = host
        self.port = port
        self.server = None
        self.thread = None
        self.running = False
        self.registry: Dict[FunctionNames, Callable] = {}
        self._setup_registry()
        self._initialized = True

    def _setup_registry(self):
        """Registers available functions."""
        self.registry = {
            FunctionNames.LIST_FUNCTIONS: lambda: handle_request(FunctionNames.LIST_FUNCTIONS),
            FunctionNames.DECOMPILE_FUNCTION: lambda func_name, filename: handle_request(
                FunctionNames.DECOMPILE_FUNCTION, func_name=func_name, filename=filename
            ),
            FunctionNames.DISASSEMBLE_FUNCTION: lambda func_name, filename: handle_request(
                FunctionNames.DISASSEMBLE_FUNCTION, func_name=func_name, filename=filename
            ),
            FunctionNames.GET_PATH: lambda func_name: handle_request(FunctionNames.GET_PATH, func_name=func_name),
        }

    def start_server(self):
        """
        Starts the XML-RPC server in a separate thread. Only starts if the server is not already running.
        """
        try:
            # Create and start the server
            self.server = SimpleXMLRPCServer((self.host, self.port), allow_none=True)
            for name, func in self.registry.items():
                self.server.register_function(func, name.value)
            self.server.register_function(self.stop_server_method, "stop_server")

            self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
            self.thread.start()

            self.running = True
            idaapi.msg(f"[IDAScope] XML-RPC server started on {self.host}:{self.port}\n")
        except Exception as e:
            idaapi.msg(f"[IDAScope] Failed to start server: {e}\n")
            self.running = False

    def stop_server(self):
        """
        Stops the XML-RPC server and ensures all resources are cleaned up.
        """
        if not self.running and self.server is None:
            return

        idaapi.msg("[IDAScope] Stopping server...\n")
        try:
            if self.server:
                self.server.shutdown()
                self.server.server_close()
                self.server = None
        except Exception as e:
            idaapi.msg(f"[IDAScope] Error during server shutdown: {e}\n")

        if self.thread:
            try:
                self.thread.join(timeout=1)
                if self.thread.is_alive():
                    idaapi.msg("[IDAScope] Server thread did not terminate cleanly. Forcing shutdown.\n")
                    self.thread = None
            except Exception as e:
                idaapi.msg(f"[IDAScope] Error joining server thread: {e}\n")
            finally:
                self.thread = None

        self.running = False
        idaapi.msg("[IDAScope] Server stopped.\n")

    def stop_server_method(self):
        """
        XML-RPC accessible method to stop the server.
        """
        idaapi.msg("[IDAScope] XML-RPC shutdown signal received.\n")
        self.stop_server()

    def toggle_server(self):
        """
        Toggles the server state between running and stopped.
        """
        if self.running:
            self.stop_server()
        else:
            self.start_server()

class IDAScopePlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_KEEP
    comment = "IDA Pro Python Plugin for Telescope Integration"
    help = "This plugin integrates IDA Pro with Neovim's Telescope."
    wanted_name = "IDAScope"
    wanted_hotkey = "Meta-Ctrl-t"

    def init(self):
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        self.manager = XMLRPCServerManager()
        self.manager.toggle_server()

    def term(self):
        self.manager.toggle_server() if XMLRPCServerManager().running else None

def PLUGIN_ENTRY():
    return IDAScopePlugin()

if __name__ == "__main__":
    manager = XMLRPCServerManager()
    manager.toggle_server()


repository: rigmar/ida_type_storage
file: example/print_ctree.py
wanted_name: cobject_helper
comment: This is a comment
created_at: 2015-12-07 18:24:40+00:00
pushed_at: 2024-10-27 10:28:46+00:00
url:
https://github.com/rigmar/ida_type_storage/blob/981b29475fb4bbbf5fb6d1f54d2ff9a2f192199c/example/print_ctree.py

content:
import idautils
import idaapi
from ida_idaapi import *
from ida_hexrays import *
from ida_kernwin import *
import idc
import traceback
import networkx as nx
import matplotlib.image as mpimg
from networkx.drawing.nx_agraph import *

fDebug = False
if fDebug:
    import pydevd

try:
    import matplotlib.pyplot as plt
except:
    raise

print_ctree_actname = "test:print_ctree"

NETNODE_NAME = '$ hexrays-print-ctree'

def hierarchy_pos(G, root, width=1., vert_gap = 0.2, vert_loc = 0, xcenter = 0.5,
                  pos = None, parent = None):
    '''If there is a cycle that is reachable from root, then this will see infinite recursion.
       G: the graph
       root: the root node of current branch
       width: horizontal space allocated for this branch - avoids overlap with other branches
       vert_gap: gap between levels of hierarchy
       vert_loc: vertical location of root
       xcenter: horizontal location of root
       pos: a dict saying where all nodes go if they have been assigned
       parent: parent of this branch.'''
    if pos == None:
        pos = {root:(xcenter,vert_loc)}
    else:
        pos[root] = (xcenter, vert_loc)
    neighbors = G.neighbors(root)
    if parent != None:
        neighbors.remove(parent)
    if len(neighbors)!=0:
        dx = width/len(neighbors)
        nextx = xcenter - width/2 - dx/2
        for neighbor in neighbors:
            nextx += dx
            pos = hierarchy_pos(G,neighbor, width = dx, vert_gap = vert_gap,
                                vert_loc = vert_loc-vert_gap, xcenter=nextx, pos=pos,
                                parent = root)
    return pos

class print_ctree_action_handler_t(idaapi.action_handler_t):
    def __init__(self,obj):
        idaapi.action_handler_t.__init__(self)
        self.hrCbObj = obj

    def activate(self, ctx):
        #print "Activate"
        if fDebug:
            pydevd.settrace('localhost', port=2255, stdoutToServer=True, stderrToServer=True, suspend = False)
        vdui = idaapi.get_tform_vdui(ctx.form)
        vdui.get_current_item(USE_KEYBOARD)
        if vdui.item.is_citem():
            print vdui.item.it.index
        self.hrCbObj.walk_ctree(vdui.cfunc)


        return 1

    def update(self, ctx):
        #print "Update"
        vdui = idaapi.get_tform_vdui(ctx.form)
        if vdui:
            return idaapi.AST_ENABLE_FOR_FORM
        else:
            return idaapi.AST_DISABLE_FOR_FORM

class ItemContainer(object):
    def __init__(self,item,parent,item_is_expr,parent_is_expr):
        self.parent = parent
        self.item = item
        self.parent_is_expr = parent_is_expr
        self.item_is_expr = item_is_expr

class hexrays_callback_info(object):

    def __init__(self):
        self.edges = []
        self.vu = None
        self.items = []
        self.G=nx.DiGraph()
        self.cfunc = None


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

        return

    def load(self):

        self.stored = []

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

        return

    def save(self):

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

        return

    def get_parent_idx(self,target_item):
        for item in self.items:
            if target_item.parent is not None:
                if target_item.parent_is_expr == item.item_is_expr:
                    if item.item == target_item.parent:
                        return self.items.index(item)
            else:
                break
        return None

    def get_item_container(self,item):
        for it in self.items:
            if it.item_is_expr == item.is_expr() and it.item == item:
                return it
        return None

    def get_label(self,n):
        item = self.items[n]
        label = "%d: "%n
        label += idaapi.get_ctype_name(item.op)
        if item.op == idaapi.cot_ptr:
            label += ".%d"%item.cexpr.ptrsize
        elif item.op == idaapi.cot_memptr:
            label += ".%d"%item.cexpr.ptrsize
            label += "m=%d\n"%item.cexpr.m
            label += idaapi.tag_remove(item.cexpr.print1(None))
        elif item.op == idaapi.cot_memref:
            label += "m=%d"%item.cexpr.m
        elif item.op == idaapi.cot_obj or item.op == idaapi.cot_var:
            label += ".%d"%item.cexpr.refwidth
            label += " "
            label += idaapi.tag_remove(item.cexpr.print1(None))
            # print
        elif item.op == idaapi.cot_num or item.op == idaapi.cot_helper or item.op == idaapi.cot_str:
            label += " "
            label += idaapi.tag_remove(item.cexpr.print1(None))
        elif item.op == idaapi.cit_goto:
            label += "LABEL_%d"%item.cinsn.cgoto.label_num
        elif item.op == idaapi.cit_asm:
            label += "%s"%item.cinsn.casm.__repr__()

        label += "\nea: 0x%08X"%item.ea
        if item.is_expr() and not item.cexpr.type.empty():
            label += "\n"
            #print item.cexpr.type._print()
            prefix = ""
            indent = 0
            cmtindent = 0
            flags = idaapi.PRTYPE_1LINE
            name = ""
            cmt = ""
            idaapi.print_tinfo(prefix,indent,cmtindent,flags,item.cexpr.type,name,cmt)
            #print prefix,indent,cmtindent,name,cmt
            label += "%s"%item.cexpr.type
        return label

    def walk_ctree(self, cfunc):
        self.items = []
        self.cfunc = cfunc
        self.G.clear()
        #if fDebug:
        #    pydevd.settrace('localhost', port=2255, stdoutToServer=True, stderrToServer=True)
        print "walk_ctree"
        root = None

        class visitor(idaapi.ctree_visitor_t):

            def __init__(self,obj, cfunc):
                idaapi.ctree_visitor_t.__init__(self,idaapi.CV_PARENTS)
                self.cfunc = cfunc
                self.obj = obj

                return

            def visit_insn(self, i):
                #print "\nvisit_insn"
                #print "item:"
                #print i,'\n'
                parent = None
                #print "Parents:\n",self.parents
                #print "Items:\n",self.obj.items
                if len(self.parents) > 1:

                    parent = self.parents[len(self.parents)-1]
                    # f = self.obj.cfunc.body.find_parent_of(i)
                    #parent =
                    #print parent.cinsn
                    #print parent in self.obj.items
                if parent is not None:
                    self.obj.edges.append((parent.index, i.index))
                self.obj.items.append(i)
                self.obj.G.add_node(i.index)
                return 0 # continue enumeration
            def visit_expr(self, i):
                #print "\nvisit_expr"
                #print "item:"
                #print i,'\n'
                parent = None
                #self.obj.items.append(i)
                #self.obj.G.add_node(self.obj.items.index(i))
                #print "Parents:\n",self.parents
                #print "Items:\n",self.obj.items
                if len(self.parents) > 1:
                    parent = self.parents[len(self.parents)-1]
                    # f = self.obj.cfunc.body.find_parent_of(i)
                    # parent = i.find_parent_of(i)
                    #print parent.cinsn
                if parent is not None:
                    self.obj.edges.append((parent.index, i.index))
                self.obj.items.append(i)
                self.obj.G.add_node(i.index)
                return 0

        visitor(self,cfunc).apply_to(cfunc.body, None)
        #print len(self.items)
        root = None
        # for item in self.items:
        #     #print item.item.opname
        #     #print dir(self)
        #     p = self.get_parent_idx(item)
        #     if p is None:
        #         assert root is None
        #         root = item
        #     else:
        #         self.G.add_edge(p,self.items.index(item))
        for ed in self.edges:
            self.G.add_edge(*ed)
        #pos = hierarchy_pos(self.G,self.items.index(root))
        #nx.draw(self.G,pos,node_shape="s")
        #plt.show() # display
        for u,v,d in self.G.edges(data=True):
            a = self.items[u]
            b = self.items[v]
            if a.is_expr():
                if type(a.x) == type(b) and a.x == b: d["label"] = "x"
                if type(a.y) == type(b) and a.y == b: d["label"] = "y"
                if type(a.z) == type(b) and a.z == b: d["label"] = "z"
        labels = {}
        for j in range(len(self.items)):
            labels[j] = self.get_label(j)
        #print labels
        nx.relabel_nodes(self.G,labels,copy=False)
        A = to_agraph(self.G)
        A.layout('dot')
        # A.draw('graph_test.ps')
        A.draw('graph_test.png')
        # image = mpimg.imread('graph_test.png')
        # plt.axis("off")
        # plt.imshow(image)
        # plt.show()




        return

    def event_callback(self, event, *args):
        #print "event_callback"
        if event == idaapi.hxe_populating_popup:
            #print "event_callback: hxe_populating_popup"
            form, phandle, vu = args
            res = idaapi.attach_action_to_popup(vu.ct, None, print_ctree_actname)

        elif event == idaapi.hxe_maturity:
            #print "event_callback: hxe_maturity"
            cfunc, maturity = args
            #if maturity == idaapi.CMAT_FINAL:
                #print "event_callback: hxe_maturity: CMAT_FINAL"
                #self.walk_ctree(cfunc)

        return 0


class print_ctree_plugin_t(plugin_t):
    flags = PLUGIN_HIDE
    comment = "This is a comment"

    help = "This is help"
    wanted_name = "cobject_helper"
    #wanted_hotkey = "Alt-F8"
    wanted_hotkey = ""

    def init(self):
        #msg("StructTyper init() called!\n")
        if fDebug:
            pydevd.settrace('localhost', port=2255, stdoutToServer=True, stderrToServer=True, suspend= False)
        ret = idaapi.init_hexrays_plugin()
        # print ret
        if not idaapi.init_hexrays_plugin():
            ret = idaapi.load_plugin("hexrays")
            # print ret
            ret = idaapi.init_hexrays_plugin()
            # print ret
        if idaapi.init_hexrays_plugin():
            i = hexrays_callback_info()
            idaapi.register_action(
                idaapi.action_desc_t(
                    print_ctree_actname,
                    "Print HexRays ctree",
                    print_ctree_action_handler_t(i),
                    ""))
            idaapi.install_hexrays_callback(i.event_callback)
            return PLUGIN_KEEP
        else:
            print 'print_ctree: hexrays is not available.'
            return PLUGIN_SKIP

    def run(self, arg):
        #msg("StructTyper run() called with %d!\n" % arg)
        #require('flare')
        #require('flare.struct_typer')
        #struct_typer.main()
        pass


    def term(self):
        #idaapi.msg("StructTyper term() called!\n")
        pass

def PLUGIN_ENTRY():
    return print_ctree_plugin_t()


repository: doronz88/fa
file: fa/ida_plugin.py
description: Automation tool for locating symbols & structs in binary (primary IDA focused)
url: https://github.com/doronz88/fa/blob/1fd337496635719fb3658feddc67d16b9d20c87e/fa/ida_plugin.py
wanted_name: FA
created_at: 2020-04-16 12:35:00+00:00
pushed_at: 2025-04-04 18:19:38+00:00
content:
import binascii
import os
import re
import subprocess
import sys
import tempfile
import traceback
from collections import namedtuple

import rpyc
from rpyc import OneShotServer

sys.path.append('.')  # noqa: E402

import click
import hjson
import ida_auto
import ida_bytes
import ida_ida
import ida_kernwin
import ida_pro
import ida_segment
import ida_typeinf
import idaapi
import idautils
import idc
from ida_kernwin import Form

from fa import fa_types, fainterp

# Filename for the temporary created signature
TEMP_SIG_FILENAME = os.path.join(tempfile.gettempdir(), 'fa_tmp_sig.sig')

# IDA fa plugin filename
PLUGIN_FILENAME = 'fa_ida_plugin.py'


def open_file(filename):
    """
    Attempt to open the given filename by OS' default editor
    :param filename: filename to open
    :return: None
    """
    if sys.platform == "win32":
        try:
            os.startfile(filename)
        except Exception as error_code:
            if error_code[0] == 1155:
                os.spawnl(os.P_NOWAIT,
                          os.path.join(os.environ['WINDIR'],
                                       'system32', 'Rundll32.exe'),
                          'Rundll32.exe SHELL32.DLL, OpenAs_RunDLL {}'
                          .format(filename))
            else:
                print("other error")
    else:
        opener = "open" if sys.platform == "darwin" else "xdg-open"
        subprocess.call([opener, filename])


class IdaLoader(fainterp.FaInterp):
    """
    IDA loader
    Includes improved GUI interaction for accessing the interpreter
    functionality.
    """

    def __init__(self):
        super(IdaLoader, self).__init__()
        self._create_template_symbol = eval(self.config_get(
            'global',
            'create_symbol_template',
            'False'
        ))

    def set_symbol_template(self, status):
        """
        Should the create-temp-signature feature attempt to create a default
        signature by predefined template?
        :param status: new boolean setting
        :return: None
        """
        self._create_template_symbol = status
        self.config_set('global', 'create_symbol_template', str(status))

    def create_symbol(self):
        """
        Create a temporary symbol signature from the current function on the
        IDA screen.
        """
        self.log('creating temporary signature')

        current_ea = idc.get_screen_ea()

        signature = {
            'name': idc.get_func_name(current_ea),
            'instructions': []
        }

        if self._create_template_symbol:
            find_bytes_ida = "find-bytes-ida '"

            for ea in idautils.FuncItems(current_ea):
                mnem = idc.print_insn_mnem(ea).lower()
                opcode_size = idc.get_item_size(ea)

                # ppc
                if mnem.startswith('b') or mnem in ('lis', 'lwz', 'addi'):
                    # relative opcodes
                    find_bytes_ida += '?? ' * opcode_size
                    continue

                # arm
                if mnem.startswith('b') or mnem in ('ldr', 'str'):
                    # relative opcodes
                    find_bytes_ida += '?? ' * opcode_size
                    continue

                opcode = binascii.hexlify(idc.get_bytes(ea, opcode_size))
                formatted_hex = ' '.join(opcode[i:i + 2].decode() for i in
                                         range(0, len(opcode), 2))
                find_bytes_ida += formatted_hex + ' '

            find_bytes_ida += "'"

            signature['instructions'].append(find_bytes_ida)
            signature['instructions'].append('function-start')
            signature['instructions'].append('set-name "{}"'.format(
                idc.get_func_name(current_ea)))

        with open(TEMP_SIG_FILENAME, 'w') as f:
            hjson.dump(signature, f, indent=4)

        self.log('Signature created at {}'.format(TEMP_SIG_FILENAME))
        return TEMP_SIG_FILENAME

    def extended_create_symbol(self):
        """
        Creates a temporary symbol of the currently active function
        and open it using OS default editor
        :return: None
        """
        filename = self.create_symbol()
        open_file(filename)

    def find_symbol(self):
        """
        Find the last create symbol signature.
        :return: None
        """
        with open(TEMP_SIG_FILENAME) as f:
            sig = hjson.load(f)

        results = self.find_from_sig_json(sig, decremental=False)

        for address in results:
            self.log('Search result: 0x{:x}'.format(address))
        self.log('Search done')

        if len(results) == 1:
            # if remote sig has a proper name, but current one is not
            ida_kernwin.jumpto(results[0])

    def verify_project(self):
        """
        Verify a valid project is currently active.
        Show IDA warning if not.
        :return: None
        """
        try:
            super(IdaLoader, self).verify_project()
        except IOError as e:
            ida_kernwin.warning(str(e))
            raise e

    def prompt_save_signature(self):
        """
        Save last-created-temp-signature if user agrees to in IDA prompt
        :return: None
        """
        self.verify_project()

        if ida_kernwin.ask_yn(1, 'Are you sure you want '
                                 'to save this signature?') != 1:
            return

        self.save_signature(TEMP_SIG_FILENAME)

    def find(self, symbol_name, use_cache=False):
        """
        Find symbol by name (as specified in SIG file)
        Show an IDA waitbox while doing so
        :param symbol_name: symbol name
        :return: output address list
        """
        ida_kernwin.replace_wait_box('Searching symbol: \'{}\'...'
                                     .format(symbol_name))
        return super(IdaLoader, self).find(symbol_name, use_cache=use_cache)

    def get_python_symbols(self, file_name=None):
        """
        Run all python scripts inside the currently active project.
        Show an IDA waitbox while doing so
        :param file_name: filter a specific filename to execute
        :return: dictionary of all found symbols
        """
        ida_kernwin.replace_wait_box('Running python scripts...')
        return super(IdaLoader, self).get_python_symbols(file_name=file_name)

    @staticmethod
    def extract_all_user_names(filename=None):
        """
        Get all user-named labels inside IDA. Also prints into output window.
        :return: dictionary of all user named labels: label_name -> ea
        """
        results = {}
        output = ''

        for ea, name in idautils.Names():
            if ida_kernwin.user_cancelled():
                return results

            if '_' in name:
                if name.split('_')[0] in ('def', 'sub', 'loc', 'jpt', 'j',
                                          'nullsub'):
                    continue
            flags = ida_bytes.get_full_flags(ea)
            if ida_bytes.has_user_name(flags):
                results[name] = ea
                output += '{} = 0x{:08x};\n'.format(name, ea)

        if filename is not None:
            with open(filename, 'w') as f:
                f.write(output)

        return results

    def set_const(self, name, value):
        super(IdaLoader, self).set_const(name, value)
        fa_types.add_const(name, value)

    def set_symbol(self, symbol_name, value):
        super(IdaLoader, self).set_symbol(symbol_name, value)
        idc.set_name(value, symbol_name, idc.SN_CHECK)

    def symbols(self, output_file_path=None):
        """
        Run find for all SIG files in currently active project.
        Show an IDA waitbox while doing so
        :param output_file_path: optional, save found symbols into output file
        :return: dictionary of found symbols
        """
        self.verify_project()
        results = {}

        try:
            ida_kernwin.show_wait_box('Searching...')
            results = super(IdaLoader, self).symbols()

            ida_kernwin.replace_wait_box('Extracting...')
            ida_symbols = IdaLoader.extract_all_user_names(output_file_path)

            results.update(ida_symbols)

        except Exception:
            traceback.print_exc()
        finally:
            ida_kernwin.hide_wait_box()

        return results

    def export(self):
        """
        Show an export dialog to export symbols and header file for given
        IDB.
        :return: None
        """

        class ExportForm(Form):
            def __init__(self):
                description = '''
                <h2>Export</h2>

                Select a directory to export IDB data into.
                '''

                Form.__init__(self,
                              r"""BUTTON YES* Save
                              Export
                              {StringLabel}
                              <#Symbols#Symbols filename:{iSymbolsFilename}>
                              <#C Header#C Header filename:{iHeaderFilename}>
                              <#ifdef_macro#ifdef'ed:{iIfdef}>
                              <#Select dir#Browse for dir:{iDir}>
                              """, {
                                  'iDir': Form.DirInput(),
                                  'StringLabel':
                                      Form.StringLabel(description,
                                                       tp=Form.FT_HTML_LABEL),
                                  'iSymbolsFilename': Form.StringInput(
                                      value='symbols.txt'),
                                  'iHeaderFilename': Form.StringInput(
                                      value='fa_structs.h'),
                                  'iIfdef': Form.StringInput(
                                      value='FA_STRUCTS_H'),
                              })
                self.__n = 0

            def OnFormChange(self, fid):
                return 1

        form = ExportForm()
        form, args = form.Compile()
        ok = form.Execute()
        if ok == 1:
            # save symbols
            symbols_filename = os.path.join(form.iDir.value,
                                            form.iSymbolsFilename.value)
            with open(symbols_filename, 'w') as f:
                results = IdaLoader.extract_all_user_names(None)
                for k, v in sorted(results.items()):
                    f.write('{} = 0x{:08x};\n'.format(k, v))

            # save c header
            idati = ida_typeinf.get_idati()
            c_header_filename = os.path.join(form.iDir.value,
                                             form.iHeaderFilename.value)

            consts_ordinal = None
            ordinals = []
            for ordinal in range(1, ida_typeinf.get_ordinal_qty(idati) + 1):
                ti = ida_typeinf.tinfo_t()
                if ti.get_numbered_type(idati, ordinal):
                    if ti.get_type_name() == 'FA_CONSTS':
                        # convert into macro definitions
                        consts_ordinal = ordinal
                    elif ti.get_type_name() in ('__va_list_tag',
                                                'va_list'):
                        continue
                    elif '$' in ti.get_type_name():
                        # skip deleted types
                        continue
                    else:
                        ordinals.append(str(ordinal))

            with open(c_header_filename, 'w') as f:
                ifdef_name = form.iIfdef.value.strip()

                if len(ifdef_name) > 0:
                    f.write('#ifndef {ifdef_name}\n'
                            '#define {ifdef_name}\n\n'
                            .format(ifdef_name=ifdef_name))

                if consts_ordinal is not None:
                    consts = re.findall('\\s*(.+?) = (.+?),',
                                        idc.print_decls(
                                            str(consts_ordinal), 0))
                    for k, v in consts:
                        f.write('#define {} ({})\n'.format(k, v))

                    # ida exports using this type
                    f.write('#define _BYTE char\n')
                    f.write('\n')

                structs_buf = idc.print_decls(','.join(ordinals),
                                              idc.PDF_DEF_BASE)

                for struct_type, struct_name in re.findall(
                        r'(struct|enum) .*?([a-zA-Z0-9_\-]+?)\s+\{',
                        structs_buf):
                    f.write(
                        'typedef {struct_type} {struct_name} {struct_name};\n'
                        .format(struct_type=struct_type,
                                struct_name=struct_name))

                structs_buf = structs_buf.replace('__fastcall', '')
                f.write('\n')
                f.write(structs_buf)
                f.write('\n')

                if len(ifdef_name) > 0:
                    f.write('#endif // {ifdef_name}\n'
                            .format(ifdef_name=ifdef_name))

        form.Free()

    def set_input(self, input_):
        """
        Mock for change_input. Just reload current loaded data settings.
        :param input_: doesn't matter
        :return: None
        """
        self.endianity = '>' if ida_ida.inf_is_be() else '<'
        self._input = input_
        self.reload_segments()

    def reload_segments(self):
        """
        memory searches will use IDA's API instead
        which is much faster so this is just a stub.
        :return: None
        """
        return

    def interactive_settings(self):
        """
        Show settings dialog
        :return: None
        """

        class SettingsForm(Form):
            def __init__(self, signatures_root, use_template):
                description = '''
                <h2>Settings</h2>
                <div>
                Here you can change global FA settings.
                </div>
                <div>
                    <a href="https://github.com/doronz88/fa">
                    For more info</a>
                </div>
                '''

                Form.__init__(self,
                              r"""BUTTON YES* Save
                              FA Settings
                              {{FormChangeCb}}
                              {{StringLabel}}
                              <Signatures root :{{signaturesRoot}}>
                              <Temp signature generation :{{signatureGeneration}}>
                              """.format(signatures_root), {
                                  'FormChangeCb':
                                      Form.FormChangeCb(self.OnFormChange),
                                  'signaturesRoot':
                                      Form.DirInput(value=signatures_root),
                                  'StringLabel':
                                      Form.StringLabel(description,
                                                       tp=Form.FT_HTML_LABEL),
                                  'signatureGeneration':
                                      Form.DropdownListControl(
                                          items=['Default',
                                                 'Using function bytes'],
                                          readonly=True,
                                          selval=use_template),
                              })
                self.__n = 0

            def OnFormChange(self, fid):
                return 1

        f = SettingsForm(self._signatures_root, self._create_template_symbol)
        f, args = f.Compile()
        ok = f.Execute()
        if ok == 1:
            self.set_signatures_root(f.signaturesRoot.value, save=True)
            self.set_symbol_template(f.signatureGeneration.value == 1)
        f.Free()

    def interactive_set_project(self):
        """
        Show set-project dialog
        :return: None
        """

        class SetProjectForm(Form):
            def __init__(self, signatures_root, projects, current):
                description = '''
                <h2>Project Selector</h2>
                <div>
                Select project you wish to work on from your
                signatures root:
                </div>
                <div><pre>{}</pre></div>
                <div><i>(Note: You may change this in config.ini)</i></div>
                <div>
                    <a href="https://github.com/doronz88/fa#projects">
                    For more info</a>
                </div>
                '''.format(signatures_root)

                Form.__init__(self,
                              r"""BUTTON YES* OK
                              FA Project Select
                              {{FormChangeCb}}
                              {{StringLabel}}
                              <Set Project :{{cbReadonly}}>
                              """.format(signatures_root), {
                                  'FormChangeCb':
                                      Form.FormChangeCb(self.OnFormChange),
                                  'cbReadonly':
                                      Form.DropdownListControl(
                                          items=projects,
                                          readonly=True,
                                          selval=projects.index(current)
                                          if current in projects else 0),
                                  'StringLabel':
                                      Form.StringLabel(description,
                                                       tp=Form.FT_HTML_LABEL),
                              })
                self.__n = 0

            def OnFormChange(self, fid):
                return 1

        projects = self.list_projects()
        f = SetProjectForm(self._signatures_root, projects, self._project)
        f, args = f.Compile()
        ok = f.Execute()
        if ok == 1:
            self.set_project(projects[f.cbReadonly.value])
        f.Free()


fa_instance = None

Action = namedtuple('action', 'name icon_filename handler label hotkey')


def add_action(action):
    """
    Add an ida-action
    :param action: action given as the `Action` namedtuple
    :return: None
    """

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

        def activate(self, ctx):
            action.handler()
            return 1

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

    act_icon = -1
    if action.icon_filename:
        icon_full_filename = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'res', 'icons', action.icon_filename)
        with open(icon_full_filename, 'rb') as f:
            icon_data = f.read()
        act_icon = ida_kernwin.load_custom_icon(data=icon_data, format="png")

    act_name = action.name

    ida_kernwin.unregister_action(act_name)
    if ida_kernwin.register_action(ida_kernwin.action_desc_t(
            act_name,  # Name. Acts as an ID. Must be unique.
            action.label,  # Label. That's what users see.
            Handler(),  # Handler. Called when activated, and for updating
            action.hotkey,  # Shortcut (optional)
            None,  # Tooltip (optional)
            act_icon)):  # Icon ID (optional)

        # Insert the action in the menu
        if not ida_kernwin.attach_action_to_menu(
                "FA/", act_name, ida_kernwin.SETMENU_APP):
            print("Failed attaching to menu.")

        # Insert the action in a toolbar
        if not ida_kernwin.attach_action_to_toolbar("fa", act_name):
            print("Failed attaching to toolbar.")

        class Hooks(ida_kernwin.UI_Hooks):
            def finish_populating_widget_popup(self, widget, popup):
                if ida_kernwin.get_widget_type(widget) == \
                        ida_kernwin.BWN_DISASM:
                    ida_kernwin.attach_action_to_popup(widget,
                                                       popup,
                                                       act_name,
                                                       None)

        hooks = Hooks()
        hooks.hook()


def load_ui():
    """
    Load FA's GUI buttons
    :return: None
    """
    actions = [
        Action(name='fa:settings',
               icon_filename='settings.png',
               handler=fa_instance.interactive_settings,
               label='Settings',
               hotkey=None),

        Action(name='fa:set-project',
               icon_filename='suitcase.png',
               handler=fa_instance.interactive_set_project,
               label='Set project...',
               hotkey='Ctrl+6'),

        Action(name='fa:symbols', icon_filename='find_all.png',
               handler=fa_instance.symbols,
               label='Find all project\'s symbols',
               hotkey='Ctrl+7'),

        Action(name='fa:export', icon_filename='export.png',
               handler=fa_instance.export,
               label='Export symbols',
               hotkey=None),

        Action(name='fa:extended-create-signature',
               icon_filename='create_sig.png',
               handler=fa_instance.extended_create_symbol,
               label='Create temp signature...',
               hotkey='Ctrl+8'),

        Action(name='fa:find-symbol',
               icon_filename='find.png',
               handler=fa_instance.find_symbol,
               label='Find last created temp signature',
               hotkey='Ctrl+9'),

        Action(name='fa:prompt-save',
               icon_filename='save.png',
               handler=fa_instance.prompt_save_signature,
               label='Save last created temp signature',
               hotkey='Ctrl+0'),
    ]

    # init toolbar
    ida_kernwin.delete_toolbar('fa')
    ida_kernwin.create_toolbar('fa', 'FA Toolbar')

    # init menu
    ida_kernwin.delete_menu('fa')
    ida_kernwin.create_menu('fa', 'FA')

    for action in actions:
        add_action(action)


def install():
    """
    Install FA ida plugin
    :return: None
    """
    fa_plugin_dir = os.path.join(
        idaapi.get_user_idadir(), 'plugins')

    if not os.path.exists(fa_plugin_dir):
        os.makedirs(fa_plugin_dir)

    fa_plugin_filename = os.path.join(fa_plugin_dir, PLUGIN_FILENAME)
    if os.path.exists(fa_plugin_filename):
        IdaLoader.log('already installed')
        return

    with open(fa_plugin_filename, 'w') as f:
        f.writelines("""from __future__ import print_function
try:
    from fa.ida_plugin import PLUGIN_ENTRY, FAIDAPlugIn
except ImportError:
    print("[WARN] Could not load FA plugin. "
          "FA Python package doesn\'t seem to be installed.")
""")

    idaapi.load_plugin(PLUGIN_FILENAME)

    IdaLoader.log('Successfully installed :)')


@click.command()
@click.option('-s', '--service', type=click.IntRange(1024, 65535), help='execute in rpyc service mode at given port')
def main(service):
    plugin_main(service)


class FaService(rpyc.Service):
    ida_segment = ida_segment
    ida_kernwin = ida_kernwin
    ida_typeinf = ida_typeinf
    ida_bytes = ida_bytes
    idautils = idautils
    ida_auto = ida_auto
    ida_pro = ida_pro
    ida_ida = ida_ida
    idaapi = idaapi
    idc = idc

    @staticmethod
    def load_module(name, filename):
        return fainterp.FaInterp.get_module(name, filename)


def plugin_main(service=None):
    global fa_instance

    fa_instance = IdaLoader()
    fa_instance.set_input('ida')

    load_ui()

    IdaLoader.log('''
    ---------------------------------
    FA Loaded successfully

    Quick usage:
    print(fa_instance.find(symbol_name)) # searches for the specific symbol
    fa_instance.get_python_symbols(filename=None) # run project's python
                                                    scripts (all or single)
    fa_instance.set_symbol_template(status) # enable/disable template temp
                                              signature
    fa_instance.symbols() # searches for the symbols in the current project
    ---------------------------------
    ''')

    if service:
        t = OneShotServer(FaService, port=service, protocol_config={
            'allow_all_attrs': True,
            'allow_setattr': True,
        })
        t.start()

    # TODO: consider adding as autostart script
    # install()


try:
    class FAIDAPlugIn(idaapi.plugin_t):
        wanted_name = "FA"
        wanted_hotkey = "Shift-,"
        flags = 0
        comment = ""
        help = "Load FA in IDA Pro"

        def init(self):
            plugin_main()
            return idaapi.PLUGIN_KEEP

        def run(self, args):
            pass

        def term(self):
            pass
except TypeError:
    print('ignoring rpyc bug')


def PLUGIN_ENTRY():
    """
    Entry point for IDA plugins
    :return:
    """
    return FAIDAPlugIn()


if __name__ == '__main__':
    # Entry point for IDA in script mode (-S)
    main(standalone_mode=False, args=idc.ARGV[1:])


repository: airbus-seclab/AutoResolv
file: AutoResolv.py
wanted_name: AutoResolv
comment: Resolve imports to find libname/paths and refactor call type from signature file
created_at: 2022-09-27 08:05:43+00:00
pushed_at: 2022-10-15 20:36:55+00:00
url:
https://github.com/airbus-seclab/AutoResolv/blob/de3169609a8630151a6c2382b407493b49f4a202/AutoResolv.py

content:
# This file is part of AutoResolv.
# Copyright 2022 - Airbus, thibault poncetta
# AutoResolv 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.
# AutoResolv 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 AutoResolv.  If not, see <http://www.gnu.org/licenses/>.


import idaapi
import idc
import sys
import idautils
import ida_funcs
import time
import os
import subprocess
import sys
from collections import defaultdict

from libautoresolv.resultshower import ResultShower
from libautoresolv.util import *
from libautoresolv.error import *
from libautoresolv.dbcache import *

from libautoresolv.GUI.gui_main import *
from libautoresolv.GUI.gui_start import *



VERSION = "dev-v0.90p"
idaapi.require("libautoresolv.GUI.gui_main")
idaapi.require("libautoresolv.GUI.gui_export")
idaapi.require("libautoresolv.GUI.gui_start")



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

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

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

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

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

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

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

class Searcher(Kp_Menu_Context):
    def activate(self, ctx):
        self.plugin.search()
        return 1



p_initialized = False
#--------------------------------------------------------------------------
# Plugin
#--------------------------------------------------------------------------
class AutoResolv(idaapi.plugin_t):
    comment = "Resolve imports to find libname/paths and refactor call type from signature file"
    help = "AutoResolv Help"
    wanted_name = "AutoResolv"
    flags = idaapi.PLUGIN_KEEP
    wanted_hotkey= "Ctrl+Shift-A"

    def init(self):
        global p_initialized

        # register popup menu handlers
        try:
            Searcher.register(self, "AutoResolv")
        except:
            pass

        if p_initialized is False:
            p_initialized = True
            idaapi.register_action(idaapi.action_desc_t(
                "AutoResolv",
                "AutoResolv imports",
                Searcher(),
                None,
                None,
                0))
            idaapi.attach_action_to_menu("Search", "AutoResolv", idaapi.SETMENU_APP)
            print("*" * 80)
            print(f"AutoResolv {VERSION} | Thibault Poncetta")
            print("*" * 80)

        return idaapi.PLUGIN_KEEP

    def term(self):
        pass

    def _decompile_then_write_on_fd(self, fd, funs):
        lenf = len(funs)
        cpt = 0
        for fun in funs:
            ea = funs[fun]
            cfunc = str(idaapi.decompile(ea)).split("\n")[0] + ";"
            print(f"Exporting {cfunc}, Progression: {(cpt/lenf)*100}%")
            cpt += 1
            fd.write(cfunc + "\n")

    def _signature_export(self, fd):
        
        start,end = get_seg(".text")   
        if start == None or end == None:
            print("[AutoResolv] Error when parsing Segments() address.")
            
        funs = get_funs(start, end)
        nb_fun = len(funs)
        print(f"[AutoResolv] Exporting {nb_fun} function from .text segment")

        self._decompile_then_write_on_fd(fd, funs)
        print(f"[AutoResolv] Done")


    def main(self):
        print(f"AutoResolv {VERSION}")



        module_path = os.path.dirname(__file__) + "/libautoresolv/db/"
        binary_name = idaapi.get_root_filename()
        db_path = module_path + ".cache_" + binary_name + ".db"
        bin_path = os.getcwd() + "/" +  binary_name

        cache = DB_CACHE_MANAGER(db_path, module_path, bin_path)
        con = cache.check_cache_con()
        if con:
            cache.parse_conf_cache()
            cache.parse_libinfo_cache()
            cache.parse_rpath_cache()
            cache.parse_data_cache()
            cache.parse_bininfo_cache()
            bin_path = cache.bin_path

        if not os.path.exists(bin_path):
            

            gui_start = GUI_START(bin_path)
            gui_start.exec_()

            newpath = gui_start.newpath
            if newpath is None:
                raise Exception("[AutoResolv] No path given  ! Aborting")
            if not os.path.exists(newpath):
                raise Exception("[AutoResolv] Given path for binary location incorrect ! Aborting")
            else:
                print(f"[AutoResolv] Set new binary path to {newpath}")
                bin_path = newpath

        if not con:
            print(f"[AutoResolv] DB Cache seem empty, Creating cache")
            libs,rpath = getLibsFromBin(bin_path)
            if rpath is not None:
                print(f"[AutoResolv] I got a rpath : {rpath}")
                cache.create_cache(libs, bin_path, rpath)
                cache.rpath = rpath.split(":")
            else:
                cache.create_cache(libs, bin_path)
                cache.rpath = []

        cache.cache_save_bininfo(bin_path)
        gui_main = GUI_MAIN(cache)
        gui_main.exec_()
        
    
    def run(self, arg):
        self.main()


def PLUGIN_ENTRY():
    return AutoResolv()


