Repository: mandiant/capa
File: scripts/show-unused-features.py (Python)

Code before:
            if args.function not in [format_address(fh.address) for fh in function_handles]:
                print(f"{args.function} not a function")
                return -1

        if len(function_handles) == 0:
            print(f"{args.function} not a function")
            return -1

    feature_map.update(get_file_features(function_handles, extractor))

    rules_feature_set = get_rules_feature_set(args.rules)

    highlight_unused_features(feature_map, rules_feature_set)
    return 0


def ida_main():
    import idc

    import capa.main
    import capa.features.extractors.ida.extractor

    function = idc.get_func_attr(idc.here(), idc.FUNCATTR_START)
    print(f"getting features for current function {hex(function)}")

    extractor = capa.features.extractors.ida.extractor.IdaFeatureExtractor()
    feature_map: dict = {}

    if not function:
        for feature, _ in extractor.extract_file_features():
            feature_map[feature] = feature_map.get(feature, 0) + 1
        return

    function_handles = tuple(extractor.get_functions())

    if function:
        function_handles = tuple(filter(lambda fh: fh.inner.start_ea == function, function_handles))

        if len(function_handles) == 0:
            print(f"{hex(function)} not a function")
            return -1

    feature_map.update(get_file_features(function_handles, extractor))

    rules_path = capa.main.get_default_root() / "rules"
    rules_feature_set = get_rules_feature_set([rules_path])

    highlight_unused_features(feature_map, rules_feature_set)

    return 0


Review comment (suggestion):
i regret how much code is needed to initialize capa and fetch the rules and features. we havent yet found a simpler way to reduce this boilerplate, but maybe we should look into this more.

its a lot of code and is repeated a bunch of times in different scripts. that means its hard to keep up to date and is at risk of bitrotting.

nothing to do here, just commentary.

Code after:

    feature_map: typing.Counter[Feature] = Counter()

    feature_map.update([feature for feature, _ in extractor.extract_global_features()])

    function_handles: Tuple[FunctionHandle, ...]
    if isinstance(extractor, capa.features.extractors.pefile.PefileFeatureExtractor):
        # pefile extractor doesn't extract function features
        function_handles = ()
    else:
        function_handles = tuple(extractor.get_functions())

    if args.function:
        if args.format == "freeze":
            function_handles = tuple(filter(lambda fh: fh.address == args.function, function_handles))
        else:
            function_handles = tuple(filter(lambda fh: format_address(fh.address) == args.function, function_handles))

            if args.function not in [format_address(fh.address) for fh in function_handles]:
                print(f"{args.function} not a function")
                return -1

        if len(function_handles) == 0:
            print(f"{args.function} not a function")
            return -1

    feature_map.update(get_file_features(function_handles, extractor))

    rules_feature_set = get_rules_feature_set(args.rules)

    print_unused_features(feature_map, rules_feature_set)
    return 0


def ida_main():
    import idc

    import capa.main
    import capa.features.extractors.ida.extractor

    function = idc.get_func_attr(idc.here(), idc.FUNCATTR_START)
    print(f"getting features for current function {hex(function)}")

    extractor = capa.features.extractors.ida.extractor.IdaFeatureExtractor()
    feature_map: typing.Counter[Feature] = Counter()

    feature_map.update([feature for feature, _ in extractor.extract_file_features()])

    function_handles = tuple(extractor.get_functions())

    if function: