diff --git a/capa/ida/plugin/form.py b/capa/ida/plugin/form.py index 655d60fa..82e3771c 100644 --- a/capa/ida/plugin/form.py +++ b/capa/ida/plugin/form.py @@ -24,7 +24,7 @@ from capa.ida.plugin.icon import QICON 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 @@ class CapaExplorerForm(idaapi.PluginForm): # 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 @@ class CapaExplorerForm(idaapi.PluginForm): """ 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 @@ class CapaExplorerForm(idaapi.PluginForm): # 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 @@ class CapaExplorerForm(idaapi.PluginForm): 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 @@ class CapaExplorerForm(idaapi.PluginForm): """ 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 @@ class CapaExplorerForm(idaapi.PluginForm): 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 @@ class CapaExplorerForm(idaapi.PluginForm): 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 @@ class CapaExplorerForm(idaapi.PluginForm): @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 """ diff --git a/capa/ida/plugin/proxy.py b/capa/ida/plugin/proxy.py index fceb1d41..349266c6 100644 --- a/capa/ida/plugin/proxy.py +++ b/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 @@ class CapaExplorerSortFilterProxyModel(QtCore.QSortFilterProxyModel): 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("") diff --git a/capa/ida/plugin/view.py b/capa/ida/plugin/view.py index c4dd86a6..49c3956d 100644 --- a/capa/ida/plugin/view.py +++ b/capa/ida/plugin/view.py @@ -72,7 +72,24 @@ class CapaExplorerQtreeView(QtWidgets.QTreeView): @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 @@ class CapaExplorerQtreeView(QtWidgets.QTreeView): @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