23s2s

✏️ Editing: completions.py

"""
Shell auto-completion script generators for the CLI.

Introspects an argparse parser to enumerate all commands, subcommands, and
flags, then emits completion scripts for bash, zsh, and fish.
"""

import argparse
import re
from typing import Dict, List, Tuple

def _safe_identifier(prog: str) -> str:
    """Convert a prog name to a safe shell identifier (letters, digits, _)."""
    return re.sub(r"[^a-zA-Z0-9]", "_", prog)

def _collect_commands(
    parser: argparse.ArgumentParser,
) -> Dict[Tuple[str, ...], List[str]]:
    """Walk the parser tree and return {command_path: [flags]} mapping."""
    result: Dict[Tuple[str, ...], List[str]] = {}
    _walk_parser(parser, (), result)
    return result

def _get_flags(parser: argparse.ArgumentParser) -> List[str]:
    """Extract all optional flags from a parser (excluding help)."""
    flags = []
    for action in parser._actions:
        if isinstance(action, argparse._HelpAction):
            continue
        if isinstance(action, argparse._SubParsersAction):
            continue
        for opt in action.option_strings:
            flags.append(opt)
    return sorted(flags)

def _walk_parser(
    parser: argparse.ArgumentParser,
    path: Tuple[str, ...],
    result: Dict[Tuple[str, ...], List[str]],
):
    """Recursively walk subparsers and collect command paths + flags."""
    flags = _get_flags(parser)
    result[path] = flags

for action in parser._actions:
        if isinstance(action, argparse._SubParsersAction):
            for name, subparser in action.choices.items():
                _walk_parser(subparser, path + (name,), result)

def _get_subcommands(
    commands: Dict[Tuple[str, ...], List[str]],
    prefix: Tuple[str, ...],
) -> List[str]:
    """Get immediate subcommands of a given prefix."""
    subs = set()
    for path in commands:
        if len(path) == len(prefix) + 1 and path[: len(prefix)] == prefix:
            subs.add(path[-1])
    return sorted(subs)

def generate_bash(
    parser: argparse.ArgumentParser, prog: str = "imunify360-agent"
) -> str:
    """Generate a bash completion script."""
    commands = _collect_commands(parser)
    lines = []
    lines.append(f"# bash completion for {prog}")
    lines.append(f"# Auto-generated by {prog} completions bash")
    lines.append("")
    lines.append(f"_{_safe_identifier(prog)}_completions() {{")
    lines.append("    local cur prev words cword")
    lines.append("    if type _init_completion &>/dev/null; then")
    lines.append("        _init_completion || return")
    lines.append("    else")
    lines.append("        COMPREPLY=()")
    lines.append('        cur="${COMP_WORDS[COMP_CWORD]}"')
    lines.append('        prev="${COMP_WORDS[COMP_CWORD-1]}"')
    lines.append('        words=("${COMP_WORDS[@]}")')
    lines.append("        cword=$COMP_CWORD")
    lines.append("    fi")
    lines.append("")
    lines.append("    # Build the command path from words")
    lines.append('    local cmd_path=""')
    lines.append("    local i")
    lines.append("    for (( i=1; i < cword; i++ )); do")
    lines.append('        case "${words[i]}" in')
    lines.append("            -*) continue ;;")
    lines.append(
        '            *)  cmd_path="${cmd_path:+${cmd_path} }${words[i]}" ;;'
    )
    lines.append("        esac")
    lines.append("    done")
    lines.append("")
    lines.append('    case "$cmd_path" in')

# Sort by depth (deepest first) so more specific paths match first
    all_paths = sorted(commands.keys(), key=lambda p: (-len(p), p))
    for path in all_paths:
        if not path:
            continue
        subs = _get_subcommands(commands, path)
        flags = commands[path]
        completions = " ".join(subs + flags)
        pattern = " ".join(path)
        lines.append(f'        "{pattern}")')
        lines.append(
            f'            COMPREPLY=($(compgen -W "{completions}" -- "$cur"))'
        )
        lines.append("            return ;;")

# Root level
    root_subs = _get_subcommands(commands, ())
    root_flags = commands.get((), [])
    root_completions = " ".join(root_subs + root_flags)
    lines.append('        "")')
    lines.append(
        f'            COMPREPLY=($(compgen -W "{root_completions}" -- "$cur"))'
    )
    lines.append("            return ;;")
    lines.append("    esac")
    lines.append("}")
    lines.append("")
    lines.append(f"complete -F _{_safe_identifier(prog)}_completions {prog}")
    lines.append("")
    return "\n".join(lines)

