From a18e632ec851532c89e25cf7a43947c305ca756d Mon Sep 17 00:00:00 2001 From: besendorf Date: Wed, 1 Jul 2026 12:35:21 +0200 Subject: [PATCH] Add shell completion command (#817) --- README.md | 17 +++++++ docs/command_completion.md | 49 ++++++++++++++----- src/mvt/android/cli.py | 50 +++++++++++++++++-- src/mvt/common/completion.py | 94 ++++++++++++++++++++++++++++++++++++ src/mvt/common/help.py | 1 + src/mvt/ios/cli.py | 50 +++++++++++++++++-- tests/test_completion.py | 78 ++++++++++++++++++++++++++++++ 7 files changed, 318 insertions(+), 21 deletions(-) create mode 100644 src/mvt/common/completion.py create mode 100644 tests/test_completion.py diff --git a/README.md b/README.md index cb651a1..adcee7a 100644 --- a/README.md +++ b/README.md @@ -60,6 +60,23 @@ 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. Module-running `check-*` commands can load custom Python modules with `--load-module PATH` or from a folder set in `MVT_CUSTOM_MODULES`. See the [development documentation](https://docs.mvt.re/en/latest/development/) for diff --git a/docs/command_completion.md b/docs/command_completion.md index 1cd4eb7..7204df5 100644 --- a/docs/command_completion.md +++ b/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 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). - diff --git a/src/mvt/android/cli.py b/src/mvt/android/cli.py index 364f4ef..d13f86c 100644 --- a/src/mvt/android/cli.py +++ b/src/mvt/android/cli.py @@ -8,6 +8,12 @@ import logging 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 @@ from mvt.common.help import ( 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, @@ -85,10 +92,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, + ) # ============================================================================== @@ -99,6 +107,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) # ============================================================================== diff --git a/src/mvt/common/completion.py b/src/mvt/common/completion.py new file mode 100644 index 0000000..6466a6d --- /dev/null +++ b/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}") diff --git a/src/mvt/common/help.py b/src/mvt/common/help.py index 4cba131..5a2dc25 100644 --- a/src/mvt/common/help.py +++ b/src/mvt/common/help.py @@ -21,6 +21,7 @@ HELP_MSG_CHECK_IOCS = "Compare stored JSON results to provided indicators" 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" diff --git a/src/mvt/ios/cli.py b/src/mvt/ios/cli.py index d9f634c..176b036 100644 --- a/src/mvt/ios/cli.py +++ b/src/mvt/ios/cli.py @@ -11,6 +11,12 @@ import click 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 @@ -40,6 +46,7 @@ from mvt.common.help import ( HELP_MSG_CHECK_IOS_BACKUP, HELP_MSG_DISABLE_UPDATE_CHECK, HELP_MSG_DISABLE_INDICATOR_UPDATE_CHECK, + HELP_MSG_COMPLETION, ) from mvt.common.module_loader import CustomModuleLoadError, load_custom_modules from .cmd_check_backup import CmdIOSCheckBackup @@ -91,10 +98,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, + ) # ============================================================================== @@ -105,6 +113,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 # ============================================================================== diff --git a/tests/test_completion.py b/tests/test_completion.py new file mode 100644 index 0000000..48c0177 --- /dev/null +++ b/tests/test_completion.py @@ -0,0 +1,78 @@ +# 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 click.testing import CliRunner + +from mvt.android.cli import cli as android_cli +from mvt.ios.cli import cli as ios_cli + + +class TestCompletionCommand: + def test_completion_prints_instructions_by_default(self): + runner = CliRunner() + result = runner.invoke(ios_cli, ["completion"]) + + assert result.exit_code == 0 + assert "Shell completion for mvt-ios" in result.output + assert "mvt-ios completion bash > ~/.mvt-ios-complete.bash" in result.output + assert "Mobile Verification Toolkit" not in result.output + + def test_completion_prints_bash_script(self): + runner = CliRunner() + result = runner.invoke(ios_cli, ["completion", "bash"]) + + assert result.exit_code == 0 + assert "_MVT_IOS_COMPLETE=bash_complete" in result.output + assert "complete -o nosort" in result.output + assert "mvt-ios" in result.output + assert "Mobile Verification Toolkit" not in result.output + + def test_completion_prints_fish_script(self): + runner = CliRunner() + result = runner.invoke(android_cli, ["completion", "fish"]) + + assert result.exit_code == 0 + assert "_MVT_ANDROID_COMPLETE=fish_complete" in result.output + assert "complete --no-files --command mvt-android" in result.output + assert "Mobile Verification Toolkit" not in result.output + + def test_completion_install_updates_bashrc_once(self, tmp_path, monkeypatch): + monkeypatch.setenv("HOME", str(tmp_path)) + runner = CliRunner() + + result = runner.invoke(ios_cli, ["completion", "bash", "--install"]) + assert result.exit_code == 0 + + script_path = tmp_path / ".mvt-ios-complete.bash" + bashrc_path = tmp_path / ".bashrc" + assert script_path.exists() + assert "_MVT_IOS_COMPLETE=bash_complete" in script_path.read_text( + encoding="utf-8" + ) + bashrc = bashrc_path.read_text(encoding="utf-8") + assert "[ -f" in bashrc + assert ".mvt-ios-complete.bash" in bashrc + + result = runner.invoke(ios_cli, ["completion", "bash", "--install"]) + assert result.exit_code == 0 + assert bashrc_path.read_text(encoding="utf-8") == bashrc + + def test_completion_install_fish_does_not_update_shell_rc( + self, tmp_path, monkeypatch + ): + monkeypatch.setenv("HOME", str(tmp_path)) + runner = CliRunner() + + result = runner.invoke(android_cli, ["completion", "fish", "--install"]) + + assert result.exit_code == 0 + script_path = ( + tmp_path / ".config" / "fish" / "completions" / "mvt-android.fish" + ) + assert script_path.exists() + assert "_MVT_ANDROID_COMPLETE=fish_complete" in script_path.read_text( + encoding="utf-8" + ) + assert not (tmp_path / ".fishrc").exists()