diff --git a/capa/features/extractors/__init__.py b/capa/features/extractors/__init__.py index 2cdbd4af..9b3e2089 100644 --- a/capa/features/extractors/__init__.py +++ b/capa/features/extractors/__init__.py @@ -8,6 +8,8 @@ import abc +from capa.helpers import oint + class FeatureExtractor(object): """ @@ -35,6 +37,12 @@ class FeatureExtractor(object): # super(FeatureExtractor, self).__init__() + def block_offset(self, bb): + return oint(bb) + + def function_offset(self, f): + return oint(f) + @abc.abstractmethod def get_base_address(self): """ diff --git a/capa/features/extractors/miasm/__init__.py b/capa/features/extractors/miasm/__init__.py index 5881edae..54235afa 100644 --- a/capa/features/extractors/miasm/__init__.py +++ b/capa/features/extractors/miasm/__init__.py @@ -54,6 +54,9 @@ class MiasmFeatureExtractor(FeatureExtractor): def block_offset(self, bb): return bb.lines[0].offset + def function_offset(self, f): + return self.cfg.loc_key_to_block(f).lines[0].offset + def get_basic_blocks(self, loc_key): """ get the basic blocks of the function represented by lock_key diff --git a/capa/main.py b/capa/main.py index 27d5a9c3..62f75855 100644 --- a/capa/main.py +++ b/capa/main.py @@ -29,7 +29,7 @@ import capa.version import capa.features import capa.features.freeze import capa.features.extractors -from capa.helpers import oint, get_file_taste +from capa.helpers import get_file_taste RULES_PATH_DEFAULT_STRING = "(embedded rules)" SUPPORTED_FILE_MAGIC = set(["MZ"]) @@ -72,14 +72,14 @@ def find_function_capabilities(ruleset, extractor, f): bb_features[feature].add(va) function_features[feature].add(va) - _, matches = capa.engine.match(ruleset.basic_block_rules, bb_features, oint(bb)) + _, matches = capa.engine.match(ruleset.basic_block_rules, bb_features, extractor.block_offset(bb)) for rule_name, res in matches.items(): bb_matches[rule_name].extend(res) for va, _ in res: function_features[capa.features.MatchedRule(rule_name)].add(va) - _, function_matches = capa.engine.match(ruleset.function_rules, function_features, oint(f)) + _, function_matches = capa.engine.match(ruleset.function_rules, function_features, extractor.function_offset(f)) return function_matches, bb_matches, len(function_features) @@ -123,8 +123,8 @@ def find_capabilities(ruleset, extractor, disable_progress=None): for f in pbar(list(extractor.get_functions()), desc="matching", unit=" functions"): function_matches, bb_matches, feature_count = find_function_capabilities(ruleset, extractor, f) - meta["feature_counts"]["functions"][f.__int__()] = feature_count - logger.debug("analyzed function 0x%x and extracted %d features", f.__int__(), feature_count) + meta["feature_counts"]["functions"][extractor.function_offset(f)] = feature_count + logger.debug("analyzed function 0x%x and extracted %d features", extractor.function_offset(f), feature_count) for rule_name, res in function_matches.items(): all_function_matches[rule_name].extend(res) diff --git a/tests/fixtures.py b/tests/fixtures.py index 4261408b..cb844e6f 100644 --- a/tests/fixtures.py +++ b/tests/fixtures.py @@ -242,14 +242,14 @@ def sample(request): def get_function(extractor, fva): for f in extractor.get_functions(): - if f.__int__() == fva: + if extractor.function_offset(f) == fva: return f raise ValueError("function not found") def get_basic_block(extractor, f, va): for bb in extractor.get_basic_blocks(f): - if bb.__int__() == va: + if extractor.block_offset(bb) == va: return bb raise ValueError("basic block not found")