From 63e4d3d5eb36d45b90d9d674f8a7e8029a05e1ac Mon Sep 17 00:00:00 2001 From: Yacine Elhamer Date: Mon, 26 Jun 2023 21:14:17 +0100 Subject: [PATCH] fix TypeAlias importing: import from typing_extensions to support Python 3.9 and lower --- capa/features/extractors/base_extractor.py | 4 +++- capa/main.py | 14 ++++---------- scripts/profile-time.py | 5 +++-- 3 files changed, 10 insertions(+), 13 deletions(-) diff --git a/capa/features/extractors/base_extractor.py b/capa/features/extractors/base_extractor.py index 75db33fa..798fa8be 100644 --- a/capa/features/extractors/base_extractor.py +++ b/capa/features/extractors/base_extractor.py @@ -8,9 +8,11 @@ import abc import dataclasses -from typing import Any, Dict, Tuple, Union, Iterator, TypeAlias +from typing import Any, Dict, Tuple, Union, Iterator from dataclasses import dataclass +from typing_extensions import TypeAlias + import capa.features.address from capa.features.common import Feature from capa.features.address import Address, AbsoluteVirtualAddress diff --git a/capa/main.py b/capa/main.py index 22766ede..85abb942 100644 --- a/capa/main.py +++ b/capa/main.py @@ -231,14 +231,7 @@ def find_code_capabilities( def find_file_capabilities(ruleset: RuleSet, extractor: FeatureExtractor, function_features: FeatureSet): file_features = collections.defaultdict(set) # type: FeatureSet - if isinstance(extractor, StaticFeatureExtractor): - extractor_: StaticFeatureExtractor = cast(StaticFeatureExtractor, extractor) - elif isinstance(extractor, DynamicFeatureExtractor): - extractor_: DynamicFeatureExtractor = cast(DynamicFeatureExtractor, extractor) - else: - raise ValueError(f"unexpected extractor type: {extractor.__class__.__name__}") - - for feature, va in itertools.chain(extractor_.extract_file_features(), extractor_.extract_global_features()): + for feature, va in itertools.chain(extractor.extract_file_features(), extractor.extract_global_features()): # not all file features may have virtual addresses. # if not, then at least ensure the feature shows up in the index. # the set of addresses will still be empty. @@ -1251,7 +1244,8 @@ def main(argv=None): if format_ == FORMAT_FREEZE: # freeze format deserializes directly into an extractor with open(args.sample, "rb") as f: - extractor = frz.load(f.read()) + extractor: FeatureExtractor = frz.load(f.read()) + assert isinstance(extractor, StaticFeatureExtractor) else: # all other formats we must create an extractor, # such as viv, binary ninja, etc. workspaces @@ -1270,7 +1264,7 @@ def main(argv=None): should_save_workspace = os.environ.get("CAPA_SAVE_WORKSPACE") not in ("0", "no", "NO", "n", None) try: - extractor: FeatureExtractor = get_extractor( + extractor = get_extractor( args.sample, format_, args.os, diff --git a/scripts/profile-time.py b/scripts/profile-time.py index 0bd4e389..2566a0fe 100644 --- a/scripts/profile-time.py +++ b/scripts/profile-time.py @@ -46,7 +46,7 @@ import capa.helpers import capa.features import capa.features.common import capa.features.freeze -from capa.features.extractors.base_extractor import FeatureExtractor +from capa.features.extractors.base_extractor import StaticFeatureExtractor logger = logging.getLogger("capa.profile") @@ -105,8 +105,9 @@ def main(argv=None): ): with open(args.sample, "rb") as f: extractor = capa.features.freeze.load(f.read()) + assert isinstance(extractor, StaticFeatureExtractor) else: - extractor: FeatureExtractor = capa.main.get_extractor( + extractor = capa.main.get_extractor( args.sample, args.format, args.os, capa.main.BACKEND_VIV, sig_paths, should_save_workspace=False )