From a4f0c1c04c06eadd324d74c78a63d92c3e0d2341 Mon Sep 17 00:00:00 2001 From: Moritz Raabe Date: Thu, 19 May 2022 20:43:02 +0200 Subject: [PATCH] fix: rule generator handles --- capa/ida/plugin/form.py | 45 +++++++++++++++++++++-------------------- capa/ida/plugin/view.py | 21 ++++++++++--------- 2 files changed, 35 insertions(+), 31 deletions(-) diff --git a/capa/ida/plugin/form.py b/capa/ida/plugin/form.py index 7c1f4b9d..305f41b4 100644 --- a/capa/ida/plugin/form.py +++ b/capa/ida/plugin/form.py @@ -27,6 +27,7 @@ import capa.render.json import capa.features.common import capa.render.result_document import capa.features.extractors.ida.extractor +from capa.engine import FeatureSet from capa.features.common import Feature from capa.ida.plugin.icon import QICON from capa.ida.plugin.view import ( @@ -35,7 +36,7 @@ from capa.ida.plugin.view import ( CapaExplorerRulgenPreview, CapaExplorerRulegenFeatures, ) -from capa.features.address import Address +from capa.features.address import NO_ADDRESS, Address from capa.ida.plugin.hooks import CapaExplorerIdaHooks from capa.ida.plugin.model import CapaExplorerDataModel from capa.ida.plugin.proxy import CapaExplorerRangeProxyModel, CapaExplorerSearchProxyModel @@ -95,7 +96,7 @@ def find_func_features(fh: FunctionHandle, extractor): return func_features, bb_features -def find_func_matches(f, ruleset, func_features, bb_features): +def find_func_matches(f: FunctionHandle, ruleset, func_features, bb_features): """ """ func_matches = collections.defaultdict(list) bb_matches = collections.defaultdict(list) @@ -112,7 +113,7 @@ def find_func_matches(f, ruleset, func_features, bb_features): func_features[capa.features.common.MatchedRule(name)].add(ea) # find rule matches for function, function features include rule matches for basic blocks - _, matches = capa.engine.match(ruleset.function_rules, func_features, int(f)) + _, matches = capa.engine.match(ruleset.function_rules, func_features, f.address) for (name, res) in matches.items(): func_matches[name].extend(res) @@ -121,19 +122,19 @@ def find_func_matches(f, ruleset, func_features, bb_features): def find_file_features(extractor): """ """ - file_features = collections.defaultdict(set) - for (feature, ea) in extractor.extract_file_features(): - if ea: - file_features[feature].add(ea) + file_features = collections.defaultdict(set) # type: FeatureSet + for (feature, addr) in extractor.extract_file_features(): + if addr: + file_features[feature].add(addr) else: if feature not in file_features: file_features[feature] = set() return file_features -def find_file_matches(ruleset, file_features): +def find_file_matches(ruleset, file_features: FeatureSet): """ """ - _, matches = capa.engine.match(ruleset.file_rules, file_features, 0x0) + _, matches = capa.engine.match(ruleset.file_rules, file_features, NO_ADDRESS) return matches @@ -876,10 +877,10 @@ class CapaExplorerForm(idaapi.PluginForm): try: f = idaapi.get_func(idaapi.get_screen_ea()) if f: - f = extractor.get_function(f.start_ea) - self.rulegen_current_function = f + fh: FunctionHandle = extractor.get_function(f.start_ea) + self.rulegen_current_function = fh - func_features, bb_features = find_func_features(f, extractor) + func_features, bb_features = find_func_features(fh, extractor) self.rulegen_func_features_cache = collections.defaultdict(set, copy.copy(func_features)) self.rulegen_bb_features_cache = collections.defaultdict(dict, copy.copy(bb_features)) @@ -890,13 +891,13 @@ class CapaExplorerForm(idaapi.PluginForm): try: # add function and bb rule matches to function features, for display purposes - func_matches, bb_matches = find_func_matches(f, self.ruleset_cache, func_features, bb_features) - for (name, res) in itertools.chain(func_matches.items(), bb_matches.items()): + func_matches, bb_matches = find_func_matches(fh, self.ruleset_cache, func_features, bb_features) + for (name, addrs) in itertools.chain(func_matches.items(), bb_matches.items()): rule = self.ruleset_cache[name] if rule.meta.get("capa/subscope-rule"): continue - for (ea, _) in res: - func_features[capa.features.common.MatchedRule(name)].add(ea) + for (addr, _) in addrs: + func_features[capa.features.common.MatchedRule(name)].add(addr) except Exception as e: logger.error("Failed to match function/basic block rule scope (error: %s)", e) return False @@ -916,7 +917,7 @@ class CapaExplorerForm(idaapi.PluginForm): try: file_features = find_file_features(extractor) - self.rulegen_file_features_cache = collections.defaultdict(dict, copy.copy(file_features)) + self.rulegen_file_features_cache = copy.copy(file_features) if ida_kernwin.user_cancelled(): logger.info("User cancelled analysis.") @@ -925,12 +926,12 @@ class CapaExplorerForm(idaapi.PluginForm): try: # add file matches to file features, for display purposes - for (name, res) in find_file_matches(self.ruleset_cache, file_features).items(): + for (name, addrs) in find_file_matches(self.ruleset_cache, file_features).items(): rule = self.ruleset_cache[name] if rule.meta.get("capa/subscope-rule"): continue - for (ea, _) in res: - file_features[capa.features.common.MatchedRule(name)].add(ea) + for (addr, _) in addrs: + file_features[capa.features.common.MatchedRule(name)].add(addr) except Exception as e: logger.error("Failed to match file scope rules (error: %s)", e) return False @@ -946,7 +947,7 @@ class CapaExplorerForm(idaapi.PluginForm): try: # load preview and feature tree self.view_rulegen_preview.load_preview_meta( - f.start_ea if f else None, + fh.address if fh else None, settings.user.get(CAPA_SETTINGS_RULEGEN_AUTHOR, ""), settings.user.get(CAPA_SETTINGS_RULEGEN_SCOPE, "function"), ) @@ -957,7 +958,7 @@ class CapaExplorerForm(idaapi.PluginForm): "capa rules directory: %s (%d rules)" % (settings.user[CAPA_SETTINGS_RULE_PATH], len(self.rules_cache)) ) except Exception as e: - logger.error("Failed to render views (error: %s)", e) + logger.error("Failed to render views (error: %s)", e, exc_info=True) return False return True diff --git a/capa/ida/plugin/view.py b/capa/ida/plugin/view.py index 3a1d324d..b68d7f28 100644 --- a/capa/ida/plugin/view.py +++ b/capa/ida/plugin/view.py @@ -18,6 +18,7 @@ import capa.ida.helpers import capa.features.common import capa.features.basicblock from capa.ida.plugin.item import CapaExplorerFunctionItem +from capa.features.address import NO_ADDRESS from capa.ida.plugin.model import CapaExplorerDataModel MAX_SECTION_SIZE = 750 @@ -1010,6 +1011,8 @@ class CapaExplorerRulegenFeatures(QtWidgets.QTreeWidget): self.parent_items = {} def format_address(e): + if e == NO_ADDRESS: + return "" return "%X" % e if e else "" def format_feature(feature): @@ -1020,7 +1023,7 @@ class CapaExplorerRulegenFeatures(QtWidgets.QTreeWidget): value = '"%s"' % capa.features.common.escape_string(value) return "%s(%s)" % (name, value) - for (feature, eas) in sorted(features.items(), key=lambda k: sorted(k[1])): + for (feature, addrs) in sorted(features.items(), key=lambda k: sorted(k[1])): if isinstance(feature, capa.features.basicblock.BasicBlock): # filter basic blocks for now, we may want to add these back in some time # in the future @@ -1032,7 +1035,7 @@ class CapaExplorerRulegenFeatures(QtWidgets.QTreeWidget): # level 1 if feature not in self.parent_items: - if len(eas) > 1: + if len(addrs) > 1: self.parent_items[feature] = self.new_parent_node( self.parent_items[type(feature)], (format_feature(feature),), feature=feature ) @@ -1042,18 +1045,18 @@ class CapaExplorerRulegenFeatures(QtWidgets.QTreeWidget): ) # level n > 1 - if len(eas) > 1: - for ea in sorted(eas): + if len(addrs) > 1: + for addr in sorted(addrs): self.new_leaf_node( - self.parent_items[feature], (format_feature(feature), format_address(ea)), feature=feature + self.parent_items[feature], (format_feature(feature), format_address(addr)), feature=feature ) else: - if eas: - ea = eas.pop() + if addrs: + addr = addrs.pop() else: # some features may not have an address e.g. "format" - ea = "" - for (i, v) in enumerate((format_feature(feature), format_address(ea))): + addr = "" + for (i, v) in enumerate((format_feature(feature), format_address(addr))): self.parent_items[feature].setText(i, v) self.parent_items[feature].setData(0, 0x100, feature)