FIX: CLI bug fixes and minor updates#1559
Conversation
There was a problem hiding this comment.
Pull request overview
This PR improves the PyRIT CLI/shell UX by expanding list-targets support (initializers + initialization scripts), tightening argument validation for no-arg commands, fixing help lookup for hyphenated commands, and cleaning up unused CLI plumbing.
Changes:
- Add
list-targetsoption parsing inpyrit_shelland enable--initialization-scriptshandling inpyrit_scan. - Reject unexpected args for commands that should not take arguments and normalize hyphenated command names for
help. - Refactor target listing initialization flow (remove unused
initializer_namesparam) and add unit tests covering new behaviors.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
pyrit/cli/pyrit_shell.py |
Adds arg rejection for no-arg commands; adds option parsing + context override for list-targets; fixes help for hyphenated commands. |
pyrit/cli/pyrit_scan.py |
Resolves --initialization-scripts for --list-targets and surfaces missing-script errors. |
pyrit/cli/frontend_core.py |
Removes unused initializer_names param from list_targets_async and supports initialization-scripts-only population. |
pyrit/cli/_cli_args.py |
Introduces parse_list_targets_arguments for shell-mode parsing. |
pyrit/cli/_banner.py |
Updates banner text to reflect list-targets [opts]. |
tests/unit/cli/test_pyrit_shell.py |
Adds coverage for arg rejection, list-targets option flow, and hyphenated help. |
tests/unit/cli/test_pyrit_scan.py |
Adds coverage for --list-targets with initializers and initialization scripts. |
tests/unit/cli/test_frontend_core.py |
Adds coverage for parse_list_targets_arguments and scripts-only target listing initialization. |
| return result | ||
|
|
||
|
|
||
| def parse_list_targets_arguments(*, args_string: str) -> dict[str, Any]: |
There was a problem hiding this comment.
There's some opportunity to share a lot of code here. One idea claude helped with might be something like this:
from dataclasses import dataclass
from enum import Enum, auto
from typing import Callable
class _FlagKind(Enum):
MULTI = auto() # collect values until next --flag
SINGLE = auto() # consume exactly one value
POSITIONAL = auto() # first non-flag token (scenario_name)
@dataclass(frozen=True)
class _FlagSpec:
flag: str | tuple[str, ...] # e.g. "--initializers" or ("--strategies", "-s")
key: str # result dict key
kind: _FlagKind
transform: Callable | None = None # per-value transform (e.g. _parse_initializer_arg, validate_integer)
# ---- shared flag definitions (define once, reuse everywhere) ----
_INITIALIZERS = _FlagSpec("--initializers", "initializers", _FlagKind.MULTI, _parse_initializer_arg)
_INIT_SCRIPTS = _FlagSpec("--initialization-scripts", "initialization_scripts", _FlagKind.MULTI)
_LOG_LEVEL = _FlagSpec("--log-level", "log_level", _FlagKind.SINGLE, validate_log_level)
# ... etc for each flag
# ---- command specs built from shared pieces ----
_RUN_SPEC = [
_FlagSpec(None, "scenario_name", _FlagKind.POSITIONAL), # required first arg
_INITIALIZERS,
_INIT_SCRIPTS,
_FlagSpec(("--strategies", "-s"), "scenario_strategies", _FlagKind.MULTI),
_FlagSpec("--max-concurrency", "max_concurrency", _FlagKind.SINGLE,
lambda v: validate_integer(v, name="--max-concurrency", min_value=1)),
# ... remaining flags
]
_LIST_TARGETS_SPEC = [
_INITIALIZERS,
_INIT_SCRIPTS,
]
Then one generic parse function drives everything:
def _parse_args(*, args_string: str, spec: list[_FlagSpec]) -> dict[str, Any]:
"""Generic table-driven parser used by all commands."""
parts = args_string.split()
# build flag→spec lookup, init result dict, walk parts once
...
def parse_run_arguments(*, args_string: str) -> dict[str, Any]:
return _parse_args(args_string=args_string, spec=_RUN_SPEC)
def parse_list_targets_arguments(*, args_string: str) -> dict[str, Any]:
return _parse_args(args_string=args_string, spec=_LIST_TARGETS_SPEC)
| self._raise_init_error() | ||
|
|
||
| def _rebuild_context( | ||
| self, |
There was a problem hiding this comment.
Could we add most of this into FrontEndCore?
Right now we're setting a bunch of private variables and whatnot. I think FrontEndCore should be able to merge contexts. And call it from here like this:
# with_overrides is the FrontEndCore function that overrides the context values
run_context = self.context.with_overrides(
initializer_names=args["initializers"],
initialization_scripts=resolved_scripts,
log_level=args["log_level"],
)
Description
This PR makes minor updates to improve the CLI experience, including the following:
list-targetslist_targetslist-scenarios)help commanddid not work properly for hyphenated commands such aslist-targetsorlist-initializersinitializer_namesinlist_targets_asyncpyrit_shell, whendo_runanddo_list_targets(with args) create a per-commandFrontendCore, they didn't propagate the shell's startup config (config_file, database, env_files). The new context fell back to the default ~/.pyrit/.pyrit_conf, silently losing any custom config passed via --config-file. This caused per-command contexts to use the wrong database type, env files, operator/operation labels, etc. The fix implemented here is to use a_rebuild_context()helper that copies the shell's_context_kwargs(which includes config_file, database, env_files) as the base, then overrides only the per-command fields (initializer_names,initialization_scripts,log_level). Bothdo_runanddo_list_targetsnow use this helper. Also sets_silent_reinit = Trueon the derived context to suppress redundant initialization output.Tests and Documentation
Unit tests added to ensure behavior of all of the above changes