Merge pull request mandiant#285 from fireeye/fix-212-2

ida plugin: add search bar

capa/ida/plugin/form.py

@@ -24,7 +24,7 @@
 from capa.ida.plugin.view import CapaExplorerQtreeView
 from capa.ida.plugin.hooks import CapaExplorerIdaHooks
 from capa.ida.plugin.model import CapaExplorerDataModel
-from capa.ida.plugin.proxy import CapaExplorerSortFilterProxyModel
+from capa.ida.plugin.proxy import CapaExplorerRangeProxyModel, CapaExplorerSearchProxyModel
 
 logger = logging.getLogger(__name__)
 settings = ida_settings.IDASettings("capa")
@@ -44,10 +44,12 @@ def __init__(self, name):
 
         # models
         self.model_data = None
-        self.model_proxy = None
+        self.range_model_proxy = None
+        self.search_model_proxy = None
 
         # user interface elements
         self.view_limit_results_by_function = None
+        self.view_search_bar = None
         self.view_tree = None
         self.view_summary = None
         self.view_attack = None
@@ -83,11 +85,17 @@ def load_interface(self):
         """ load user interface """
         # load models
         self.model_data = CapaExplorerDataModel()
-        self.model_proxy = CapaExplorerSortFilterProxyModel()
-        self.model_proxy.setSourceModel(self.model_data)
+
+        # model <- filter range <- filter search <- view
+
+        self.range_model_proxy = CapaExplorerRangeProxyModel()
+        self.range_model_proxy.setSourceModel(self.model_data)
+
+        self.search_model_proxy = CapaExplorerSearchProxyModel()
+        self.search_model_proxy.setSourceModel(self.range_model_proxy)
 
         # load tree
-        self.view_tree = CapaExplorerQtreeView(self.model_proxy, self.parent)
+        self.view_tree = CapaExplorerQtreeView(self.search_model_proxy, self.parent)
 
         # load summary table
         self.load_view_summary()
@@ -96,6 +104,7 @@ def load_interface(self):
         # load parent tab and children tab views
         self.load_view_tabs()
         self.load_view_checkbox_limit_by()
+        self.load_view_search_bar()
         self.load_view_summary_tab()
         self.load_view_attack_tab()
         self.load_view_tree_tab()
@@ -171,6 +180,14 @@ def load_view_checkbox_limit_by(self):
 
         self.view_limit_results_by_function = check
 
+    def load_view_search_bar(self):
+        """ load the search bar control """
+        line = QtWidgets.QLineEdit()
+        line.setPlaceholderText("search...")
+        line.textChanged.connect(self.search_model_proxy.set_query)
+
+        self.view_search_bar = line
+
     def load_view_parent(self):
         """ load view parent """
         layout = QtWidgets.QVBoxLayout()
@@ -184,6 +201,7 @@ def load_view_tree_tab(self):
         """ load capa tree tab view """
         layout = QtWidgets.QVBoxLayout()
         layout.addWidget(self.view_limit_results_by_function)
+        layout.addWidget(self.view_search_bar)
         layout.addWidget(self.view_tree)
 
         tab = QtWidgets.QWidget()
@@ -484,12 +502,14 @@ def ida_reset(self):
         self.model_data.reset()
         self.view_tree.reset()
         self.view_limit_results_by_function.setChecked(False)
+        self.view_search_bar.setText("")
         self.set_view_tree_default_sort_order()
 
     def reload(self):
         """ reload views and re-run capa analysis """
         self.ida_reset()
-        self.model_proxy.invalidate()
+        self.range_model_proxy.invalidate()
+        self.search_model_proxy.invalidate()
         self.model_data.clear()
         self.view_summary.setRowCount(0)
         self.load_capa_results()
@@ -527,7 +547,7 @@ def slot_checkbox_limit_by_changed(self, state):
         if state == QtCore.Qt.Checked:
             self.limit_results_to_function(idaapi.get_func(idaapi.get_screen_ea()))
         else:
-            self.model_proxy.reset_address_range_filter()
+            self.range_model_proxy.reset_address_range_filter()
 
         self.view_tree.reset()
 
