def multi_collection()

in detection_rules/cli_utils.py [0:0]
39 lines of code
18 McCabe index (conditional complexity)

def multi_collection(f):
    """Add arguments to get a RuleCollection by file, directory or a list of IDs"""
    from .misc import client_error

    @click.option("--rule-file", "-f", multiple=True, type=click.Path(dir_okay=False), required=False)
    @click.option("--directory", "-d", multiple=True, type=click.Path(file_okay=False), required=False,
                  help="Recursively load rules from a directory")
    @click.option("--rule-id", "-id", multiple=True, required=False)
    @functools.wraps(f)
    def get_collection(*args, **kwargs):
        rule_id: List[str] = kwargs.pop("rule_id", [])
        rule_files: List[str] = kwargs.pop("rule_file")
        directories: List[str] = kwargs.pop("directory")

        rules = RuleCollection()

        if not (directories or rule_id or rule_files or (DEFAULT_PREBUILT_RULES_DIRS + DEFAULT_PREBUILT_BBR_DIRS)):
            client_error("Required: at least one of --rule-id, --rule-file, or --directory")

        rules.load_files(Path(p) for p in rule_files)
        rules.load_directories(Path(d) for d in directories)

        if rule_id:
            rules.load_directories(
                DEFAULT_PREBUILT_RULES_DIRS + DEFAULT_PREBUILT_BBR_DIRS, obj_filter=dict_filter(rule__rule_id=rule_id)
            )
            found_ids = {rule.id for rule in rules}
            missing = set(rule_id).difference(found_ids)

            if missing:
                client_error(f'Could not find rules with IDs: {", ".join(missing)}')
        elif not rule_files and not directories:
            rules.load_directories(Path(d) for d in (DEFAULT_PREBUILT_RULES_DIRS + DEFAULT_PREBUILT_BBR_DIRS))

        if len(rules) == 0:
            client_error("No rules found")

        # Warn that if the path does not match the expected path, it will be saved to the expected path
        for rule in rules:
            threat = rule.contents.data.get("threat")
            first_tactic = threat[0].tactic.name if threat else ""
            rule_name = rulename_to_filename(rule.contents.data.name, tactic_name=first_tactic)
            if rule.path.name != rule_name:
                click.secho(
                    f"WARNING: Rule path does not match required path: {rule.path.name} != {rule_name}", fg="yellow"
                )

        kwargs["rules"] = rules
        return f(*args, **kwargs)

    return get_collection