Add shell completion command

README.md

@@ -60,6 +60,24 @@ For alternative installation options and known issues, please refer to the [docu
 
 MVT provides two commands `mvt-ios` and `mvt-android`. [Check out the documentation to learn how to use them!](https://docs.mvt.re/)
 
+### Shell completion
+
+MVT can generate shell completion scripts for Bash, Zsh, and Fish:
+
+```bash
+mvt-ios completion
+mvt-android completion
+```
+
+The commands print setup instructions by default. To generate a completion script directly, pass the shell name:
+
+```bash
+mvt-ios completion bash
+mvt-android completion zsh
+```
+
+MVT only writes completion files or shell configuration when `--install` is passed. See the [command completion documentation](https://docs.mvt.re/en/latest/command_completion/) for details.
+
 
 ## License
 

docs/command_completion.md

@@ -1,43 +1,66 @@
-# Command Completion 
+# Command Completion
 
-MVT utilizes the [Click](https://click.palletsprojects.com/en/stable/) library for creating its command line interface. 
+MVT utilizes the [Click](https://click.palletsprojects.com/en/stable/) library for creating its command line interface.
 
 Click provides tab completion support for Bash (version 4.4 and up), Zsh, and Fish.
 
-To enable it, you need to manually register a special function with your shell, which varies depending on the shell you are using.
+To enable it, you need to register a completion script with your shell, which varies depending on the shell you are using.
 
-The following describes how to generate the command completion scripts and add them to your shell configuration. 
+The following describes how to generate the command completion scripts and add them to your shell configuration.
 
 > **Note: You will need to start a new shell for the changes to take effect.**
 
 ### For Bash
 
 ```bash
-# Generates bash completion scripts
-echo "$(_MVT_IOS_COMPLETE=bash_source mvt-ios)" > ~/.mvt-ios-complete.bash &&
-echo "$(_MVT_ANDROID_COMPLETE=bash_source mvt-android)" > ~/.mvt-android-complete.bash
+# Generate bash completion scripts
+mvt-ios completion bash > ~/.mvt-ios-complete.bash
+mvt-android completion bash > ~/.mvt-android-complete.bash
 ```
 
 Add the following to `~/.bashrc`:
 ```bash
 # source mvt completion scripts
-. ~/.mvt-ios-complete.bash && .  ~/.mvt-android-complete.bash
+[ -f ~/.mvt-ios-complete.bash ] && . ~/.mvt-ios-complete.bash
+[ -f ~/.mvt-android-complete.bash ] && . ~/.mvt-android-complete.bash
 ```
 
 ### For Zsh
 
 ```bash
-# Generates zsh completion scripts
-echo "$(_MVT_IOS_COMPLETE=zsh_source mvt-ios)" >  ~/.mvt-ios-complete.zsh &&
-echo "$(_MVT_ANDROID_COMPLETE=zsh_source mvt-android)" > ~/.mvt-android-complete.zsh
+# Generate zsh completion scripts
+mvt-ios completion zsh > ~/.mvt-ios-complete.zsh
+mvt-android completion zsh > ~/.mvt-android-complete.zsh
 ```
 
 Add the following to `~/.zshrc`:
 ```bash
 # source mvt completion scripts
-.  ~/.mvt-ios-complete.zsh  && .  ~/.mvt-android-complete.zsh
+[ -f ~/.mvt-ios-complete.zsh ] && . ~/.mvt-ios-complete.zsh
+[ -f ~/.mvt-android-complete.zsh ] && . ~/.mvt-android-complete.zsh
 ```
 
-For more information, visit the official [Click Docs](https://click.palletsprojects.com/en/stable/shell-completion/#enabling-completion).
+### For Fish
+
+```bash
+# Generate fish completion scripts
+mkdir -p ~/.config/fish/completions
+mvt-ios completion fish > ~/.config/fish/completions/mvt-ios.fish
+mvt-android completion fish > ~/.config/fish/completions/mvt-android.fish
+```
+
+Fish loads completion files from `~/.config/fish/completions` automatically.
+
+### Automatic Installation
 
+MVT can write the completion file and update the relevant shell configuration for Bash and Zsh when you pass `--install`:
+
+```bash
+mvt-ios completion bash --install
+mvt-android completion bash --install
+```
+
+Replace `bash` with `zsh` or `fish` as needed. For Fish, `--install` writes the completion file into `~/.config/fish/completions`.
+
+For more information, visit the official [Click Docs](https://click.palletsprojects.com/en/stable/shell-completion/#enabling-completion).
 

src/mvt/android/cli.py

@@ -8,6 +8,12 @@
 import click
 
 from mvt.common.cmd_check_iocs import CmdCheckIOCS
+from mvt.common.completion import (
+    SUPPORTED_SHELLS,
+    completion_instructions,
+    generate_completion_script,
+    install_completion_script,
+)
 from mvt.common.help import (
     HELP_MSG_ANDROID_BACKUP_PASSWORD,
     HELP_MSG_CHECK_ADB_REMOVED,
@@ -17,6 +23,7 @@
     HELP_MSG_CHECK_BUGREPORT,
     HELP_MSG_CHECK_IOCS,
     HELP_MSG_CHECK_INTRUSION_LOGS,
+    HELP_MSG_COMPLETION,
     HELP_MSG_DISABLE_INDICATOR_UPDATE_CHECK,
     HELP_MSG_DISABLE_UPDATE_CHECK,
     HELP_MSG_HASHES,
@@ -76,10 +83,11 @@ def cli(ctx, disable_update_check, disable_indicator_update_check):
     ctx.ensure_object(dict)
     ctx.obj["disable_version_check"] = disable_update_check
     ctx.obj["disable_indicator_check"] = disable_indicator_update_check
-    logo(
-        disable_version_check=disable_update_check,
-        disable_indicator_check=disable_indicator_update_check,
-    )
+    if ctx.invoked_subcommand != "completion":
+        logo(
+            disable_version_check=disable_update_check,
+            disable_indicator_check=disable_indicator_update_check,
+        )
 
 
 # ==============================================================================
@@ -90,6 +98,40 @@ def version():
     return
 
 
+# ==============================================================================
+# Command: completion
+# ==============================================================================
+@cli.command("completion", context_settings=CONTEXT_SETTINGS, help=HELP_MSG_COMPLETION)
+@click.argument("shell", required=False, type=click.Choice(SUPPORTED_SHELLS))
+@click.option(
+    "--install",
+    is_flag=True,
+    help="Write completion files and update shell configuration.",
+)
+@click.pass_context
+def completion(ctx, shell, install):
+    program_name = "mvt-android"
+
+    if shell is None:
+        if install:
+            raise click.UsageError("A shell is required when using --install.")
+        click.echo(completion_instructions(program_name))
+        return
+
+    root_cli = ctx.find_root().command
+
+    if install:
+        script_path = install_completion_script(root_cli, program_name, shell)
+        click.echo(f"Installed {shell} completion to {script_path}")
+        if shell in ("bash", "zsh"):
+            click.echo(f"Updated ~/.{shell}rc")
+        else:
+            click.echo("Fish loads completion files automatically.")
+        return
+
+    click.echo(generate_completion_script(root_cli, program_name, shell))
+
+
 # ==============================================================================
 # Command: check-adb (removed)
 # ==============================================================================

src/mvt/common/completion.py

@@ -0,0 +1,94 @@
+# Mobile Verification Toolkit (MVT)
+# Copyright (c) 2021-2023 The MVT Authors.
+# Use of this software is governed by the MVT License 1.1 that can be found at
+#   https://license.mvt.re/1.1/
+
+from pathlib import Path
+import shlex
+
+import click
+from click.shell_completion import get_completion_class
+
+
+SUPPORTED_SHELLS = ("bash", "zsh", "fish")
+
+
+def completion_instructions(program_name: str) -> str:
+    return f"""Shell completion for {program_name}
+
+Print a completion script:
+  {program_name} completion bash > ~/.{program_name}-complete.bash
+  {program_name} completion zsh > ~/.{program_name}-complete.zsh
+  mkdir -p ~/.config/fish/completions
+  {program_name} completion fish > ~/.config/fish/completions/{program_name}.fish
+
+Load the generated Bash script from ~/.bashrc:
+  [ -f ~/.{program_name}-complete.bash ] && . ~/.{program_name}-complete.bash
+
+Load the generated Zsh script from ~/.zshrc:
+  [ -f ~/.{program_name}-complete.zsh ] && . ~/.{program_name}-complete.zsh
+
+Fish loads completion files from ~/.config/fish/completions automatically.
+
+To write these files and update Bash/Zsh shell configuration automatically:
+  {program_name} completion bash --install
+  {program_name} completion zsh --install
+  {program_name} completion fish --install
+"""
+
+
+def generate_completion_script(cli: click.Command, program_name: str, shell: str) -> str:
+    completion_class = get_completion_class(shell)
+    if completion_class is None:
+        raise click.ClickException(f"Unsupported shell: {shell}")
+
+    complete_var = f"_{program_name.upper().replace('-', '_')}_COMPLETE"
+    return completion_class(cli, {}, program_name, complete_var).source()
+
+
+def install_completion_script(
+    cli: click.Command,
+    program_name: str,
+    shell: str,
+) -> Path:
+    script = generate_completion_script(cli, program_name, shell)
+    script_path = _completion_script_path(program_name, shell)
+    script_path.parent.mkdir(parents=True, exist_ok=True)
+    script_path.write_text(script, encoding="utf-8")
+
+    if shell in ("bash", "zsh"):
+        _install_shell_source_line(program_name, shell, script_path)
+
+    return script_path
+
+
+def _completion_script_path(program_name: str, shell: str) -> Path:
+    home = Path.home()
+
+    if shell == "fish":
+        return home / ".config" / "fish" / "completions" / f"{program_name}.fish"
+
+    return home / f".{program_name}-complete.{shell}"
+
+
+def _install_shell_source_line(program_name: str, shell: str, script_path: Path) -> None:
+    shell_config_path = Path.home() / f".{shell}rc"
+    source_line = (
+        f"[ -f {shlex.quote(str(script_path))} ] && "
+        f". {shlex.quote(str(script_path))}"
+    )
+    block = (
+        f"# MVT shell completion for {program_name}\n"
+        f"{source_line}\n"
+    )
+
+    if shell_config_path.exists():
+        shell_config = shell_config_path.read_text(encoding="utf-8")
+        if source_line in shell_config:
+            return
+    else:
+        shell_config = ""
+
+    separator = "" if not shell_config or shell_config.endswith("\n") else "\n"
+    with shell_config_path.open("a", encoding="utf-8") as handle:
+        handle.write(f"{separator}{block}")

src/mvt/common/help.py

@@ -17,6 +17,7 @@
 HELP_MSG_STIX2 = "Download public STIX2 indicators"
 HELP_MSG_DISABLE_UPDATE_CHECK = "Disable MVT version update check"
 HELP_MSG_DISABLE_INDICATOR_UPDATE_CHECK = "Disable indicators update check"
+HELP_MSG_COMPLETION = "Generate or install shell completion"
 
 # IOS Specific
 HELP_MSG_DECRYPT_BACKUP = "Decrypt an encrypted iTunes backup"

src/mvt/ios/cli.py

@@ -11,6 +11,12 @@
 from rich.prompt import Prompt
 
 from mvt.common.cmd_check_iocs import CmdCheckIOCS
+from mvt.common.completion import (
+    SUPPORTED_SHELLS,
+    completion_instructions,
+    generate_completion_script,
+    install_completion_script,
+)
 from mvt.common.logo import logo
 from mvt.common.options import MutuallyExclusiveOption
 from mvt.common.updates import IndicatorsUpdates
@@ -39,6 +45,7 @@
     HELP_MSG_CHECK_IOS_BACKUP,
     HELP_MSG_DISABLE_UPDATE_CHECK,
     HELP_MSG_DISABLE_INDICATOR_UPDATE_CHECK,
+    HELP_MSG_COMPLETION,
 )
 from .cmd_check_backup import CmdIOSCheckBackup
 from .cmd_check_fs import CmdIOSCheckFS
@@ -82,10 +89,11 @@ def cli(ctx, disable_update_check, disable_indicator_update_check):
     ctx.ensure_object(dict)
     ctx.obj["disable_version_check"] = disable_update_check
     ctx.obj["disable_indicator_check"] = disable_indicator_update_check
-    logo(
-        disable_version_check=disable_update_check,
-        disable_indicator_check=disable_indicator_update_check,
-    )
+    if ctx.invoked_subcommand != "completion":
+        logo(
+            disable_version_check=disable_update_check,
+            disable_indicator_check=disable_indicator_update_check,
+        )
 
 
 # ==============================================================================
@@ -96,6 +104,40 @@ def version():
     return
 
 
+# ==============================================================================
+# Command: completion
+# ==============================================================================
+@cli.command("completion", context_settings=CONTEXT_SETTINGS, help=HELP_MSG_COMPLETION)
+@click.argument("shell", required=False, type=click.Choice(SUPPORTED_SHELLS))
+@click.option(
+    "--install",
+    is_flag=True,
+    help="Write completion files and update shell configuration.",
+)
+@click.pass_context
+def completion(ctx, shell, install):
+    program_name = "mvt-ios"
+
+    if shell is None:
+        if install:
+            raise click.UsageError("A shell is required when using --install.")
+        click.echo(completion_instructions(program_name))
+        return
+
+    root_cli = ctx.find_root().command
+
+    if install:
+        script_path = install_completion_script(root_cli, program_name, shell)
+        click.echo(f"Installed {shell} completion to {script_path}")
+        if shell in ("bash", "zsh"):
+            click.echo(f"Updated ~/.{shell}rc")
+        else:
+            click.echo("Fish loads completion files automatically.")
+        return
+
+    click.echo(generate_completion_script(root_cli, program_name, shell))
+
+
 # ==============================================================================
 # Command: decrypt-backup
 # ==============================================================================