Replace HexRaysCodeXplorer and HexraysInvertIf with HexRaysPyTools
This commit is contained in:
ecx86
536
plugins/HexRaysPyTools/Core/Classes.py
Normal file
536
plugins/HexRaysPyTools/Core/Classes.py
Normal file
@@ -0,0 +1,536 @@
|
||||
import idc
|
||||
import idaapi
|
||||
import HexRaysPyTools.Forms
|
||||
import Helper
|
||||
|
||||
# from PySide import QtGui, QtCore
|
||||
from HexRaysPyTools.Cute import *
|
||||
|
||||
all_virtual_functions = {} # name -> VirtualMethod
|
||||
all_virtual_tables = {} # ordinal -> VirtualTable
|
||||
|
||||
|
||||
class VirtualMethod(object):
|
||||
def __init__(self, tinfo, name, parent):
|
||||
self.tinfo = tinfo
|
||||
self.tinfo_modified = False
|
||||
self.name = name
|
||||
self.class_name = None
|
||||
self.name_modified = False
|
||||
self.parents = [parent]
|
||||
self.base_address = Helper.get_virtual_func_address(name)
|
||||
if self.base_address:
|
||||
self.base_address -= idaapi.get_imagebase()
|
||||
|
||||
self.rowcount = 0
|
||||
self.children = []
|
||||
|
||||
@staticmethod
|
||||
def create(tinfo, name, parent):
|
||||
result = all_virtual_functions.get(name)
|
||||
if result:
|
||||
result.parents.append(parent)
|
||||
return result
|
||||
result = VirtualMethod(tinfo, name, parent)
|
||||
all_virtual_functions[name] = result
|
||||
return result
|
||||
|
||||
def update(self, name, tinfo):
|
||||
self.name = name
|
||||
self.tinfo = tinfo
|
||||
self.name_modified = False
|
||||
self.tinfo_modified = False
|
||||
|
||||
self.base_address = idc.LocByName(self.name)
|
||||
if self.base_address != idaapi.BADADDR:
|
||||
self.base_address -= idaapi.get_imagebase()
|
||||
|
||||
def data(self, column):
|
||||
if column == 0:
|
||||
return self.name
|
||||
elif column == 1:
|
||||
return self.tinfo.get_pointed_object().dstr()
|
||||
elif column == 2:
|
||||
return "0x{0:08X}".format(self.address) if self.address else None
|
||||
|
||||
def setData(self, column, value):
|
||||
if column == 0:
|
||||
if idaapi.isident(value) and self.name != value:
|
||||
self.name = value
|
||||
self.name_modified = True
|
||||
for parent in self.parents:
|
||||
parent.modified = True
|
||||
return True
|
||||
elif column == 1:
|
||||
tinfo = idaapi.tinfo_t()
|
||||
split = value.split('(')
|
||||
if len(split) == 2:
|
||||
value = split[0] + ' ' + self.name + '(' + split[1] + ';'
|
||||
if idaapi.parse_decl2(idaapi.cvar.idati, value, '', tinfo, idaapi.PT_TYP):
|
||||
if tinfo.is_func():
|
||||
tinfo.create_ptr(tinfo)
|
||||
if tinfo.dstr() != self.tinfo.dstr():
|
||||
self.tinfo = tinfo
|
||||
self.tinfo_modified = True
|
||||
for parent in self.parents:
|
||||
parent.modified = True
|
||||
return True
|
||||
return False
|
||||
|
||||
def font(self, column):
|
||||
if column == 0 and self.name_modified:
|
||||
return QtGui.QFont("Consolas", 10, italic=True)
|
||||
elif column == 1 and self.tinfo_modified:
|
||||
return QtGui.QFont("Consolas", 10, italic=True)
|
||||
return QtGui.QFont("Consolas", 10, 0)
|
||||
|
||||
def flags(self, column):
|
||||
if column != 2:
|
||||
if self.address:
|
||||
return QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsEditable
|
||||
return QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEnabled
|
||||
|
||||
|
||||
@property
|
||||
def color(self):
|
||||
# return QtGui.QBrush(QtGui.QColor("#ffffb3"))
|
||||
return QtGui.QColor("#fefbd8")
|
||||
|
||||
@property
|
||||
def tooltip(self):
|
||||
return None
|
||||
|
||||
@property
|
||||
def address(self):
|
||||
return self.base_address + idaapi.get_imagebase() if self.base_address else None
|
||||
|
||||
def set_first_argument_type(self, name):
|
||||
func_data = idaapi.func_type_data_t()
|
||||
func_tinfo = self.tinfo.get_pointed_object()
|
||||
class_tinfo = idaapi.tinfo_t()
|
||||
if func_tinfo.get_func_details(func_data) and func_tinfo.get_nargs() and \
|
||||
class_tinfo.get_named_type(idaapi.cvar.idati, name):
|
||||
class_tinfo.create_ptr(class_tinfo)
|
||||
first_arg_tinfo = func_data[0].type
|
||||
if (first_arg_tinfo.is_ptr() and first_arg_tinfo.get_pointed_object().is_udt()) or \
|
||||
Helper.is_legal_type(func_data[0].type):
|
||||
func_data[0].type = class_tinfo
|
||||
func_data[0].name = "this"
|
||||
func_tinfo.create_func(func_data)
|
||||
func_tinfo.create_ptr(func_tinfo)
|
||||
if func_tinfo.dstr() != self.tinfo.dstr():
|
||||
self.tinfo = func_tinfo
|
||||
self.tinfo_modified = True
|
||||
for parent in self.parents:
|
||||
parent.modified = True
|
||||
else:
|
||||
print "[Warning] function {0} probably have wrong type".format(self.name)
|
||||
|
||||
def open_function(self):
|
||||
if self.address:
|
||||
if idaapi.decompile(self.address):
|
||||
idaapi.open_pseudocode(self.address, 0)
|
||||
else:
|
||||
idaapi.jumpto(self.address)
|
||||
|
||||
def commit(self):
|
||||
if self.name_modified:
|
||||
self.name_modified = False
|
||||
if self.address:
|
||||
idaapi.set_name(self.address, self.name)
|
||||
if self.tinfo_modified:
|
||||
self.tinfo_modified = False
|
||||
if self.address:
|
||||
idaapi.apply_tinfo2(self.address, self.tinfo.get_pointed_object(), idaapi.TINFO_DEFINITE)
|
||||
|
||||
def __eq__(self, other):
|
||||
return self.address == other.address
|
||||
|
||||
def __repr__(self):
|
||||
return self.name
|
||||
|
||||
|
||||
class VirtualTable(object):
|
||||
def __init__(self, ordinal, tinfo, class_):
|
||||
self.ordinal = ordinal
|
||||
self.tinfo = tinfo
|
||||
self.class_ = [class_]
|
||||
self.class_name = None
|
||||
self.virtual_functions = []
|
||||
self.name = self.tinfo.dstr()
|
||||
self._modified = False
|
||||
|
||||
def update(self):
|
||||
if self.modified:
|
||||
vtable_tinfo = idaapi.tinfo_t()
|
||||
udt_data = idaapi.udt_type_data_t()
|
||||
vtable_tinfo.get_numbered_type(idaapi.cvar.idati, self.ordinal)
|
||||
vtable_tinfo.get_udt_details(udt_data)
|
||||
self.tinfo = vtable_tinfo
|
||||
self.name = vtable_tinfo.dstr()
|
||||
self.modified = False
|
||||
if len(self.virtual_functions) == len(udt_data):
|
||||
for current_function, other_function in zip(self.virtual_functions, udt_data):
|
||||
current_function.update(other_function.name, other_function.type)
|
||||
else:
|
||||
print "[ERROR] Something have been modified in Local types. Please refresh this view"
|
||||
|
||||
def update_local_type(self):
|
||||
if self.modified:
|
||||
final_tinfo = idaapi.tinfo_t()
|
||||
udt_data = idaapi.udt_type_data_t()
|
||||
self.tinfo.get_udt_details(udt_data)
|
||||
if len(udt_data) == len(self.virtual_functions):
|
||||
for udt_member, virtual_function in zip(udt_data, self.virtual_functions):
|
||||
udt_member.name = virtual_function.name
|
||||
udt_member.type = virtual_function.tinfo
|
||||
virtual_function.commit()
|
||||
final_tinfo.create_udt(udt_data, idaapi.BTF_STRUCT)
|
||||
final_tinfo.set_numbered_type(idaapi.cvar.idati, self.ordinal, idaapi.NTF_REPLACE, self.name)
|
||||
self.modified = False
|
||||
else:
|
||||
print "[ERROR] Something have been modified in Local types. Please refresh this view"
|
||||
|
||||
def set_first_argument_type(self, class_name):
|
||||
for function in self.virtual_functions:
|
||||
function.set_first_argument_type(class_name)
|
||||
|
||||
@property
|
||||
def modified(self):
|
||||
return self._modified
|
||||
|
||||
@modified.setter
|
||||
def modified(self, value):
|
||||
self._modified = value
|
||||
if value:
|
||||
for class_ in self.class_:
|
||||
class_.modified = True
|
||||
|
||||
@staticmethod
|
||||
def create(tinfo, class_):
|
||||
ordinal = idaapi.get_type_ordinal(idaapi.cvar.idati, tinfo.dstr())
|
||||
result = all_virtual_tables.get(ordinal)
|
||||
if result:
|
||||
result.class_.append(class_)
|
||||
else:
|
||||
udt_data = idaapi.udt_type_data_t()
|
||||
tinfo.get_udt_details(udt_data)
|
||||
result = VirtualTable(ordinal, tinfo, class_)
|
||||
virtual_functions = [VirtualMethod.create(func.type, func.name, result) for func in udt_data]
|
||||
result.virtual_functions = virtual_functions
|
||||
all_virtual_functions[ordinal] = result
|
||||
return result
|
||||
|
||||
def get_class_tinfo(self):
|
||||
if len(self.class_) == 1:
|
||||
return self.class_.tinfo
|
||||
|
||||
def setData(self, column, value):
|
||||
if column == 0:
|
||||
if idaapi.isident(value) and self.name != value:
|
||||
self.name = value
|
||||
self.modified = True
|
||||
return True
|
||||
return False
|
||||
|
||||
def data(self, column):
|
||||
if column == 0:
|
||||
return self.name
|
||||
|
||||
@property
|
||||
def color(self):
|
||||
return QtGui.QColor("#d5f4e6")
|
||||
|
||||
@property
|
||||
def tooltip(self):
|
||||
pass
|
||||
|
||||
def font(self, column):
|
||||
if self.modified:
|
||||
return QtGui.QFont("Consolas", 12, italic=True)
|
||||
return QtGui.QFont("Consolas", 12)
|
||||
|
||||
def flags(self, column):
|
||||
if column == 0:
|
||||
return QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEditable | QtCore.Qt.ItemIsEnabled
|
||||
return QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEnabled
|
||||
|
||||
@property
|
||||
def children(self):
|
||||
return self.virtual_functions
|
||||
|
||||
def __repr__(self):
|
||||
return str(self.virtual_functions)
|
||||
|
||||
|
||||
class Class(object):
|
||||
def __init__(self, name, tinfo, ordinal):
|
||||
self.name = name
|
||||
self.class_name = name
|
||||
self.ordinal = ordinal
|
||||
self.parent = None
|
||||
self.vtables = {}
|
||||
self.modified = False
|
||||
|
||||
@staticmethod
|
||||
def create_class(ordinal):
|
||||
tinfo = idaapi.tinfo_t()
|
||||
tinfo.get_numbered_type(idaapi.cvar.idati, ordinal)
|
||||
vtables = {}
|
||||
if tinfo.is_struct():
|
||||
udt_data = idaapi.udt_type_data_t()
|
||||
tinfo.get_udt_details(udt_data)
|
||||
for field_udt in udt_data:
|
||||
if field_udt.type.is_ptr():
|
||||
possible_vtable = field_udt.type.get_pointed_object()
|
||||
if possible_vtable.is_struct():
|
||||
v_udt_data = idaapi.udt_type_data_t()
|
||||
possible_vtable.get_udt_details(v_udt_data)
|
||||
for possible_func_udt in v_udt_data:
|
||||
if not possible_func_udt.type.is_funcptr():
|
||||
break
|
||||
else:
|
||||
vtables[field_udt.offset / 8] = possible_vtable
|
||||
if vtables:
|
||||
class_ = Class(tinfo.dstr(), tinfo, ordinal)
|
||||
for offset, vtable_tinfo in vtables.iteritems():
|
||||
vtables[offset] = VirtualTable.create(vtable_tinfo, class_)
|
||||
class_.vtables = vtables
|
||||
return class_
|
||||
|
||||
def update_from_local_types(self):
|
||||
try:
|
||||
if self.modified:
|
||||
class_ = self.create_class(self.ordinal)
|
||||
if class_:
|
||||
self.name = class_.name
|
||||
self.modified = False
|
||||
for offset, vtable in class_.vtables.iteritems():
|
||||
self.vtables[offset].update()
|
||||
else:
|
||||
# TODO: drop class
|
||||
raise IndexError
|
||||
except IndexError:
|
||||
print "[ERROR] Something have been modified in Local types. Please refresh this view"
|
||||
|
||||
def update_local_type(self):
|
||||
if self.modified:
|
||||
for vtable in self.vtables.values():
|
||||
vtable.update_local_type()
|
||||
udt_data = idaapi.udt_type_data_t()
|
||||
tinfo = idaapi.tinfo_t()
|
||||
self.tinfo.get_udt_details(udt_data)
|
||||
tinfo.create_udt(udt_data, idaapi.BTF_STRUCT)
|
||||
tinfo.set_numbered_type(idaapi.cvar.idati, self.ordinal, idaapi.NTF_REPLACE, self.name)
|
||||
self.modified = False
|
||||
|
||||
def set_first_argument_type(self, class_name):
|
||||
if 0 in self.vtables:
|
||||
self.vtables[0].set_first_argument_type(class_name)
|
||||
|
||||
def has_function(self, regexp):
|
||||
for vtable in self.vtables.values():
|
||||
if filter(lambda func: regexp.indexIn(func.name) >= 0, vtable.virtual_functions):
|
||||
return True
|
||||
return False
|
||||
|
||||
def data(self, column):
|
||||
if column == 0:
|
||||
return self.name
|
||||
|
||||
def setData(self, column, value):
|
||||
if column == 0:
|
||||
if idaapi.isident(value) and self.name != value:
|
||||
self.name = value
|
||||
self.modified = True
|
||||
return True
|
||||
return False
|
||||
|
||||
def font(self, column):
|
||||
if self.modified:
|
||||
return QtGui.QFont("Consolas", 12, QtGui.QFont.Bold, italic=True)
|
||||
return QtGui.QFont("Consolas", 12, QtGui.QFont.Bold)
|
||||
|
||||
@property
|
||||
def color(self):
|
||||
# return QtGui.QBrush(QtGui.QColor("#ffb3ff")):
|
||||
return QtGui.QColor("#80ced6")
|
||||
|
||||
@property
|
||||
def tooltip(self):
|
||||
return None
|
||||
|
||||
def flags(self, column):
|
||||
if column == 0:
|
||||
return QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEditable | QtCore.Qt.ItemIsEnabled
|
||||
return QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEnabled
|
||||
|
||||
@property
|
||||
def children(self):
|
||||
return self.vtables.values()
|
||||
|
||||
@property
|
||||
def tinfo(self):
|
||||
tinfo = idaapi.tinfo_t()
|
||||
tinfo.get_numbered_type(idaapi.cvar.idati, self.ordinal)
|
||||
return tinfo
|
||||
|
||||
def open_function(self):
|
||||
pass
|
||||
|
||||
def __repr__(self):
|
||||
return self.name + " ^_^ " + str(self.vtables)
|
||||
|
||||
|
||||
class TreeItem:
|
||||
def __init__(self, item, row, parent):
|
||||
self.item = item
|
||||
self.parent = parent
|
||||
self.row = row
|
||||
self.children = []
|
||||
|
||||
def __repr__(self):
|
||||
return str(self.item)
|
||||
|
||||
|
||||
class TreeModel(QtCore.QAbstractItemModel):
|
||||
# TODO: Add higlighting if eip in function, consider setting breakpoints
|
||||
|
||||
refreshed = QtCore.Signal()
|
||||
|
||||
def __init__(self):
|
||||
super(TreeModel, self).__init__()
|
||||
self.tree_data = []
|
||||
self.headers = ["Name", "Declaration", "Address"]
|
||||
|
||||
self.init()
|
||||
# import pydevd
|
||||
# pydevd.settrace("localhost", port=12345, stdoutToServer=True, stderrToServer=True)
|
||||
|
||||
def init(self):
|
||||
idaapi.show_wait_box("Looking for classes...")
|
||||
all_virtual_functions.clear()
|
||||
all_virtual_tables.clear()
|
||||
|
||||
classes = []
|
||||
for ordinal in xrange(1, idaapi.get_ordinal_qty(idaapi.cvar.idati)):
|
||||
result = Class.create_class(ordinal)
|
||||
if result:
|
||||
classes.append(result)
|
||||
|
||||
for class_row, class_ in enumerate(classes):
|
||||
class_item = TreeItem(class_, class_row, None)
|
||||
for vtable_row, vtable in class_.vtables.iteritems():
|
||||
vtable_item = TreeItem(vtable, vtable_row, class_item)
|
||||
vtable_item.children = [TreeItem(function, 0, vtable_item) for function in vtable.virtual_functions]
|
||||
class_item.children.append(vtable_item)
|
||||
self.tree_data.append(class_item)
|
||||
|
||||
idaapi.hide_wait_box()
|
||||
|
||||
def flags(self, index):
|
||||
if index.isValid():
|
||||
return index.internalPointer().item.flags(index.column())
|
||||
|
||||
def index(self, row, column, parent=QtCore.QModelIndex()):
|
||||
if parent.isValid():
|
||||
node = parent.internalPointer()
|
||||
return self.createIndex(row, column, node.children[row])
|
||||
else:
|
||||
return self.createIndex(row, column, self.tree_data[row])
|
||||
|
||||
def parent(self, index):
|
||||
if index.isValid():
|
||||
node = index.internalPointer()
|
||||
if node.parent:
|
||||
return self.createIndex(0, 0, node.parent)
|
||||
return QtCore.QModelIndex()
|
||||
|
||||
def rowCount(self, index=QtCore.QModelIndex()):
|
||||
if index.isValid():
|
||||
node = index.internalPointer()
|
||||
if node:
|
||||
return len(node.children)
|
||||
return len(self.tree_data)
|
||||
|
||||
def columnCount(self, index=QtCore.QModelIndex()):
|
||||
return 3
|
||||
|
||||
def data(self, index, role=QtCore.Qt.DisplayRole):
|
||||
if role == QtCore.Qt.DisplayRole or role == QtCore.Qt.EditRole:
|
||||
node = index.internalPointer()
|
||||
return node.item.data(index.column())
|
||||
elif role == QtCore.Qt.FontRole:
|
||||
return index.internalPointer().item.font(index.column())
|
||||
elif role == QtCore.Qt.ToolTipRole:
|
||||
return index.internalPointer().item.tooltip
|
||||
elif role == QtCore.Qt.BackgroundRole:
|
||||
return index.internalPointer().item.color
|
||||
elif role == QtCore.Qt.ForegroundRole:
|
||||
return QtGui.QBrush(QtGui.QColor("#191919"))
|
||||
return None
|
||||
|
||||
def setData(self, index, value, role=QtCore.Qt.DisplayRole):
|
||||
result = False
|
||||
if role == QtCore.Qt.EditRole and value != "":
|
||||
node = index.internalPointer()
|
||||
result = node.item.setData(index.column(), str(value))
|
||||
return result
|
||||
|
||||
def headerData(self, section, orientation, role):
|
||||
if role == QtCore.Qt.DisplayRole and orientation == QtCore.Qt.Horizontal:
|
||||
return self.headers[section]
|
||||
|
||||
def set_first_argument_type(self, indexes):
|
||||
indexes = filter(lambda x: x.column() == 0, indexes)
|
||||
class_name = indexes[0].internalPointer().item.class_name
|
||||
if not class_name:
|
||||
classes = [[x.item.name] for x in self.tree_data]
|
||||
class_chooser = HexRaysPyTools.Forms.MyChoose(classes, "Select Class", [["Name", 25]])
|
||||
idx = class_chooser.Show(True)
|
||||
if idx != -1:
|
||||
class_name = classes[idx][0]
|
||||
if class_name:
|
||||
for index in indexes:
|
||||
index.internalPointer().item.set_first_argument_type(class_name)
|
||||
|
||||
def refresh(self):
|
||||
self.tree_data = []
|
||||
self.modelReset.emit()
|
||||
self.init()
|
||||
self.refreshed.emit()
|
||||
|
||||
def rollback(self):
|
||||
for class_item in self.tree_data:
|
||||
class_item.item.update_from_local_types()
|
||||
self.dataChanged.emit(QtCore.QModelIndex(), QtCore.QModelIndex())
|
||||
|
||||
def commit(self):
|
||||
for class_item in self.tree_data:
|
||||
if class_item.item.modified:
|
||||
class_item.item.update_local_type()
|
||||
|
||||
def open_function(self, index):
|
||||
if index.column() == 2:
|
||||
index.internalPointer().item.open_function()
|
||||
|
||||
|
||||
class ProxyModel(QtGui.QSortFilterProxyModel):
|
||||
def __init__(self):
|
||||
super(ProxyModel, self).__init__()
|
||||
self.filter_by_function = False
|
||||
|
||||
def set_regexp_filter(self, regexp):
|
||||
if regexp and regexp[0] == '!':
|
||||
self.filter_by_function = True
|
||||
self.setFilterRegExp(regexp[1:])
|
||||
else:
|
||||
self.filter_by_function = False
|
||||
self.setFilterRegExp(regexp)
|
||||
|
||||
def filterAcceptsRow(self, row, parent):
|
||||
if not parent.isValid() and self.filterRegExp():
|
||||
if self.filter_by_function:
|
||||
return self.sourceModel().tree_data[row].item.has_function(self.filterRegExp())
|
||||
return self.filterRegExp().indexIn(self.sourceModel().tree_data[row].item.class_name) >= 0
|
||||
return True
|
||||
Reference in New Issue
Block a user