Repository: mandiant/capa
File: capa/capabilities/common.py (Python)

Code before:


def find_file_capabilities(ruleset: RuleSet, extractor: FeatureExtractor, function_features: FeatureSet):
    file_features: FeatureSet = collections.defaultdict(set)

    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.
        if va:
            file_features[feature].add(va)
        else:
            if feature not in file_features:
                file_features[feature] = set()

    logger.debug("analyzed file and extracted %d features", len(file_features))

    file_features.update(function_features)

    _, matches = ruleset.match(Scope.FILE, file_features, NO_ADDRESS)
    return matches, len(file_features)


def has_static_limitation(rules: RuleSet, capabilities: MatchResults, is_standalone=True) -> bool:
    def is_static_limitation_rule(r: Rule) -> bool:
        return r.meta.get("namespace", "") == "internal/limitation/static/file"

    file_limitation_rules = list(filter(lambda r: is_static_limitation_rule(r), rules.rules.values()))

    return has_limitation(file_limitation_rules, capabilities, is_standalone)


def has_dynamic_limitation(rules: RuleSet, capabilities: MatchResults, is_standalone=True) -> bool:
    def is_dynamic_limitation_rule(r: Rule) -> bool:
        return r.meta.get("namespace", "") == "internal/limitation/dynamic"

    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:
    # given list of rules and capabilities, finds the limitation if it exists, logs warning and returns True
    # else logs nothing and returns False

    for rule in rules:
        if rule.name not in capabilities:
            continue
        logger.warning("-" * 80)
        for line in rule.meta.get("description", "").split("\n"):
            logger.warning(" %s", line)
        logger.warning(" Identified via rule: %s", rule.name)

Review comment (suggestion):
```suggestion
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:
```

Code after:

logger = logging.getLogger(__name__)


@dataclass
class FileCapabilities:
    features: FeatureSet
    matches: MatchResults
    feature_count: int


def find_file_capabilities(
    ruleset: RuleSet, extractor: FeatureExtractor, function_features: FeatureSet
) -> FileCapabilities:
    file_features: FeatureSet = collections.defaultdict(set)

    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.
        if va:
            file_features[feature].add(va)
        else:
            if feature not in file_features:
                file_features[feature] = set()

    logger.debug("analyzed file and extracted %d features", len(file_features))

    file_features.update(function_features)

    features, matches = ruleset.match(Scope.FILE, file_features, NO_ADDRESS)
    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):