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

Code before:
        raise ShouldExitError(E_INVALID_FILE_ARCH) from e
    except UnsupportedOSError as e:
        log_unsupported_os_error()
        raise ShouldExitError(E_INVALID_FILE_OS) from e


def get_extractor_filters_from_cli(args, input_format) -> Optional[Set]:
    if input_format in STATIC_FORMATS:
        if args.processes:
            raise InvalidArgument("Cannot filter processes with static analysis.")
        return set(map(str_to_number, args.functions))
    elif input_format in DYNAMIC_FORMATS:
        if args.functions:
            raise InvalidArgument("Cannot filter functions with dynamic analysis.")
        return set(map(str_to_number, args.processes))
    else:
        raise ShouldExitError(E_INVALID_INPUT_FORMAT)


def apply_extractor_filters(extractor: FeatureExtractor, elements: Set):
    if isinstance(extractor, StaticFeatureExtractor):
        return FunctionFilter(extractor, elements)
    elif isinstance(extractor, DynamicFeatureExtractor):
        return ProcessFilter(extractor, elements)
    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+")

    if argv is None:
        argv = sys.argv[1:]

    desc = "The FLARE team's open-source tool to identify capabilities in executable files."
    epilog = textwrap.dedent(
        """
        By default, capa uses a default set of embedded rules.
        You can see the rule set here:
          https://github.com/mandiant/capa-rules

        To provide your own rule set, use the `-r` flag:
          capa  --rules /path/to/rules  suspicious.exe
          capa  -r      /path/to/rules  suspicious.exe

        examples:
          identify capabilities in a binary
            capa suspicious.exe

          identify capabilities in 32-bit shellcode, see `-f` for all supported formats

Review comment (suggestion):
I don't think this routine is designed so that it can be easily extended, such as also supporting thread/basic block filters.

Code after:
    should_save_workspace = os.environ.get("CAPA_SAVE_WORKSPACE") not in ("0", "no", "NO", "n", None)

    os_ = get_os_from_cli(args, backend)
    sample_path = get_sample_path_from_cli(args, backend)
    extractor_filters = get_extractor_filters_from_cli(args, input_format)

    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: