explorer: fix multiple bugs

CHANGELOG.md

@@ -112,6 +112,7 @@
 - fix: improve exception handling to prevent IDA from locking up when errors occur #1262 @mike-hunhoff
 - verify rule metadata using Pydantic #1167 @mr-tz
 - extractor: make read consistent with file object behavior #1254 @mr-tz
+- fix: UnboundLocalError x2 #1302 @mike-hunhoff
 
 ### Development
 

capa/ida/plugin/form.py

@@ -526,9 +526,13 @@ def ida_hook_rebase(self, meta, post=False):
             self.model_data.reset()
 
     def ensure_capa_settings_rule_path(self):
+        path: str
+
         try:
+            path = settings.user.get(CAPA_SETTINGS_RULE_PATH, "")
+
             # resolve rules directory - check self and settings first, then ask user
-            if not os.path.exists(settings.user.get(CAPA_SETTINGS_RULE_PATH, "")):
+            if not os.path.exists(path):
                 # configure rules selection messagebox
                 rules_message = QtWidgets.QMessageBox()
                 rules_message.setIcon(QtWidgets.QMessageBox.Information)
@@ -566,6 +570,7 @@ def ensure_capa_settings_rule_path(self):
             logger.info("User cancelled analysis.")
             return False
 
+        path = settings.user.get(CAPA_SETTINGS_RULE_PATH, "")
         if not os.path.exists(path):
             logger.error("rule path %s does not exist or cannot be accessed" % path)
             return False
@@ -580,7 +585,7 @@ def load_capa_rules(self):
         rule_path: str = settings.user.get(CAPA_SETTINGS_RULE_PATH, "")
         try:
 
-            def on_load_rule(rule_path, i, total):
+            def on_load_rule(_, i, total):
                 update_wait_box("loading capa rules from %s (%d of %d)" % (rule_path, i, total))
                 if ida_kernwin.user_cancelled():
                     raise UserCancelledError("user cancelled")
@@ -723,10 +728,12 @@ def slot_progress_feature_extraction(text):
             # either the results are cached and the doc already exists,
             # or the doc was just created above
             assert self.resdoc_cache is not None
+            assert self.ruleset_cache is not None
 
             self.model_data.render_capa_doc(self.resdoc_cache, self.view_show_results_by_function.isChecked())
             self.set_view_status_label(
-                "capa rules: %s (%d rules)" % (settings.user[CAPA_SETTINGS_RULE_PATH], ruleset.source_rule_count)
+                "capa rules: %s (%d rules)"
+                % (settings.user[CAPA_SETTINGS_RULE_PATH], self.ruleset_cache.source_rule_count)
             )
         except Exception as e:
             logger.error("Failed to render results (error: %s)", e, exc_info=True)

capa/ida/plugin/model.py

@@ -6,7 +6,7 @@
 #  is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and limitations under the License.
 
-from typing import Set, Dict, List, Tuple
+from typing import Set, Dict, List, Tuple, Optional
 from collections import deque
 
 import idc
@@ -444,34 +444,40 @@ def render_capa_doc_match(self, parent: CapaExplorerDataItem, match: rd.Match, d
             self.render_capa_doc_match(parent2, child, doc)
 
     def render_capa_doc_by_function(self, doc: rd.ResultDocument):
-        """ """
-        matches_by_function: Dict[int, Tuple[CapaExplorerFunctionItem, Set[str]]] = {}
+        """render rule matches by function meaning each rule match is nested under function where it was found"""
+        matches_by_function: Dict[AbsoluteVirtualAddress, Tuple[CapaExplorerFunctionItem, Set[str]]] = {}
         for rule in rutils.capability_rules(doc):
             for location_, _ in rule.matches:
                 location = location_.to_capa()
 
                 if not isinstance(location, AbsoluteVirtualAddress):
                     # only handle matches with a VA
                     continue
-                ea = int(location)
 
-                ea = capa.ida.helpers.get_func_start_ea(ea)
-                if ea is None:
-                    # file scope, skip rendering in this mode
+                func_ea: Optional[int] = capa.ida.helpers.get_func_start_ea(int(location))
+                if func_ea is None:
+                    # rule match address is not located in a defined function
                     continue
-                if not matches_by_function.get(ea, ()):
-                    # new function root
-                    matches_by_function[ea] = (
-                        CapaExplorerFunctionItem(self.root_node, location, can_check=False),
+
+                func_address: AbsoluteVirtualAddress = AbsoluteVirtualAddress(func_ea)
+                if not matches_by_function.get(func_address, ()):
+                    # create a new function root to nest its rule matches; Note: we must use the address of the
+                    # function here so everything is displayed properly
+                    matches_by_function[func_address] = (
+                        CapaExplorerFunctionItem(self.root_node, func_address, can_check=False),
                         set(),
                     )
-                function_root, match_cache = matches_by_function[ea]
-                if rule.meta.name in match_cache:
-                    # rule match already rendered for this function root, skip it
+
+                func_root, func_match_cache = matches_by_function[func_address]
+                if rule.meta.name in func_match_cache:
+                    # only nest each rule once, so if found, skip
                     continue
-                match_cache.add(rule.meta.name)
+
+                # add matched rule to its function cache; create a new rule node whose parent is the matched
+                # function node
+                func_match_cache.add(rule.meta.name)
                 CapaExplorerRuleItem(
-                    function_root,
+                    func_root,
                     rule.meta.name,
                     rule.meta.namespace or "",
                     len(rule.matches),

capa/main.py

@@ -637,7 +637,7 @@ def get_rules(
 
     total_rule_count = len(rule_file_paths)
     for i, (path, content) in enumerate(zip(rule_file_paths, rule_contents)):
-        on_load_rule(path, i, total_rule_count)
+        on_load_rule(path, i + 1, total_rule_count)
 
         try:
             rule = capa.rules.Rule.from_yaml(content.decode("utf-8"))