247 lines
10 KiB
Python
247 lines
10 KiB
Python
import idaapi
|
|
# import PySide.QtGui as QtGui
|
|
# import PySide.QtCore as QtCore
|
|
from HexRaysPyTools.Cute import *
|
|
import Core.Classes
|
|
|
|
|
|
class MyChoose(idaapi.Choose2):
|
|
def __init__(self, items, title, cols, icon=-1):
|
|
idaapi.Choose2.__init__(self, title, cols, flags=idaapi.Choose2.CH_MODAL, icon=icon)
|
|
self.items = items
|
|
|
|
def OnClose(self):
|
|
pass
|
|
|
|
def OnGetLine(self, n):
|
|
return self.items[n]
|
|
|
|
def OnGetSize(self):
|
|
return len(self.items)
|
|
|
|
|
|
class StructureBuilder(idaapi.PluginForm):
|
|
def __init__(self, structure_model):
|
|
super(StructureBuilder, self).__init__()
|
|
self.structure_model = structure_model
|
|
self.parent = None
|
|
|
|
def OnCreate(self, form):
|
|
self.parent = form_to_widget(form)
|
|
self.init_ui()
|
|
|
|
def init_ui(self):
|
|
self.parent.setStyleSheet(
|
|
"QTableView {background-color: transparent; selection-background-color: #87bdd8;}"
|
|
"QHeaderView::section {background-color: transparent; border: 0.5px solid;}"
|
|
"QPushButton {width: 50px; height: 20px;}"
|
|
# "QPushButton::pressed {background-color: #ccccff}"
|
|
)
|
|
self.parent.resize(400, 600)
|
|
self.parent.setWindowTitle('Structure Builder')
|
|
|
|
btn_finalize = QtGui.QPushButton("&Finalize")
|
|
btn_disable = QtGui.QPushButton("&Disable")
|
|
btn_enable = QtGui.QPushButton("&Enable")
|
|
btn_origin = QtGui.QPushButton("&Origin")
|
|
btn_array = QtGui.QPushButton("&Array")
|
|
btn_pack = QtGui.QPushButton("&Pack")
|
|
btn_unpack = QtGui.QPushButton("&Unpack")
|
|
btn_remove = QtGui.QPushButton("&Remove")
|
|
btn_resolve = QtGui.QPushButton("Resolve")
|
|
btn_clear = QtGui.QPushButton("Clear") # Clear button doesn't have shortcut because it can fuck up all work
|
|
btn_recognize = QtGui.QPushButton("Recognize Shape")
|
|
btn_recognize.setStyleSheet("QPushButton {width: 100px; height: 20px;}")
|
|
|
|
btn_finalize.setShortcut("f")
|
|
btn_disable.setShortcut("d")
|
|
btn_enable.setShortcut("e")
|
|
btn_origin.setShortcut("o")
|
|
btn_array.setShortcut("a")
|
|
btn_pack.setShortcut("p")
|
|
btn_unpack.setShortcut("u")
|
|
btn_remove.setShortcut("r")
|
|
|
|
struct_view = QtGui.QTableView()
|
|
struct_view.setModel(self.structure_model)
|
|
# struct_view.setSelectionMode(QtGui.QAbstractItemView.ExtendedSelection)
|
|
|
|
struct_view.verticalHeader().setVisible(False)
|
|
struct_view.verticalHeader().setDefaultSectionSize(24)
|
|
struct_view.horizontalHeader().setStretchLastSection(True)
|
|
struct_view.horizontalHeader().setSectionResizeMode(QtGui.QHeaderView.ResizeToContents)
|
|
|
|
grid_box = QtGui.QGridLayout()
|
|
grid_box.setSpacing(0)
|
|
grid_box.addWidget(btn_finalize, 0, 0)
|
|
grid_box.addWidget(btn_enable, 0, 1)
|
|
grid_box.addWidget(btn_disable, 0, 2)
|
|
grid_box.addWidget(btn_origin, 0, 3)
|
|
grid_box.addItem(QtGui.QSpacerItem(20, 20, QtGui.QSizePolicy.Expanding), 0, 5)
|
|
grid_box.addWidget(btn_array, 1, 0)
|
|
grid_box.addWidget(btn_pack, 1, 1)
|
|
grid_box.addWidget(btn_unpack, 1, 2)
|
|
grid_box.addWidget(btn_remove, 1, 3)
|
|
grid_box.addWidget(btn_resolve, 0, 4)
|
|
grid_box.addItem(QtGui.QSpacerItem(20, 20, QtGui.QSizePolicy.Expanding), 1, 5)
|
|
grid_box.addWidget(btn_recognize, 0, 6)
|
|
grid_box.addWidget(btn_clear, 1, 6)
|
|
|
|
vertical_box = QtGui.QVBoxLayout()
|
|
vertical_box.addWidget(struct_view)
|
|
vertical_box.addLayout(grid_box)
|
|
self.parent.setLayout(vertical_box)
|
|
|
|
btn_finalize.clicked.connect(lambda: self.structure_model.finalize())
|
|
btn_disable.clicked.connect(lambda: self.structure_model.disable_rows(struct_view.selectedIndexes()))
|
|
btn_enable.clicked.connect(lambda: self.structure_model.enable_rows(struct_view.selectedIndexes()))
|
|
btn_origin.clicked.connect(lambda: self.structure_model.set_origin(struct_view.selectedIndexes()))
|
|
btn_array.clicked.connect(lambda: self.structure_model.make_array(struct_view.selectedIndexes()))
|
|
btn_pack.clicked.connect(lambda: self.structure_model.pack_substructure(struct_view.selectedIndexes()))
|
|
btn_unpack.clicked.connect(lambda: self.structure_model.unpack_substructure(struct_view.selectedIndexes()))
|
|
btn_remove.clicked.connect(lambda: self.structure_model.remove_items(struct_view.selectedIndexes()))
|
|
btn_resolve.clicked.connect(lambda: self.structure_model.resolve_types())
|
|
btn_clear.clicked.connect(lambda: self.structure_model.clear())
|
|
btn_recognize.clicked.connect(lambda: self.structure_model.recognize_shape(struct_view.selectedIndexes()))
|
|
struct_view.activated[QtCore.QModelIndex].connect(self.structure_model.activated)
|
|
self.structure_model.dataChanged.connect(struct_view.clearSelection)
|
|
|
|
def OnClose(self, form):
|
|
pass
|
|
|
|
def Show(self, caption=None, options=0):
|
|
return idaapi.PluginForm.Show(self, caption, options=options)
|
|
|
|
|
|
class StructureGraphViewer(idaapi.GraphViewer):
|
|
def __init__(self, title, graph):
|
|
idaapi.GraphViewer.__init__(self, title)
|
|
self.graph = graph
|
|
self.nodes_id = {}
|
|
|
|
def OnRefresh(self):
|
|
self.Clear()
|
|
self.nodes_id.clear()
|
|
for node in self.graph.get_nodes():
|
|
self.nodes_id[node] = self.AddNode(node)
|
|
for first, second in self.graph.get_edges():
|
|
self.AddEdge(self.nodes_id[first], self.nodes_id[second])
|
|
return True
|
|
|
|
def OnGetText(self, node_id):
|
|
return self.graph.local_types[self[node_id]].name_and_color
|
|
|
|
def OnHint(self, node_id):
|
|
""" Try-catch clause because IDA sometimes attempts to use old information to get hint """
|
|
try:
|
|
ordinal = self[node_id]
|
|
return self.graph.local_types[ordinal].hint
|
|
except KeyError:
|
|
return
|
|
|
|
def OnDblClick(self, node_id):
|
|
self.change_selected([self[node_id]])
|
|
|
|
def change_selected(self, ordinals):
|
|
self.graph.change_selected(ordinals)
|
|
self.Refresh()
|
|
self.Select(self.nodes_id[ordinals[0]])
|
|
|
|
|
|
class ClassViewer(idaapi.PluginForm):
|
|
def __init__(self):
|
|
super(ClassViewer, self).__init__()
|
|
self.parent = None
|
|
self.class_tree = QtGui.QTreeView()
|
|
self.line_edit_filter = QtGui.QLineEdit()
|
|
|
|
self.action_collapse = QtGui.QAction("Collapse all", self.class_tree)
|
|
self.action_expand = QtGui.QAction("Expand all", self.class_tree)
|
|
self.action_set_arg = QtGui.QAction("Set First Argument Type", self.class_tree)
|
|
self.action_rollback = QtGui.QAction("Rollback", self.class_tree)
|
|
self.action_refresh = QtGui.QAction("Refresh", self.class_tree)
|
|
self.action_commit = QtGui.QAction("Commit", self.class_tree)
|
|
|
|
self.menu = QtGui.QMenu(self.parent)
|
|
|
|
self.proxy_model = Core.Classes.ProxyModel()
|
|
|
|
def OnCreate(self, form):
|
|
# self.parent = self.FormToPySideWidget(form)
|
|
self.parent = form_to_widget(form)
|
|
self.init_ui()
|
|
|
|
def init_ui(self):
|
|
self.parent.setWindowTitle('Classes')
|
|
self.parent.setStyleSheet(
|
|
# "QTreeView::item:!has-children { background-color: #fefbd8; border: 0.5px solid lightgray ;}"
|
|
# "QTreeView::item:has-children { background-color: #80ced6; border-top: 1px solid black ;}"
|
|
# "QTreeView::item:selected { background-color: #618685; show-decoration-selected: 1;}"
|
|
"QTreeView {background-color: transparent; }"
|
|
"QHeaderView::section {background-color: transparent; border: 1px solid;}"
|
|
)
|
|
|
|
hbox_layout = QtGui.QHBoxLayout()
|
|
label_filter = QtGui.QLabel("&Filter:")
|
|
label_filter.setBuddy(self.line_edit_filter)
|
|
hbox_layout.addWidget(label_filter)
|
|
hbox_layout.addWidget(self.line_edit_filter)
|
|
|
|
class_model = Core.Classes.TreeModel()
|
|
self.proxy_model.setSourceModel(class_model)
|
|
self.proxy_model.setFilterCaseSensitivity(QtCore.Qt.CaseInsensitive)
|
|
self.class_tree.setModel(self.proxy_model)
|
|
self.class_tree.setContextMenuPolicy(QtCore.Qt.CustomContextMenu)
|
|
self.class_tree.expandAll()
|
|
self.class_tree.header().setStretchLastSection(True)
|
|
self.class_tree.header().setSectionResizeMode(QtGui.QHeaderView.ResizeToContents)
|
|
self.class_tree.setSelectionMode(QtGui.QAbstractItemView.ExtendedSelection)
|
|
|
|
self.action_collapse.triggered.connect(self.class_tree.collapseAll)
|
|
self.action_expand.triggered.connect(self.class_tree.expandAll)
|
|
self.action_set_arg.triggered.connect(
|
|
lambda: class_model.set_first_argument_type(
|
|
map(self.proxy_model.mapToSource, self.class_tree.selectedIndexes())
|
|
)
|
|
)
|
|
self.action_rollback.triggered.connect(lambda: class_model.rollback())
|
|
self.action_refresh.triggered.connect(lambda: class_model.refresh())
|
|
self.action_commit.triggered.connect(lambda: class_model.commit())
|
|
class_model.refreshed.connect(self.class_tree.expandAll)
|
|
|
|
self.menu.addAction(self.action_collapse)
|
|
self.menu.addAction(self.action_expand)
|
|
self.menu.addAction(self.action_refresh)
|
|
self.menu.addAction(self.action_set_arg)
|
|
self.menu.addAction(self.action_rollback)
|
|
self.menu.addAction(self.action_commit)
|
|
|
|
vertical_box = QtGui.QVBoxLayout()
|
|
vertical_box.addWidget(self.class_tree)
|
|
vertical_box.addLayout(hbox_layout)
|
|
self.parent.setLayout(vertical_box)
|
|
|
|
self.class_tree.activated[QtCore.QModelIndex].connect(
|
|
lambda x: class_model.open_function(self.proxy_model.mapToSource(x))
|
|
)
|
|
self.class_tree.customContextMenuRequested[QtCore.QPoint].connect(self.show_menu)
|
|
self.line_edit_filter.textChanged[str].connect(self.proxy_model.set_regexp_filter)
|
|
# proxy_model.rowsInserted[object].connect(lambda: self.class_tree.setExpanded(object, True))
|
|
|
|
def OnClose(self, form):
|
|
pass
|
|
|
|
def Show(self, caption=None, options=0):
|
|
return idaapi.PluginForm.Show(self, caption, options=options)
|
|
|
|
def show_menu(self, point):
|
|
self.action_set_arg.setEnabled(True)
|
|
indexes = map(
|
|
self.proxy_model.mapToSource,
|
|
filter(lambda x: x.column() == 0, self.class_tree.selectedIndexes())
|
|
)
|
|
if len(indexes) > 1:
|
|
if filter(lambda x: len(x.internalPointer().children) > 0, indexes):
|
|
self.action_set_arg.setEnabled(False)
|
|
self.menu.exec_(self.class_tree.mapToGlobal(point))
|