Repository: mandiant/capa
File: capa/main.py (Python)

Code before:

    all_file_matches, feature_count = find_file_capabilities(ruleset, extractor, function_and_lower_features)
    feature_counts.file = feature_count

    matches = {
        rule_name: results
        for rule_name, results in itertools.chain(
            # each rule exists in exactly one scope,
            # so there won't be any overlap among these following MatchResults,
            # and we can merge the dictionaries naively.
            all_insn_matches.items(),
            all_bb_matches.items(),
            all_function_matches.items(),
            all_file_matches.items(),
        )
    }

    meta = {
        "feature_counts": feature_counts,
        "library_functions": library_functions,
    }

    return matches, meta


def find_capabilities(ruleset: RuleSet, extractor: FeatureExtractor, **kwargs) -> Tuple[MatchResults, Any]:
    if isinstance(extractor, StaticFeatureExtractor):
        extractor_: StaticFeatureExtractor = cast(StaticFeatureExtractor, extractor)
        return find_capabilities_static(ruleset, extractor_, kwargs)
    else:
        # extractor_ = cast(DynamicFeatureExtractor, extractor)
        print("nni")


# TODO move all to helpers?
def has_rule_with_namespace(rules, capabilities, rule_cat):
    for rule_name in capabilities.keys():
        if rules.rules[rule_name].meta.get("namespace", "").startswith(rule_cat):
            return True
    return False


def is_internal_rule(rule: Rule) -> bool:
    return rule.meta.get("namespace", "").startswith("internal/")


def is_file_limitation_rule(rule: Rule) -> bool:
    return rule.meta.get("namespace", "") == "internal/limitation/file"


def has_file_limitation(rules: RuleSet, capabilities: MatchResults, is_standalone=True) -> bool:

Review comment (suggestion):
the `isinstance(...)` type assertion *should* be enough for mypy to understand the types, but using `cast` might be even more explicit. i dont feel strongly that it should or shouldn't be present, just info.

Code after:
    feature_counts.file = feature_count

    matches = {
        rule_name: results
        for rule_name, results in itertools.chain(
            # each rule exists in exactly one scope,
            # so there won't be any overlap among these following MatchResults,
            # and we can merge the dictionaries naively.
            all_insn_matches.items(),
            all_bb_matches.items(),
            all_function_matches.items(),
            all_file_matches.items(),
        )
    }

    meta = {
        "feature_counts": feature_counts,
        "library_functions": library_functions,
    }

    return matches, meta


def find_capabilities(ruleset: RuleSet, extractor: FeatureExtractor, **kwargs) -> Tuple[MatchResults, Any]:
    if isinstance(extractor, StaticFeatureExtractor):
        extractor_: StaticFeatureExtractor = cast(StaticFeatureExtractor, extractor)
        return find_static_capabilities(ruleset, extractor_, kwargs)
    elif isinstance(extractor, DynamicFeatureExtractor):
        # extractor_ = cast(DynamicFeatureExtractor, extractor)
        raise NotImplementedError()
    else:
        raise ValueError(f"unexpected extractor type: {extractor.__class__.__name__}")


# TODO move all to helpers?
def has_rule_with_namespace(rules, capabilities, rule_cat):
    for rule_name in capabilities.keys():
        if rules.rules[rule_name].meta.get("namespace", "").startswith(rule_cat):
            return True
    return False


def is_internal_rule(rule: Rule) -> bool:
    return rule.meta.get("namespace", "").startswith("internal/")


def is_file_limitation_rule(rule: Rule) -> bool:
    return rule.meta.get("namespace", "") == "internal/limitation/file"


def has_file_limitation(rules: RuleSet, capabilities: MatchResults, is_standalone=True) -> bool: