adding code to include basic block scope when limiting results by a function

capa/ida/explorer/proxy.py

@@ -16,6 +16,9 @@ class CapaExplorerSortFilterProxyModel(QtCore.QSortFilterProxyModel):
         """ """
         super(CapaExplorerSortFilterProxyModel, self).__init__(parent)
 
+        self.min_ea = None
+        self.max_ea = None
+
     def lessThan(self, left, right):
         """true if the value of the left item is less than value of right item
 
@@ -62,15 +65,6 @@ class CapaExplorerSortFilterProxyModel(QtCore.QSortFilterProxyModel):
 
         return False
 
-    def add_single_string_filter(self, column, string):
-        """add fixed string filter
-
-        @param column: key column
-        @param string: string to sort
-        """
-        self.setFilterKeyColumn(column)
-        self.setFilterFixedString(string)
-
     def index_has_accepted_children(self, row, parent):
         """ """
         model_index = self.sourceModel().index(row, 0, parent)
@@ -86,4 +80,33 @@ class CapaExplorerSortFilterProxyModel(QtCore.QSortFilterProxyModel):
 
     def filter_accepts_row_self(self, row, parent):
         """ """
-        return super(CapaExplorerSortFilterProxyModel, self).filterAcceptsRow(row, parent)
+        # filter not set
+        if self.min_ea is None and self.max_ea is None:
+            return True
+
+        index = self.sourceModel().index(row, 0, parent)
+        data = index.internalPointer().data(CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS)
+
+        if not data:
+            return False
+
+        ea = int(data, 16)
+
+        if self.min_ea <= ea and ea < self.max_ea:
+            return True
+
+        return False
+
+    def add_address_range_filter(self, min_ea, max_ea):
+        """ """
+        self.min_ea = min_ea
+        self.max_ea = max_ea
+
+        self.setFilterKeyColumn(CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS)
+        self.invalidateFilter()
+
+    def reset_address_range_filter(self):
+        """ """
+        self.min_ea = None
+        self.max_ea = None
+        self.invalidateFilter()

capa/ida/ida_capa_explorer.py

@@ -324,35 +324,25 @@ class CapaExplorerForm(idaapi.PluginForm):
     def ida_hook_screen_ea_changed(self, widget, new_ea, old_ea):
         """hook for IDA screen ea changed
 
+        this hook is currently only relevant for limiting results displayed in the UI
+
         @param widget: IDA widget type
         @param new_ea: destination ea
         @param old_ea: source ea
         """
         if not self.view_limit_results_by_function.isChecked():
-            # ignore if checkbox not selected
+            # ignore if limit checkbox not selected
             return
 
         if idaapi.get_widget_type(widget) != idaapi.BWN_DISASM:
-            # ignore views other than asm
+            # ignore views not the assembly view
             return
 
-        # attempt to map virtual addresses to function start addresses
-        new_func_start = capa.ida.helpers.get_func_start_ea(new_ea)
-        old_func_start = capa.ida.helpers.get_func_start_ea(old_ea)
-
-        if new_func_start and new_func_start == old_func_start:
-            # navigated within the same function - do nothing
+        if idaapi.get_func(new_ea) == idaapi.get_func(old_ea):
+            # user navigated same function - ignore
             return
 
-        if new_func_start:
-            # navigated to new function - filter for function start virtual address
-            match = capa.ida.explorer.item.location_to_hex(new_func_start)
-        else:
-            # navigated to virtual address not in valid function - clear filter
-            match = ""
-
-        # filter on virtual address to avoid updating filter string if function name is changed
-        self.model_proxy.add_single_string_filter(CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS, match)
+        self.limit_results_to_function(idaapi.get_func(new_ea))
         self.view_tree.resize_columns_to_content()
 
     def load_capa_results(self):
@@ -534,16 +524,24 @@ class CapaExplorerForm(idaapi.PluginForm):
         if checked, configure function filter if screen location is located
         in function, otherwise clear filter
         """
-        match = ""
         if self.view_limit_results_by_function.isChecked():
-            location = capa.ida.helpers.get_func_start_ea(idaapi.get_screen_ea())
-            if location:
-                match = capa.ida.explorer.item.location_to_hex(location)
-
-        self.model_proxy.add_single_string_filter(CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS, match)
+            self.limit_results_to_function(idaapi.get_func(idaapi.get_screen_ea()))
+        else:
+            self.model_proxy.reset_address_range_filter()
 
         self.view_tree.resize_columns_to_content()
 
+    def limit_results_to_function(self, f):
+        """add filter to limit results to current function
+
+        @param f: (IDA func_t)
+        """
+        if f:
+            self.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)
+
 
 def main():
     """ TODO: move to idaapi.plugin_t class """