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

Code before:
            backend,
            sig_paths,
            should_save_workspace=should_save_workspace,
            disable_progress=args.quiet or args.debug,
            sample_path=sample_path,
        )
    except UnsupportedFormatError as e:
        if input_format == FORMAT_CAPE:
            log_unsupported_cape_report_error(str(e))
        elif input_format == FORMAT_DRAKVUF:
            log_unsupported_drakvuf_report_error(str(e))
        else:
            log_unsupported_format_error()
        raise ShouldExitError(E_INVALID_FILE_TYPE) from e
    except UnsupportedArchError as e:
        log_unsupported_arch_error()
        raise ShouldExitError(E_INVALID_FILE_ARCH) from e
    except UnsupportedOSError as e:
        log_unsupported_os_error()
        raise ShouldExitError(E_INVALID_FILE_OS) from e
    except capa.loader.CorruptFile as e:
        logger.error("Input file '%s' is not a valid file: %s", args.input_file, str(e))
        raise ShouldExitError(E_CORRUPT_FILE) from e


def get_extractor_filters_from_cli(args, input_format) -> Dict[str, Set]:
    if input_format in STATIC_FORMATS:
        if args.restrict_to_processes:
            raise InvalidArgument("Cannot filter processes with static analysis.")
        return {"functions": {int(addr, 0) for addr in args.restrict_to_functions}}
    elif input_format in DYNAMIC_FORMATS:
        if args.restrict_to_functions:
            raise InvalidArgument("Cannot filter functions with dynamic analysis.")
        return {"processes": {int(pid, 0) for pid in args.restrict_to_processes}}
    else:
        raise ShouldExitError(E_INVALID_INPUT_FORMAT)


def apply_extractor_filters(extractor: FeatureExtractor, extractor_filters: Dict[str, Set]):
    if isinstance(extractor, StaticFeatureExtractor):
        return FunctionFilter(extractor, extractor_filters["functions"])
    elif isinstance(extractor, DynamicFeatureExtractor):
        return ProcessFilter(extractor, extractor_filters["processes"])
    else:
        raise ShouldExitError(E_INVALID_FEATURE_EXTRACTOR)


def main(argv: Optional[List[str]] = None):
    if sys.version_info < (3, 8):
        raise UnsupportedRuntimeError("This version of capa can only be used with Python 3.8+")


Review comment (suggestion):
please create a type for the return value, or use `typing.TypedDict` to specify the expected keys in this dict.

Code after:

    try:
        extractor = capa.loader.get_extractor(
            args.input_file,
            input_format,
            os_,
            backend,
            sig_paths,
            should_save_workspace=should_save_workspace,
            disable_progress=args.quiet or args.debug,
            sample_path=sample_path,
        )
        return apply_extractor_filters(extractor, extractor_filters)
    except UnsupportedFormatError as e:
        if input_format == FORMAT_CAPE:
            log_unsupported_cape_report_error(str(e))
        elif input_format == FORMAT_DRAKVUF:
            log_unsupported_drakvuf_report_error(str(e))
        else:
            log_unsupported_format_error()
        raise ShouldExitError(E_INVALID_FILE_TYPE) from e
    except UnsupportedArchError as e:
        log_unsupported_arch_error()
        raise ShouldExitError(E_INVALID_FILE_ARCH) from e
    except UnsupportedOSError as e:
        log_unsupported_os_error()
        raise ShouldExitError(E_INVALID_FILE_OS) from e
    except capa.loader.CorruptFile as e:
        logger.error("Input file '%s' is not a valid file: %s", args.input_file, str(e))
        raise ShouldExitError(E_CORRUPT_FILE) from e


def get_extractor_filters_from_cli(args, input_format) -> FilterConfig:
    if not hasattr(args, "restrict_to_processes") and not hasattr(args, "restrict_to_functions"):
        # no processes or function filters were installed in the args
        return {}

    if input_format in STATIC_FORMATS:
        if args.restrict_to_processes:
            raise InvalidArgument("Cannot filter processes with static analysis.")
        return {"functions": {int(addr, 0) for addr in args.restrict_to_functions}}
    elif input_format in DYNAMIC_FORMATS:
        if args.restrict_to_functions:
            raise InvalidArgument("Cannot filter functions with dynamic analysis.")
        return {"processes": {int(pid, 0) for pid in args.restrict_to_processes}}
    else:
        raise ShouldExitError(E_INVALID_INPUT_FORMAT)


def apply_extractor_filters(extractor: FeatureExtractor, extractor_filters: FilterConfig):
    if not any(extractor_filters.values()):