enhance capability handling with new Capabilities dataclass and update related functions

Signed-off-by: vibhatsu <maulikbarot2915@gmail.com>

Author: v1bh475u

capa/capabilities/common.py

@@ -58,11 +58,33 @@ def find_file_capabilities(
     return FileCapabilities(features, matches, len(file_features))
 
 
+@dataclass
+class Capabilities:
+    matches: MatchResults
+    feature_counts: StaticFeatureCounts | DynamicFeatureCounts
+    library_functions: Optional[tuple[LibraryFunction, ...]] = None
+
+
+def find_capabilities(ruleset: RuleSet, extractor: FeatureExtractor, disable_progress=None, **kwargs) -> Capabilities:
+    from capa.capabilities.static import find_static_capabilities
+    from capa.capabilities.dynamic import find_dynamic_capabilities
+
+    if isinstance(extractor, StaticFeatureExtractor):
+        # for the time being, extractors are either static or dynamic.
+        # Remove this assertion once that has changed
+        assert not isinstance(extractor, DynamicFeatureExtractor)
+        return find_static_capabilities(ruleset, extractor, disable_progress=disable_progress, **kwargs)
+    if isinstance(extractor, DynamicFeatureExtractor):
+        return find_dynamic_capabilities(ruleset, extractor, disable_progress=disable_progress, **kwargs)
+
+    raise ValueError(f"unexpected extractor type: {extractor.__class__.__name__}")
+
+
 def is_static_limitation_rule(r: Rule) -> bool:
     return r.meta.get("namespace", "") == "internal/limitation/static"
 
 
-def has_static_limitation(rules: RuleSet, capabilities: MatchResults, is_standalone=True) -> bool:
+def has_static_limitation(rules: RuleSet, capabilities: Capabilities | FileCapabilities, is_standalone=True) -> bool:
 
     file_limitation_rules = list(filter(lambda r: is_static_limitation_rule(r), rules.rules.values()))
 
@@ -73,16 +95,16 @@ def is_dynamic_limitation_rule(r: Rule) -> bool:
     return r.meta.get("namespace", "") == "internal/limitation/dynamic"
 
 
-def has_dynamic_limitation(rules: RuleSet, capabilities: MatchResults, is_standalone=True) -> bool:
+def has_dynamic_limitation(rules: RuleSet, capabilities: Capabilities | FileCapabilities, is_standalone=True) -> bool:
 
     dynamic_limitation_rules = list(filter(lambda r: is_dynamic_limitation_rule(r), rules.rules.values()))
     return has_limitation(dynamic_limitation_rules, capabilities, is_standalone)
 
 
-def has_limitation(rules: list, capabilities: MatchResults, is_standalone: bool) -> bool:
+def has_limitation(rules: list, capabilities: Capabilities | FileCapabilities, is_standalone: bool) -> bool:
 
     for rule in rules:
-        if rule.name not in capabilities:
+        if rule.name not in capabilities.matches:
             continue
         logger.warning("-" * 80)
         for line in rule.meta.get("description", "").split("\n"):
@@ -96,25 +118,3 @@ def has_limitation(rules: list, capabilities: MatchResults, is_standalone: bool)
         # bail on first file limitation
         return True
     return False
-
-
-@dataclass
-class Capabilities:
-    matches: MatchResults
-    feature_counts: StaticFeatureCounts | DynamicFeatureCounts
-    library_functions: Optional[tuple[LibraryFunction, ...]] = None
-
-
-def find_capabilities(ruleset: RuleSet, extractor: FeatureExtractor, disable_progress=None, **kwargs) -> Capabilities:
-    from capa.capabilities.static import find_static_capabilities
-    from capa.capabilities.dynamic import find_dynamic_capabilities
-
-    if isinstance(extractor, StaticFeatureExtractor):
-        # for the time being, extractors are either static or dynamic.
-        # Remove this assertion once that has changed
-        assert not isinstance(extractor, DynamicFeatureExtractor)
-        return find_static_capabilities(ruleset, extractor, disable_progress=disable_progress, **kwargs)
-    if isinstance(extractor, DynamicFeatureExtractor):
-        return find_dynamic_capabilities(ruleset, extractor, disable_progress=disable_progress, **kwargs)
-
-    raise ValueError(f"unexpected extractor type: {extractor.__class__.__name__}")

capa/main.py

@@ -100,6 +100,7 @@
     FORMAT_BINEXPORT2,
 )
 from capa.capabilities.common import (
+    Capabilities,
     find_capabilities,
     has_static_limitation,
     find_file_capabilities,
@@ -797,7 +798,7 @@ def find_dynamic_limitations_from_cli(args, rules: RuleSet, file_extractors: lis
     """
     found_dynamic_limitation = False
     for file_extractor in file_extractors:
-        pure_dynamic_capabilities, _ = find_file_capabilities(rules, file_extractor, {})
+        pure_dynamic_capabilities = find_file_capabilities(rules, file_extractor, {})
         found_dynamic_limitation = has_dynamic_limitation(rules, pure_dynamic_capabilities)
 
     # file limitations that rely on non-file scope won't be detected here.