diff --git a/capa/capabilities/common.py b/capa/capabilities/common.py index 6575fe7b..639a41b0 100644 --- a/capa/capabilities/common.py +++ b/capa/capabilities/common.py @@ -64,7 +64,9 @@ class Capabilities: library_functions: Optional[tuple[LibraryFunction, ...]] = None -def find_capabilities(ruleset: RuleSet, extractor: FeatureExtractor, disable_progress=None, **kwargs) -> Capabilities: +def find_capabilities( + ruleset: RuleSet, extractor: FeatureExtractor, disable_progress: Optional[bool] = None, **kwargs +) -> Capabilities: from capa.capabilities.static import find_static_capabilities from capa.capabilities.dynamic import find_dynamic_capabilities diff --git a/capa/capabilities/dynamic.py b/capa/capabilities/dynamic.py index e7e6594a..f8dbfc1b 100644 --- a/capa/capabilities/dynamic.py +++ b/capa/capabilities/dynamic.py @@ -15,6 +15,7 @@ import logging import itertools import collections +from typing import Optional from dataclasses import dataclass import capa.perf @@ -269,7 +270,7 @@ def find_process_capabilities( def find_dynamic_capabilities( - ruleset: RuleSet, extractor: DynamicFeatureExtractor, disable_progress: bool = False + ruleset: RuleSet, extractor: DynamicFeatureExtractor, disable_progress: Optional[bool] = None ) -> Capabilities: all_process_matches: MatchResults = collections.defaultdict(list) all_thread_matches: MatchResults = collections.defaultdict(list) @@ -285,7 +286,7 @@ def find_dynamic_capabilities( n_processes: int = len(processes) with capa.helpers.CapaProgressBar( - console=capa.helpers.log_console, transient=True, disable=disable_progress + console=capa.helpers.log_console, transient=True, disable=bool(disable_progress) ) as pbar: task = pbar.add_task("matching", total=n_processes, unit="processes") for p in processes: diff --git a/capa/capabilities/static.py b/capa/capabilities/static.py index 1047713b..a6a97faf 100644 --- a/capa/capabilities/static.py +++ b/capa/capabilities/static.py @@ -16,6 +16,7 @@ import time import logging import itertools import collections +from typing import Optional from dataclasses import dataclass import capa.perf @@ -149,7 +150,7 @@ def find_code_capabilities(ruleset: RuleSet, extractor: StaticFeatureExtractor, def find_static_capabilities( - ruleset: RuleSet, extractor: StaticFeatureExtractor, disable_progress=None + ruleset: RuleSet, extractor: StaticFeatureExtractor, disable_progress: Optional[bool] = None ) -> Capabilities: all_function_matches: MatchResults = collections.defaultdict(list) all_bb_matches: MatchResults = collections.defaultdict(list) @@ -168,7 +169,7 @@ def find_static_capabilities( percentage: float = 0 with capa.helpers.CapaProgressBar( - console=capa.helpers.log_console, transient=True, disable=disable_progress + console=capa.helpers.log_console, transient=True, disable=bool(disable_progress) ) as pbar: task = pbar.add_task( "matching", total=n_funcs, unit="functions", postfix=f"skipped {n_libs} library functions, {percentage}%" diff --git a/capa/features/extractors/common.py b/capa/features/extractors/common.py index f8918b8d..dc744c29 100644 --- a/capa/features/extractors/common.py +++ b/capa/features/extractors/common.py @@ -30,9 +30,11 @@ from capa.features.common import ( OS_ANY, OS_AUTO, ARCH_ANY, + VALID_OS, FORMAT_PE, FORMAT_ELF, OS_WINDOWS, + VALID_ARCH, FORMAT_FREEZE, FORMAT_RESULT, Arch, @@ -52,7 +54,7 @@ MATCH_RESULT = b'{"meta":' MATCH_JSON_OBJECT = b'{"' -def extract_file_strings(buf: bytes, **kwargs) -> Iterator[tuple[String, Address]]: +def extract_file_strings(buf: bytes) -> Iterator[tuple[String, Address]]: """ extract ASCII and UTF-16 LE strings from file """ @@ -97,7 +99,7 @@ def extract_arch(buf) -> Iterator[tuple[Feature, Address]]: with contextlib.closing(io.BytesIO(buf)) as f: arch = capa.features.extractors.elf.detect_elf_arch(f) - if arch not in capa.features.common.VALID_ARCH: + if arch not in VALID_ARCH: logger.debug("unsupported arch: %s", arch) return @@ -130,7 +132,7 @@ def extract_os(buf, os=OS_AUTO) -> Iterator[tuple[Feature, Address]]: with contextlib.closing(io.BytesIO(buf)) as f: os = capa.features.extractors.elf.detect_elf_os(f) - if os not in capa.features.common.VALID_OS: + if os not in VALID_OS: logger.debug("unsupported os: %s", os) return diff --git a/tests/test_capabilities.py b/tests/test_capabilities.py index 56406d62..b965b06f 100644 --- a/tests/test_capabilities.py +++ b/tests/test_capabilities.py @@ -226,7 +226,7 @@ def test_com_feature_matching(z395eb_extractor): """) ) ]) - capabilities = capa.main.find_capabilities(rules, z395eb_extractor) + capabilities = capa.capabilities.common.find_capabilities(rules, z395eb_extractor) assert "initialize IWebBrowser2" in capabilities.matches