@@ -537,10 +557,10 @@ def limit_results_to_function(self, f):
         @param f: (IDA func_t)
         """
         if f:
-            self.model_proxy.add_address_range_filter(f.start_ea, f.end_ea)
+            self.range_model_proxy.add_address_range_filter(f.start_ea, f.end_ea)
         else:
             # if function not exists don't display any results (address should not be -1)
-            self.model_proxy.add_address_range_filter(-1, -1)
+            self.range_model_proxy.add_address_range_filter(-1, -1)
 
     def ask_user_directory(self):
         """ create Qt dialog to ask user for a directory """

capa/ida/plugin/proxy.py

@@ -7,14 +7,15 @@
 # See the License for the specific language governing permissions and limitations under the License.
 
 from PyQt5 import QtCore
+from PyQt5.QtCore import Qt
 
 from capa.ida.plugin.model import CapaExplorerDataModel
 
 
-class CapaExplorerSortFilterProxyModel(QtCore.QSortFilterProxyModel):
+class CapaExplorerRangeProxyModel(QtCore.QSortFilterProxyModel):
     def __init__(self, parent=None):
         """ """
-        super(CapaExplorerSortFilterProxyModel, self).__init__(parent)
+        super(CapaExplorerRangeProxyModel, self).__init__(parent)
 
         self.min_ea = None
         self.max_ea = None
@@ -110,3 +111,85 @@ def reset_address_range_filter(self):
         self.min_ea = None
         self.max_ea = None
         self.invalidateFilter()
+
+
+class CapaExplorerSearchProxyModel(QtCore.QSortFilterProxyModel):
+    """A SortFilterProxyModel that accepts rows with a substring match for a configurable query.
+
+    Looks for matches in the RULE_INFORMATION column (e.g. column 0).
+    Displays the entire tree row if any of the tree branches,
+     that is, you can filter by rule name, or also
+     filter by "characteristic(nzxor)" to filter matches with some feature.
+    """
+
+    def __init__(self, parent=None):
+        """ """
+        super(CapaExplorerSearchProxyModel, self).__init__(parent)
+        self.query = ""
+        self.setFilterKeyColumn(-1)  # all columns
+
+    def filterAcceptsRow(self, row, parent):
+        """true if the item in the row indicated by the given row and parent
+        should be included in the model; otherwise returns false
+
+        @param row: int
+        @param parent: QModelIndex*
+
+        @retval True/False
+        """
+        # this row matches, accept it
+        if self.filter_accepts_row_self(row, parent):
+            return True
+
+        # the parent of this row matches, accept it
+        alpha = parent
+        while alpha.isValid():
+            if self.filter_accepts_row_self(alpha.row(), alpha.parent()):
+                return True
+            alpha = alpha.parent()
+
+        # this row is a parent, and a child matches, accept it
+        if self.index_has_accepted_children(row, parent):
+            return True
+
+        return False
+
+    def index_has_accepted_children(self, row, parent):
+        """returns True if the given row or its children should be accepted"""
+        source_model = self.sourceModel()
+        model_index = source_model.index(row, 0, parent)
+
+        if model_index.isValid():
+            for idx in range(source_model.rowCount(model_index)):
+                if self.filter_accepts_row_self(idx, model_index):
+                    return True
+                if self.index_has_accepted_children(idx, model_index):
+                    return True
+
+        return False
+
+    def filter_accepts_row_self(self, row, parent):
+        """returns True if the given row should be accepted"""
+        if self.query == "":
+            return True
+
+        source_model = self.sourceModel()
+
+        index = source_model.index(row, 0, parent)
+        data = source_model.data(index, Qt.DisplayRole)
+
+        if not data:
+            return False
+
+        if not isinstance(data, str):
+            # sanity check: should already be a string, but double check
+            return False
+
+        return self.query in data
+
+    def set_query(self, query):
+        self.query = query
+        self.invalidateFilter()
+
+    def reset_query(self):
+        self.set_query("")

capa/ida/plugin/view.py

@@ -72,7 +72,24 @@ def map_index_to_source_item(self, model_index):
 
         @retval QObject*
         """
-        return self.model.mapToSource(model_index).internalPointer()
+        # assume that self.model here is either:
+        #  - CapaExplorerDataModel, or
+        #  - QSortFilterProxyModel subclass
+        #
+        # The ProxyModels may be chained,
+        #  so keep resolving the index the CapaExplorerDataModel.
+
+        model = self.model
+        while not isinstance(model, CapaExplorerDataModel):
+            if not model_index.isValid():
+                raise ValueError("invalid index")
+
+            model_index = model.mapToSource(model_index)
+            model = model.sourceModel()
+
+        if not model_index.isValid():
+            raise ValueError("invalid index")
+        return model_index.internalPointer()
 
     def send_data_to_clipboard(self, data):
         """copy data to the clipboard
@@ -223,11 +240,8 @@ def slot_custom_context_menu_requested(self, pos):
         @param pos: TODO
         """
         model_index = self.indexAt(pos)
-
-        if not model_index.isValid():
-            return
-
         item = self.map_index_to_source_item(model_index)
+
         column = model_index.column()
         menu = None