diff --git a/capa/features/extractors/base_extractor.py b/capa/features/extractors/base_extractor.py index 4b6be1ee..c287e75b 100644 --- a/capa/features/extractors/base_extractor.py +++ b/capa/features/extractors/base_extractor.py @@ -58,7 +58,7 @@ class FeatureExtractor: fetch the preferred load address at which the sample was analyzed. """ raise NotImplemented - + @abc.abstractmethod def extract_global_features(self) -> Iterator[Tuple[Feature, int]]: """ @@ -74,7 +74,7 @@ class FeatureExtractor: Tuple[Feature, int]: feature and its location """ raise NotImplemented - + @abc.abstractmethod def extract_file_features(self) -> Iterator[Tuple[Feature, int]]: """ @@ -272,12 +272,12 @@ class NullFeatureExtractor(FeatureExtractor): def get_base_address(self): return self.features["base address"] - + def extract_global_features(self): for p in self.features.get("global features", []): va, feature = p yield feature, va - + def extract_file_features(self): for p in self.features.get("file features", []): va, feature = p diff --git a/capa/features/extractors/pefile.py b/capa/features/extractors/pefile.py index a204107c..8f6a3ed1 100644 --- a/capa/features/extractors/pefile.py +++ b/capa/features/extractors/pefile.py @@ -150,7 +150,7 @@ FILE_HANDLERS = ( extract_file_function_names, extract_file_format, ) - + def extract_global_features(pe, buf): """ diff --git a/capa/features/freeze.py b/capa/features/freeze.py index 42e3d8c9..9c0d0c98 100644 --- a/capa/features/freeze.py +++ b/capa/features/freeze.py @@ -192,7 +192,7 @@ def loads(s): va = int(va, 0x10) feature = deserialize_feature(feature[:2]) features["global features"].append((va, feature)) - + for feature in doc.get("scopes", {}).get("file", []): va, loc = feature[2:] va = int(va, 0x10) diff --git a/capa/main.py b/capa/main.py index ef380581..a7e1fb8a 100644 --- a/capa/main.py +++ b/capa/main.py @@ -88,12 +88,16 @@ def find_function_capabilities(ruleset: RuleSet, extractor: FeatureExtractor, f: # - basic blocks bb_features = collections.defaultdict(set) - for feature, va in itertools.chain(extractor.extract_basic_block_features(f, bb), extractor.extract_global_features()): + for feature, va in itertools.chain( + extractor.extract_basic_block_features(f, bb), extractor.extract_global_features() + ): bb_features[feature].add(va) function_features[feature].add(va) for insn in extractor.get_instructions(f, bb): - for feature, va in itertools.chain(extractor.extract_insn_features(f, bb, insn), extractor.extract_global_features()): + for feature, va in itertools.chain( + extractor.extract_insn_features(f, bb, insn), extractor.extract_global_features() + ): bb_features[feature].add(va) function_features[feature].add(va) diff --git a/tests/fixtures.py b/tests/fixtures.py index 8283502d..1b404bb2 100644 --- a/tests/fixtures.py +++ b/tests/fixtures.py @@ -139,7 +139,7 @@ def extract_global_features(extractor): for feature, va in extractor.extract_global_features(): features[feature].add(va) return features - + @lru_cache() def extract_file_features(extractor):