feat: Add Filter, Collapse and Expand to "Classes" window

Author: igogo-x86

HexRaysPyTools/Core/Classes.py

@@ -327,6 +327,12 @@ 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
@@ -504,3 +510,24 @@ def commit(self):
     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

HexRaysPyTools/Cute.py

@@ -43,6 +43,7 @@
         QtGui.QTableView = QtWidgets.QTableView
         QtGui.QAction = QtWidgets.QAction
         QtGui.QMenu = QtWidgets.QMenu
+        QtGui.QLabel = QtWidgets.QLabel
 
         QtGui.QHeaderView = QtWidgets.QHeaderView
         QtGui.QAbstractItemView = QtWidgets.QAbstractItemView
@@ -52,6 +53,8 @@
         QtGui.QPushButton = QtWidgets.QPushButton
         QtGui.QSpacerItem = QtWidgets.QSpacerItem
         QtGui.QSizePolicy = QtWidgets.QSizePolicy
+        QtGui.QSortFilterProxyModel = QtCore.QSortFilterProxyModel
+        QtGui.QLineEdit = QtWidgets.QLineEdit
 
         use_qt5 = True
     else:

HexRaysPyTools/Forms.py

@@ -140,7 +140,10 @@ 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)
@@ -163,14 +166,25 @@ def init_ui(self):
             "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.class_tree.setModel(class_model)
+        proxy_model = Core.Classes.ProxyModel()
+        proxy_model.setSourceModel(class_model)
+        proxy_model.setFilterCaseSensitivity(QtCore.Qt.CaseInsensitive)
+        self.class_tree.setModel(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(self.class_tree.selectedIndexes())
         )
@@ -179,17 +193,22 @@ def init_ui(self):
         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(class_model.open_function)
         self.class_tree.customContextMenuRequested[QtCore.QPoint].connect(self.show_menu)
+        self.line_edit_filter.textChanged[str].connect(proxy_model.set_regexp_filter)
+        # proxy_model.rowsInserted[object].connect(lambda: self.class_tree.setExpanded(object, True))
 
     def OnClose(self, form):
         pass

Img/classes.JPG

@@ -150,11 +150,14 @@ Also can be found at _View->Open Subview->Classes_. Helps to manage classes (str
 ##### !! Better to rename all functions before debugging because Ida can mess up with default names and information in virtual tables will be inconsistent
 Class, virtual table and functions names are editable. Also function's declaration can be edited. After edit, altered items change font to _italic_. Right click opens following menu options:
 
+* Expand All / Collapse All
 * Refresh - clear all and rescan local types for information again
 * Rollback - undo changes
-* Commit - apply changes. Functions will be renamed and recasted both in virtual tables in Local Types and dissasmbly code.
+* Commit - apply changes. Functions will be renamed and recasted both in virtual tables in Local Types and dissasembly code.
 * Set First Argument type - allows to select first argument for function among all classes. If right click was used on class name, than it's type will be automatically applied to virtual table at offset 0
 
+You can also filter classes using Regexp either by class_name or by existence of specific functions. Just input expression in line edit for filtering by class_name or prepend it with "!" to filter by function name.
+
 
 
 ### 7) Function signature manipulation