def generate_zsh(
    parser: argparse.ArgumentParser, prog: str = "imunify360-agent"
) -> str:
    """Generate a zsh completion script."""
    commands = _collect_commands(parser)
    func_name = f"_{_safe_identifier(prog)}"
    lines = []
    lines.append(f"#compdef {prog}")
    lines.append(f"# zsh completion for {prog}")
    lines.append(f"# Auto-generated by {prog} completions zsh")
    lines.append("")
    lines.append(f"{func_name}() {{")
    lines.append("    local -a commands flags")
    lines.append("    local cmd_path")
    lines.append("")
    lines.append("    # Build command path from words")
    lines.append("    cmd_path=()")
    lines.append("    for word in ${words[2,-1]}; do")
    lines.append("        [[ $word == -* ]] && continue")
    lines.append('        [[ $word == "$words[$CURRENT]" ]] && continue')
    lines.append("        cmd_path+=($word)")
    lines.append("    done")
    lines.append("")
    lines.append('    case "${cmd_path[*]}" in')

all_paths = sorted(commands.keys(), key=lambda p: (-len(p), p))
    for path in all_paths:
        if not path:
            continue
        subs = _get_subcommands(commands, path)
        flags = commands[path]
        pattern = " ".join(path)
        lines.append(f'        "{pattern}")')
        if subs:
            desc_list = " ".join(f'"{s}"' for s in subs)
            lines.append(f"            commands=({desc_list})")
        if flags:
            flag_list = " ".join(f'"{f}"' for f in flags)
            lines.append(f"            flags=({flag_list})")
        lines.append(
            "            _describe 'command' commands -- flags && return"
            if subs
            else f"            compadd -- {' '.join(flags)} && return"
        )
        lines.append("            ;;")

# Root level
    root_subs = _get_subcommands(commands, ())
    root_flags = commands.get((), [])
    root_desc = " ".join(f'"{s}"' for s in root_subs)
    lines.append('        "")')
    lines.append(f"            commands=({root_desc})")
    if root_flags:
        flag_list = " ".join(f'"{f}"' for f in root_flags)
        lines.append(f"            flags=({flag_list})")
    lines.append("            _describe 'command' commands -- flags && return")
    lines.append("            ;;")
    lines.append("    esac")
    lines.append("}")
    lines.append("")
    lines.append(f"{func_name}")
    lines.append("")
    return "\n".join(lines)

def generate_fish(
    parser: argparse.ArgumentParser, prog: str = "imunify360-agent"
) -> str:
    """Generate a fish completion script."""
    commands = _collect_commands(parser)
    lines = []
    lines.append(f"# fish completion for {prog}")
    lines.append(f"# Auto-generated by {prog} completions fish")
    lines.append("")

# For each command path, emit completions
    # Fish uses conditions based on what subcommands have been entered
    for path in sorted(commands.keys(), key=lambda p: (len(p), p)):
        subs = _get_subcommands(commands, path)
        flags = commands[path]

if not path:
            # Root level subcommands
            condition = (
                "not __fish_seen_subcommand_from"
                f" {' '.join(_get_subcommands(commands, ()))}"
            )
            for sub in subs:
                lines.append(
                    f"complete -c {prog} -n '{condition}' -f -a '{sub}'"
                )
            for flag in flags:
                if flag.startswith("--"):
                    lines.append(
                        f"complete -c {prog} -n '{condition}' -l '{flag[2:]}'"
                    )
                elif flag.startswith("-"):
                    lines.append(
                        f"complete -c {prog} -n '{condition}' -s '{flag[1:]}'"
                    )
        else:
            # Build condition: must have seen parent commands but not children
            seen_parts = []
            for p in path:
                seen_parts.append(f"__fish_seen_subcommand_from {p}")
            condition = " && ".join(seen_parts)

child_subs = subs
            if child_subs:
                condition += (
                    " && not __fish_seen_subcommand_from"
                    f" {' '.join(child_subs)}"
                )

for sub in subs:
                lines.append(
                    f"complete -c {prog} -n '{condition}' -f -a '{sub}'"
                )
            for flag in flags:
                if flag.startswith("--"):
                    lines.append(
                        f"complete -c {prog} -n '{condition}' -l '{flag[2:]}'"
                    )
                elif flag.startswith("-"):
                    lines.append(
                        f"complete -c {prog} -n '{condition}' -s '{flag[1:]}'"
                    )

lines.append("")
    return "\n".join(lines)

GENERATORS = {
    "bash": generate_bash,
    "zsh": generate_zsh,
    "fish": generate_fish,
}

SUPPORTED_SHELLS = sorted(GENERATORS.keys())

def generate_completions(
    parser: argparse.ArgumentParser,
    shell: str,
    prog: str = "imunify360-agent",
) -> str:
    """Generate completion script for the given shell.

Raises ValueError if shell is not supported.
    """
    generator = GENERATORS.get(shell)
    if generator is None:
        raise ValueError(
            f"Unsupported shell: {shell}. "
            f"Supported shells: {', '.join(SUPPORTED_SHELLS)}"
        )
    return generator(parser, prog)

Sid Gifari From Gifari Industries - BD Cyber Security Team