diff --git a/.pyrit_conf_example b/.pyrit_conf_example index 523ccc28e1..b1b1ecf932 100644 --- a/.pyrit_conf_example +++ b/.pyrit_conf_example @@ -127,6 +127,18 @@ max_concurrent_scenario_runs: 3 # Default: false allow_custom_initializers: false +# Local Backend Server +# -------------------- +# Client settings used by pyrit_scan when connecting to or launching a backend. +# - url: Backend URL used when --server-url is omitted. +# - startup_timeout: Seconds to wait for --start-server before cleaning up the +# spawned process and returning an error. +# +# Both settings can be overridden with --server-url and --startup-timeout. +server: + url: http://localhost:8000 + startup_timeout: 120 + # Silent Mode # ----------- # If true, suppresses print statements during initialization. diff --git a/doc/getting_started/pyrit_conf.md b/doc/getting_started/pyrit_conf.md index 153c93b9bf..cfb753ed9c 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -175,6 +175,23 @@ env_files: If `true`, suppresses print statements during initialization. Useful for non-interactive environments or when embedding PyRIT in other tools. Defaults to `false`. +### `server` + +Client settings for connecting to or launching a PyRIT backend. + +| Field | Description | Default | +|---|---|---| +| `url` | Backend URL used when `--server-url` is omitted | `http://localhost:8000` | +| `startup_timeout` | Seconds `pyrit_scan --start-server` waits for a healthy backend before terminating the spawned process | `120` | + +`startup_timeout` must be a finite number greater than zero. The `--startup-timeout` CLI option overrides the configured value for an individual scanner invocation. + +```yaml +server: + url: http://localhost:8000 + startup_timeout: 120 +``` + ## Configuration Precedence PyRIT uses a 3-layer configuration precedence model. **Later layers override earlier ones:** @@ -282,6 +299,11 @@ initializers: # Suppress initialization messages silent: false + +# Backend connection and local startup settings +server: + url: http://localhost:8000 + startup_timeout: 120 ``` ## What's Next? diff --git a/pyrit/cli/_config_reader.py b/pyrit/cli/_config_reader.py index 17feabb3d7..05179fdcf8 100644 --- a/pyrit/cli/_config_reader.py +++ b/pyrit/cli/_config_reader.py @@ -4,12 +4,14 @@ """ Lightweight config reader for the PyRIT CLI thin client. -Reads only the ``server.url`` field from ``~/.pyrit/.pyrit_conf`` (and an -optional overlay file) using ``yaml.safe_load``. No heavy pyrit imports. +Reads the ``server`` settings used by the client from ``~/.pyrit/.pyrit_conf`` +(and an optional overlay file) using ``yaml.safe_load``. No heavy pyrit imports. """ from __future__ import annotations +import math +from dataclasses import dataclass from pathlib import Path from typing import Any @@ -18,13 +20,22 @@ _DEFAULT_CONFIG_FILE = _DEFAULT_CONFIG_DIR / ".pyrit_conf" DEFAULT_SERVER_URL = "http://localhost:8000" +DEFAULT_SERVER_STARTUP_TIMEOUT = 120.0 + + +@dataclass(frozen=True) +class ServerSettings: + """Client-side settings for connecting to or launching a backend server.""" + + url: str | None = None + startup_timeout: float = DEFAULT_SERVER_STARTUP_TIMEOUT class ConfigError(Exception): """ Raised when a CLI config file exists but cannot be parsed or is structurally invalid (e.g. malformed YAML, a non-mapping root, or a wrong-typed - ``server.url``). + ``server`` setting). A *missing* config file or a *missing* field is not an error -- the CLI just falls back to its defaults. This is reserved for configs the user clearly @@ -78,6 +89,25 @@ def read_server_url(*, config_file: Path | None = None) -> str | None: Returns: str | None: The server URL, or ``None`` if not configured. + Raises: + ConfigError: If a config file exists but is malformed. + """ + return read_server_settings(config_file=config_file).url + + +def read_server_settings(*, config_file: Path | None = None) -> ServerSettings: + """ + Read client-side server settings from the default config and an optional overlay. + + A later ``server`` block replaces the earlier block. Config files that omit + ``server`` leave the prior settings unchanged. + + Args: + config_file: Optional explicit config path. + + Returns: + ServerSettings: The resolved URL and startup timeout. + Raises: ConfigError: If a config file exists but is malformed. """ @@ -89,10 +119,13 @@ def read_server_url(*, config_file: Path | None = None) -> str | None: if config_file is not None and config_file.exists(): paths.append(config_file) - url: str | None = None + settings = ServerSettings() for p in paths: - url = _extract_server_url(path=p, yaml_module=yaml) or url - return url + data = _load_config_mapping(path=p, yaml_module=yaml) + if data is None or "server" not in data: + continue + settings = _extract_server_settings(data=data, path=p) + return settings def validate_client_config(*, config_file: Path | None = None) -> None: @@ -125,38 +158,40 @@ def validate_client_config(*, config_file: Path | None = None) -> None: ) -def _extract_server_url(*, path: Path, yaml_module: Any) -> str | None: +def _extract_server_settings(*, data: dict[str, Any], path: Path) -> ServerSettings: """ - Extract ``server.url`` from a single YAML file. + Extract client settings from one parsed ``server`` block. Args: - path (Path): YAML config file path. - yaml_module (Any): The imported ``yaml`` module (passed to avoid - top-level import). + data: Parsed top-level config mapping. + path: YAML config file path used in error messages. Returns: - str | None: The URL string, or ``None`` if absent. + ServerSettings: Settings represented by this config layer. Raises: - ConfigError: If the file is malformed, or ``server`` / ``server.url`` - are present but have the wrong type. + ConfigError: If ``server`` or one of its supported fields has the wrong type. """ - data = _load_config_mapping(path=path, yaml_module=yaml_module) - if data is None: - return None - server_block = data.get("server") if server_block is None: - return None + return ServerSettings() if not isinstance(server_block, dict): - raise ConfigError( - f"Config file {path}: 'server' must be a mapping with a 'url' field, got {type(server_block).__name__}." - ) + raise ConfigError(f"Config file {path}: 'server' must be a mapping, got {type(server_block).__name__}.") raw_url = server_block.get("url") - if raw_url is None: - return None - if not isinstance(raw_url, str): + if raw_url is not None and not isinstance(raw_url, str): raise ConfigError(f"Config file {path}: 'server.url' must be a string, got {type(raw_url).__name__}.") - - return raw_url.strip() or None + url: str | None = None + if isinstance(raw_url, str): + url = raw_url.strip() or None + + startup_timeout = server_block.get("startup_timeout", DEFAULT_SERVER_STARTUP_TIMEOUT) + if ( + isinstance(startup_timeout, bool) + or not isinstance(startup_timeout, int | float) + or not math.isfinite(startup_timeout) + or startup_timeout <= 0 + ): + raise ConfigError(f"Config file {path}: 'server.startup_timeout' must be a finite number greater than 0.") + + return ServerSettings(url=url, startup_timeout=float(startup_timeout)) diff --git a/pyrit/cli/_server_launcher.py b/pyrit/cli/_server_launcher.py index cdfffa8a6f..742bae791f 100644 --- a/pyrit/cli/_server_launcher.py +++ b/pyrit/cli/_server_launcher.py @@ -11,20 +11,27 @@ from __future__ import annotations import asyncio +import json import logging +import math import os +import re import signal import subprocess import sys import tempfile -from typing import TYPE_CHECKING +import time +from pathlib import Path +from urllib.parse import urlparse +from pyrit.cli._config_reader import DEFAULT_SERVER_STARTUP_TIMEOUT from pyrit.cli.api_client import PyRITApiClient -if TYPE_CHECKING: - from pathlib import Path - _logger = logging.getLogger(__name__) +_HEALTH_PROBE_TIMEOUT = 5.0 +_PROCESS_STOP_TIMEOUT = 10.0 +_STARTUP_POLL_INTERVAL = 0.5 +_PID_DIRECTORY = Path.home() / ".pyrit" / "run" # --------------------------------------------------------------------------- @@ -32,9 +39,107 @@ # --------------------------------------------------------------------------- +def parse_local_server_address(*, base_url: str) -> tuple[str, int] | None: + """ + Parse a plain loopback HTTP URL into its host and port. + + Args: + base_url: Server root URL. + + Returns: + tuple[str, int] | None: The loopback host and port, or ``None`` for a non-local or malformed URL. + """ + parsed_url = urlparse(base_url) + if ( + parsed_url.scheme != "http" + or parsed_url.hostname not in {"localhost", "127.0.0.1"} + or parsed_url.username is not None + or parsed_url.password is not None + or parsed_url.path not in {"", "/"} + or parsed_url.query + or parsed_url.fragment + ): + return None + try: + port = parsed_url.port or 80 + except ValueError: + return None + return parsed_url.hostname, port + + +def _get_pid_directory() -> Path: + return _PID_DIRECTORY + + +def _pid_file_path(*, port: int) -> Path: + return _get_pid_directory() / f"server-{port}.json" + + +def _write_pid_record(*, host: str, port: int, pid: int) -> None: + path = _pid_file_path(port=port) + temporary_path = path.with_suffix(".tmp") + try: + path.parent.mkdir(parents=True, exist_ok=True) + temporary_path.write_text( + json.dumps({"host": host, "port": port, "pid": pid}), + encoding="utf-8", + ) + temporary_path.replace(path) + except OSError as exc: + _logger.warning("Could not persist backend process state to %s: %s", path, exc) + try: + temporary_path.unlink(missing_ok=True) + except OSError: + _logger.warning("Could not remove temporary backend process state from %s", temporary_path) + + +def _remove_pid_record(*, port: int) -> None: + path = _pid_file_path(port=port) + try: + path.unlink(missing_ok=True) + except OSError as exc: + _logger.warning("Could not remove backend process state from %s: %s", path, exc) + + +def _read_recorded_pid(*, port: int) -> int | None: + path = _pid_file_path(port=port) + try: + raw_record = json.loads(path.read_text(encoding="utf-8")) + except FileNotFoundError: + return None + except (json.JSONDecodeError, OSError) as exc: + _logger.warning("Discarding invalid backend process state from %s: %s", path, exc) + _remove_pid_record(port=port) + return None + + pid = raw_record.get("pid") if isinstance(raw_record, dict) else None + recorded_port = raw_record.get("port") if isinstance(raw_record, dict) else None + if isinstance(pid, bool) or not isinstance(pid, int) or pid <= 0 or recorded_port != port: + _logger.warning("Discarding invalid backend process state from %s", path) + _remove_pid_record(port=port) + return None + return pid + + +def _parse_netstat_listener_pids(*, output: str, port: int) -> set[int]: + pids: set[int] = set() + for line in output.splitlines(): + tokens = line.split() + if len(tokens) < 5 or tokens[0].upper() != "TCP" or tokens[3].upper() != "LISTENING": + continue + try: + local_port = int(tokens[1].rsplit(":", maxsplit=1)[1]) + pid = int(tokens[-1]) + except (IndexError, ValueError): + continue + if local_port == port: + pids.add(pid) + return pids + + def _find_pid_on_port_windows(*, port: int) -> int | None: """ - Find the PID listening on *port* on Windows via ``netstat``. + Find the PID listening on *port* on Windows. Args: port: TCP port to look up. @@ -42,6 +147,25 @@ def _find_pid_on_port_windows(*, port: int) -> int | None: Returns: int | None: The PID, or ``None`` if no listener was found. """ + powershell_command = ( + f"Get-NetTCPConnection -State Listen -LocalPort {port} -ErrorAction SilentlyContinue " + "| Select-Object -ExpandProperty OwningProcess -Unique" + ) + try: + result = subprocess.run( + ["powershell.exe", "-NoProfile", "-NonInteractive", "-Command", powershell_command], + capture_output=True, + text=True, + timeout=5, + ) + pids = {int(line.strip()) for line in result.stdout.splitlines() if line.strip().isdigit()} + if len(pids) == 1: + return next(iter(pids)) + if len(pids) > 1: + return None + except (OSError, subprocess.SubprocessError): + pass + try: result = subprocess.run( ["netstat", "-ano", "-p", "TCP"], @@ -51,23 +175,8 @@ def _find_pid_on_port_windows(*, port: int) -> int | None: ) except (OSError, subprocess.SubprocessError): return None - for line in result.stdout.splitlines(): - if "LISTENING" not in line: - continue - tokens = line.split() - if len(tokens) < 5: - continue - # Match the local-address column exactly so port 80 doesn't match - # :8000 / :8080 / etc. Handles both IPv4 ("0.0.0.0:8000") and - # IPv6 ("[::]:8000") local-address formats. - local_addr = tokens[1] - if local_addr.rsplit(":", 1)[-1] != str(port): - continue - try: - return int(tokens[-1]) - except (ValueError, IndexError): - continue - return None + pids = _parse_netstat_listener_pids(output=result.stdout, port=port) + return next(iter(pids)) if len(pids) == 1 else None def _find_pid_on_port_unix(*, port: int) -> int | None: @@ -82,42 +191,144 @@ def _find_pid_on_port_unix(*, port: int) -> int | None: """ try: result = subprocess.run( - ["lsof", "-ti", f":{port}"], + ["lsof", "-t", f"-iTCP:{port}", "-sTCP:LISTEN"], + capture_output=True, + text=True, + timeout=5, + ) + except (OSError, subprocess.SubprocessError): + result = None + if result is not None: + pids = {int(line.strip()) for line in result.stdout.splitlines() if line.strip().isdigit()} + if len(pids) == 1: + return next(iter(pids)) + if len(pids) > 1: + return None + + try: + result = subprocess.run( + ["ss", "-ltnp", f"sport = :{port}"], capture_output=True, text=True, timeout=5, ) except (OSError, subprocess.SubprocessError): return None - for pid_str in result.stdout.strip().splitlines(): + pids = {int(pid) for pid in re.findall(r"pid=(\d+)", result.stdout)} + return next(iter(pids)) if len(pids) == 1 else None + + +def _find_pid_on_port(*, port: int) -> int | None: + if sys.platform == "win32": + return _find_pid_on_port_windows(port=port) + return _find_pid_on_port_unix(port=port) + + +def _resolve_server_pid(*, port: int) -> int | None: + listener_pid = _find_pid_on_port(port=port) + recorded_pid = _read_recorded_pid(port=port) + if recorded_pid is not None: + if recorded_pid == listener_pid: + return recorded_pid + _remove_pid_record(port=port) + return listener_pid + + +def _process_exists(*, pid: int) -> bool: + if sys.platform == "win32": + powershell_command = ( + f"$process = Get-Process -Id {pid} -ErrorAction SilentlyContinue; if ($null -eq $process) {{ exit 1 }}" + ) try: - return int(pid_str) - except ValueError: - continue - return None + result = subprocess.run( + ["powershell.exe", "-NoProfile", "-NonInteractive", "-Command", powershell_command], + capture_output=True, + text=True, + timeout=5, + ) + except (OSError, subprocess.SubprocessError): + return True + return result.returncode == 0 + + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + except OSError: + return False + return True + + +def _wait_for_process_exit(*, pid: int, port: int, timeout: float) -> bool: + deadline = time.monotonic() + timeout + while _process_exists(pid=pid): + if time.monotonic() >= deadline: + return False + time.sleep(0.1) + return _find_pid_on_port(port=port) is None + + +def _terminate_process_tree(*, process: subprocess.Popen[bytes]) -> bool: + if os.name == "nt": + try: + result = subprocess.run( + ["taskkill", "/PID", str(process.pid), "/T", "/F"], + capture_output=True, + text=True, + timeout=_PROCESS_STOP_TIMEOUT, + ) + except (OSError, subprocess.SubprocessError): + return False + if result.returncode != 0 and process.poll() is None: + return False + else: + try: + os.killpg(process.pid, signal.SIGTERM) + except ProcessLookupError: + pass + except OSError: + return False + + try: + process.wait(timeout=_PROCESS_STOP_TIMEOUT) + return True + except subprocess.TimeoutExpired: + if os.name == "nt": + return False + try: + os.killpg(process.pid, signal.SIGKILL) + process.wait(timeout=_PROCESS_STOP_TIMEOUT) + except (OSError, subprocess.TimeoutExpired): + return False + return True -def stop_server_on_port(*, port: int) -> bool: +def stop_server_on_port(*, port: int, shutdown_timeout: float = _PROCESS_STOP_TIMEOUT) -> bool: """ Find and terminate the process listening on *port*. Args: port: TCP port to look up. + shutdown_timeout: Seconds to wait for the process to exit. Returns: - bool: ``True`` if a process was found and signalled, ``False`` otherwise. + bool: ``True`` if the process exits and releases the port, ``False`` otherwise. """ - if sys.platform == "win32": - pid = _find_pid_on_port_windows(port=port) - else: - pid = _find_pid_on_port_unix(port=port) + pid = _resolve_server_pid(port=port) if pid is None: return False try: os.kill(pid, signal.SIGTERM) - return True - except OSError: + except OSError as exc: + _logger.warning("Could not signal backend process %d on port %d: %s", pid, port, exc) + return False + if not _wait_for_process_exit(pid=pid, port=port, timeout=shutdown_timeout): + _logger.warning("Backend process %d did not stop within %.1f seconds", pid, shutdown_timeout) return False + _remove_pid_record(port=port) + return True class ServerLauncher: @@ -132,6 +343,7 @@ class ServerLauncher: def __init__(self) -> None: self._process: subprocess.Popen[bytes] | None = None self._pid: int | None = None + self._port: int | None = None self._log_path: str | None = None # ------------------------------------------------------------------ @@ -163,7 +375,7 @@ async def start_async( port: int = 8000, config_file: Path | None = None, log_level: str | None = None, - startup_timeout: int = 30, + startup_timeout: float = DEFAULT_SERVER_STARTUP_TIMEOUT, ) -> str: """ Start ``pyrit_backend`` as a detached subprocess and wait until healthy. @@ -180,11 +392,27 @@ async def start_async( Raises: RuntimeError: If the server did not become healthy within the timeout. + ValueError: If ``startup_timeout`` is not finite and greater than zero. """ + if ( + isinstance(startup_timeout, bool) + or not isinstance(startup_timeout, int | float) + or not math.isfinite(startup_timeout) + or startup_timeout <= 0 + ): + raise ValueError("startup_timeout must be a finite number greater than 0.") + base_url = f"http://{host}:{port}" # Already running? - if await self.probe_health_async(base_url=base_url): + try: + already_running = await asyncio.wait_for( + self.probe_health_async(base_url=base_url), + timeout=min(_HEALTH_PROBE_TIMEOUT, startup_timeout), + ) + except asyncio.TimeoutError: + already_running = False + if already_running: _logger.info("Server already running at %s", base_url) return base_url @@ -228,29 +456,62 @@ async def start_async( creationflags=creation_flags, start_new_session=start_new_session, ) - self._pid = self._process.pid - _logger.info("Backend PID: %d (logs: %s)", self._pid, self._log_path) - - # Wait for health, checking if the process crashed - for _elapsed in range(startup_timeout): - await asyncio.sleep(1) - - exit_code = self._process.poll() - if exit_code is not None: - self._print_log_tail() - raise RuntimeError( - f"Server process exited with code {exit_code} during startup. See logs: {self._log_path}" - ) - - if await self.probe_health_async(base_url=base_url): - print(f"Server ready (PID {self._pid}). Logs: {self._log_path}") - return base_url - - self._print_log_tail() - raise RuntimeError( - f"pyrit_backend did not become healthy within {startup_timeout}s. " - f"Check the server logs ({self._log_path}) or start it manually with: pyrit_backend" - ) + pid = self._process.pid + self._pid = pid + self._port = port + _write_pid_record(host=host, port=port, pid=pid) + _logger.info("Backend PID: %d (logs: %s)", pid, self._log_path) + + startup_succeeded = False + cleanup_attempted = False + deadline = time.monotonic() + startup_timeout + try: + while True: + exit_code = self._process.poll() + if exit_code is not None: + self._print_log_tail() + _remove_pid_record(port=port) + self._process = None + self._pid = None + self._port = None + raise RuntimeError( + f"Server process exited with code {exit_code} during startup. See logs: {self._log_path}" + ) + + remaining = deadline - time.monotonic() + if remaining <= 0: + break + try: + healthy = await asyncio.wait_for( + self.probe_health_async(base_url=base_url), + timeout=min(_HEALTH_PROBE_TIMEOUT, remaining), + ) + except asyncio.TimeoutError: + healthy = False + if healthy: + listener_pid = _find_pid_on_port(port=port) + if listener_pid is not None: + _write_pid_record(host=host, port=port, pid=listener_pid) + print(f"Server ready (PID {self._pid}). Logs: {self._log_path}") + startup_succeeded = True + return base_url + + remaining = deadline - time.monotonic() + if remaining <= 0: + break + await asyncio.sleep(min(_STARTUP_POLL_INTERVAL, remaining)) + + self._print_log_tail() + process_stopped = self.stop() + cleanup_attempted = True + cleanup_message = "" if process_stopped else " The spawned backend process could not be stopped." + raise RuntimeError( + f"pyrit_backend did not become healthy within {startup_timeout}s. " + f"Check the server logs ({self._log_path}) or start it manually with: pyrit_backend.{cleanup_message}" + ) + finally: + if not startup_succeeded and not cleanup_attempted and self._process is not None: + self.stop() def _read_log_tail(self, *, max_lines: int = 20) -> str: """ @@ -281,18 +542,36 @@ def _print_log_tail(self) -> None: # Stop # ------------------------------------------------------------------ - def stop(self) -> None: - """Terminate the owned subprocess (if any).""" - if self._process is not None: - try: - self._process.terminate() - self._process.wait(timeout=5) - _logger.info("Stopped server (PID %d)", self._pid) - except Exception: - _logger.warning("Failed to stop server (PID %s)", self._pid, exc_info=True) - finally: - self._process = None - self._pid = None + def stop(self) -> bool: + """ + Terminate the owned subprocess, if any. + + Returns: + bool: ``True`` when no owned process remains, ``False`` when termination fails. + """ + if self._process is None: + return True + + process = self._process + if process.poll() is not None: + if self._port is not None: + _remove_pid_record(port=self._port) + self._process = None + self._pid = None + self._port = None + return True + + if not _terminate_process_tree(process=process): + _logger.warning("Failed to stop server (PID %s)", self._pid) + return False + + _logger.info("Stopped server (PID %d)", self._pid) + if self._port is not None: + _remove_pid_record(port=self._port) + self._process = None + self._pid = None + self._port = None + return True @property def pid(self) -> int | None: diff --git a/pyrit/cli/pyrit_scan.py b/pyrit/cli/pyrit_scan.py index e2ce641a43..a6760120fb 100644 --- a/pyrit/cli/pyrit_scan.py +++ b/pyrit/cli/pyrit_scan.py @@ -14,11 +14,11 @@ import argparse import asyncio import logging +import math import sys from argparse import ArgumentParser, Namespace, RawDescriptionHelpFormatter from pathlib import Path from typing import TYPE_CHECKING, Any, get_args, get_origin -from urllib.parse import urlparse from pyrit.cli._cli_args import ( ARG_HELP, @@ -126,6 +126,16 @@ def _print_cli_exception(*, exc: BaseException) -> None: """ +def _positive_finite_float(value: str) -> float: + try: + parsed = float(value) + except ValueError as exc: + raise argparse.ArgumentTypeError(f"expected a number greater than 0, got {value!r}") from exc + if not math.isfinite(parsed) or parsed <= 0: + raise argparse.ArgumentTypeError(f"expected a finite number greater than 0, got {value!r}") + return parsed + + def _build_base_parser(*, add_help: bool = True) -> ArgumentParser: """ Build the ``pyrit_scan`` argparse parser with the built-in (non-scenario) flags. @@ -160,6 +170,13 @@ def _build_base_parser(*, add_help: bool = True) -> ArgumentParser: action="store_true", help="Stop the backend server and exit", ) + server_group.add_argument( + "--startup-timeout", + type=_positive_finite_float, + default=None, + metavar="SECONDS", + help="Seconds to wait for a local backend to start (default: server.startup_timeout or 120)", + ) server_group.add_argument( "--config-file", type=Path, @@ -432,12 +449,12 @@ async def _resolve_server_url_async(*, parsed_args: Namespace) -> str | None: Returns: str | None: The server base URL, or ``None`` if unreachable. """ - from pyrit.cli._config_reader import DEFAULT_SERVER_URL, read_server_url - from pyrit.cli._server_launcher import ServerLauncher + from pyrit.cli._config_reader import DEFAULT_SERVER_URL, read_server_settings + from pyrit.cli._server_launcher import ServerLauncher, parse_local_server_address - base_url = parsed_args.server_url - if base_url is None: - base_url = read_server_url(config_file=parsed_args.config_file) or DEFAULT_SERVER_URL + server_settings = read_server_settings(config_file=parsed_args.config_file) + base_url = parsed_args.server_url or server_settings.url or DEFAULT_SERVER_URL + startup_timeout = getattr(parsed_args, "startup_timeout", None) or server_settings.startup_timeout # Probe existing server if await ServerLauncher.probe_health_async(base_url=base_url): @@ -445,16 +462,8 @@ async def _resolve_server_url_async(*, parsed_args: Namespace) -> str | None: # Auto-start if requested if parsed_args.start_server: - parsed_url = urlparse(base_url) - if ( - parsed_url.scheme != "http" - or parsed_url.hostname not in {"localhost", "127.0.0.1"} - or parsed_url.username is not None - or parsed_url.password is not None - or parsed_url.path not in {"", "/"} - or parsed_url.query - or parsed_url.fragment - ): + local_address = parse_local_server_address(base_url=base_url) + if local_address is None: print( f"Error: cannot --start-server because the configured server URL ({base_url}) " "is not a plain local HTTP URL. Use localhost or 127.0.0.1, " @@ -462,12 +471,14 @@ async def _resolve_server_url_async(*, parsed_args: Namespace) -> str | None: file=sys.stderr, ) return None + host, port = local_address launcher = ServerLauncher() try: return await launcher.start_async( - host=parsed_url.hostname, - port=parsed_url.port or 80, + host=host, + port=port, config_file=parsed_args.config_file, + startup_timeout=startup_timeout, ) except RuntimeError as exc: print(f"Error: {exc}") @@ -512,21 +523,28 @@ async def _handle_stop_server_async(*, parsed_args: Namespace) -> int: Handle ``--stop-server``: probe, then terminate the listening process. Returns: - int: Exit code (always ``0``). + int: Zero when no server is running or shutdown succeeds; one otherwise. """ - from pyrit.cli._server_launcher import ServerLauncher, stop_server_on_port + from pyrit.cli._server_launcher import ServerLauncher, parse_local_server_address, stop_server_on_port base_url = _resolve_configured_server_url(parsed_args=parsed_args) + local_address = parse_local_server_address(base_url=base_url) + if local_address is None: + print(f"Cannot stop non-local server {base_url}. Stop it on its host instead.", file=sys.stderr) + return 1 if not await ServerLauncher.probe_health_async(base_url=base_url): print(f"No server running at {base_url}.") return 0 - port = urlparse(base_url).port or 8000 - if stop_server_on_port(port=port): - print(f"Server on port {port} stopped.") - else: - print(f"Server at {base_url} is running but could not identify the process.") + _, port = local_address + if not stop_server_on_port(port=port): + print(f"Server at {base_url} is running but could not be stopped.") print(f"Find and kill it manually: look for a process listening on port {port}.") + return 1 + if await ServerLauncher.probe_health_async(base_url=base_url): + print(f"Server process exited, but a healthy backend is still responding at {base_url}.") + return 1 + print(f"Server on port {port} stopped.") return 0 diff --git a/pyrit/cli/pyrit_shell.py b/pyrit/cli/pyrit_shell.py index 408fb105f6..668af7f72b 100644 --- a/pyrit/cli/pyrit_shell.py +++ b/pyrit/cli/pyrit_shell.py @@ -192,7 +192,10 @@ def _ensure_client(self) -> bool: if not healthy and self._start_server: self._launcher = ServerLauncher() try: - base_url = self._run_async(self._launcher.start_async(config_file=self._config_file)) + base_url = self._run_async( + self._launcher.start_async(config_file=self._config_file), + timeout=None, + ) healthy = True except RuntimeError as exc: print(f"Error starting server: {exc}") @@ -584,7 +587,10 @@ def do_start_server(self, arg: str) -> None: self._launcher = ServerLauncher() try: - new_url = self._run_async(self._launcher.start_async(config_file=self._config_file)) + new_url = self._run_async( + self._launcher.start_async(config_file=self._config_file), + timeout=None, + ) self._base_url = new_url # Create new client for the started server if self._api_client is not None: @@ -599,26 +605,30 @@ def do_stop_server(self, arg: str) -> None: if arg.strip(): print(f"Error: stop-server does not accept arguments, got: {arg.strip()}") return - from pyrit.cli._server_launcher import ServerLauncher, stop_server_on_port + from pyrit.cli._server_launcher import ServerLauncher, parse_local_server_address, stop_server_on_port # If we own the launcher, use it directly if self._launcher is not None: - self._launcher.stop() + if not self._launcher.stop(): + print("Server could not be stopped.") + return print("Server stopped.") else: # Find and kill by port. Probe first so we don't SIGTERM a non-pyrit # process that happens to be listening on this port. - from urllib.parse import urlparse - base_url = self._base_url or self._resolve_base_url() - port = urlparse(base_url).port or 8000 + local_address = parse_local_server_address(base_url=base_url) + if local_address is None: + print(f"Cannot stop non-local server {base_url}. Stop it on its host instead.") + return + _, port = local_address if not self._run_async(ServerLauncher.probe_health_async(base_url=base_url)): print(f"No pyrit backend responding at {base_url}; not stopping anything.") return if stop_server_on_port(port=port): print(f"Server on port {port} stopped.") else: - print(f"No server found on port {port}.") + print(f"Server on port {port} could not be stopped.") return # Close the API client since the server is gone diff --git a/pyrit/setup/configuration_loader.py b/pyrit/setup/configuration_loader.py index 4ed9355deb..ae4350430e 100644 --- a/pyrit/setup/configuration_loader.py +++ b/pyrit/setup/configuration_loader.py @@ -8,9 +8,11 @@ from YAML files and initializes PyRIT accordingly. """ +import copy +import math import pathlib from collections.abc import Sequence -from dataclasses import dataclass, field +from dataclasses import dataclass, field, fields from typing import TYPE_CHECKING, Any, ClassVar from pyrit.common.path import DEFAULT_CONFIG_PATH @@ -58,9 +60,11 @@ class ServerConfig: Attributes: url: Base URL of the backend (e.g. ``http://localhost:8000``). + startup_timeout: Seconds to wait for a locally launched backend to become healthy. """ url: str = "http://localhost:8000" + startup_timeout: float = 120.0 @dataclass @@ -124,6 +128,20 @@ class ConfigurationLoader(YamlLoadable): allow_custom_initializers: bool = False server: dict[str, Any] | None = None extensions: dict[str, Any] = field(default_factory=dict) + _configured_fields: frozenset[str] = field(default_factory=frozenset, init=False, repr=False, compare=False) + _configured_extension_fields: frozenset[str] = field( + default_factory=frozenset, + init=False, + repr=False, + compare=False, + ) + _configured_nested_extension_fields: frozenset[str] = field( + default_factory=frozenset, + init=False, + repr=False, + compare=False, + ) + _reset_extensions: bool = field(default=False, init=False, repr=False, compare=False) def __post_init__(self) -> None: """Validate and normalize the configuration after loading.""" @@ -193,10 +211,11 @@ def _normalize_server(self) -> None: """ Normalize the optional ``server`` block to a ``ServerConfig``. - Accepts ``None`` (no server configured) or ``{"url": "..."}`` form. + Accepts ``None`` (no server configured) or a mapping with ``url`` and + ``startup_timeout`` fields. Raises: - ValueError: If ``server`` is not ``None`` or a dict, or if ``url`` is not a string. + ValueError: If ``server`` is invalid. """ if self.server is None: self._server_config: ServerConfig | None = None @@ -206,7 +225,18 @@ def _normalize_server(self) -> None: url = self.server.get("url", "http://localhost:8000") if not isinstance(url, str): raise ValueError(f"Server 'url' must be a string. Got: {type(url).__name__}") - self._server_config = ServerConfig(url=url.rstrip("/")) + startup_timeout = self.server.get("startup_timeout", 120.0) + if ( + isinstance(startup_timeout, bool) + or not isinstance(startup_timeout, int | float) + or not math.isfinite(startup_timeout) + or startup_timeout <= 0 + ): + raise ValueError("Server 'startup_timeout' must be a finite number greater than 0.") + self._server_config = ServerConfig( + url=url.rstrip("/"), + startup_timeout=float(startup_timeout), + ) return raise ValueError(f"Server entry must be a dict, got: {type(self.server).__name__}") @@ -236,20 +266,37 @@ def from_dict(cls, data: dict[str, Any]) -> "ConfigurationLoader": "The 'scenario' configuration block is no longer supported. " "Pass the scenario name positionally and its parameters as CLI flags." ) + configured_fields = frozenset(data) # Filter out None values only - empty lists are meaningful ("load nothing") filtered_data = {k: v for k, v in data.items() if v is not None} - known_fields = set(cls.__dataclass_fields__.keys()) + known_fields = {config_field.name for config_field in fields(cls) if config_field.init} + configured_extension_fields = frozenset(data.keys() - known_fields) known_data = {k: v for k, v in filtered_data.items() if k in known_fields and k != "extensions"} extra_data = {k: v for k, v in filtered_data.items() if k not in known_fields} + configured_nested_extension_fields: frozenset[str] = frozenset() + reset_extensions = False + if "extensions" in data: + raw_extensions = data["extensions"] + reset_extensions = raw_extensions is None or raw_extensions == {} + if isinstance(raw_extensions, dict): + if not all(isinstance(key, str) for key in raw_extensions): + raise ValueError("ConfigurationLoader.extensions keys must be strings.") + configured_nested_extension_fields = frozenset(key for key in raw_extensions if isinstance(key, str)) if "extensions" in filtered_data: extensions = filtered_data["extensions"] if not isinstance(extensions, dict): raise ValueError(f"ConfigurationLoader.extensions must be a dict. Got: {type(extensions).__name__}") extra_data = {**extra_data, **extensions} - return cls(**known_data, extensions=extra_data) + config = cls(**known_data, extensions=extra_data) + config._configured_fields = configured_fields + config._configured_extension_fields = configured_extension_fields + config._configured_nested_extension_fields = configured_nested_extension_fields + config._reset_extensions = reset_extensions + return config - @staticmethod + @classmethod def load_with_overrides( + cls, config_file: pathlib.Path | None = None, *, memory_db_type: str | None = None, @@ -266,10 +313,6 @@ def load_with_overrides( 2. Explicit config_file argument if provided 3. Individual override arguments (non-None values take precedence) - This is a staticmethod (not classmethod) because it's a pure factory function - that doesn't need access to class state and can be reused by multiple interfaces - (CLI, shell, programmatic API). - Args: config_file: Optional path to a YAML-formatted configuration file. memory_db_type: Override for database type (in_memory, sqlite, azure_sql). @@ -290,37 +333,19 @@ def load_with_overrides( logger = logging.getLogger(__name__) - # Start with defaults - None means "use defaults", [] means "load nothing" - config_data: dict[str, Any] = { - "memory_db_type": "sqlite", - "initializers": [], - "initialization_scripts": None, # None = use defaults - "env_files": None, # None = use defaults - "env_akv_ref": None, - "silent": False, - } + init_fields = {config_field.name for config_field in fields(cls) if config_field.init} + + def to_init_data(config: ConfigurationLoader) -> dict[str, Any]: + return {name: copy.deepcopy(getattr(config, name)) for name in init_fields} # 1. Try loading default config file if it exists + config_data = to_init_data(cls()) default_config_path = DEFAULT_CONFIG_PATH if default_config_path.exists(): try: logger.info(f"Loading default configuration file: {default_config_path}") print(f"Loading default configuration file: {default_config_path}") - default_config = ConfigurationLoader.from_yaml_file(default_config_path) - config_data["memory_db_type"] = default_config.memory_db_type - config_data["initializers"] = [ - {"name": ic.name, "args": ic.args} if ic.args else ic.name - for ic in default_config._initializer_configs - ] - # Preserve None vs [] distinction from config file - config_data["initialization_scripts"] = default_config.initialization_scripts - config_data["env_files"] = default_config.env_files - config_data["env_akv_ref"] = default_config.env_akv_ref - config_data["silent"] = default_config.silent - if default_config.operator: - config_data["operator"] = default_config.operator - if default_config.operation: - config_data["operation"] = default_config.operation + config_data = to_init_data(cls.from_yaml_file(default_config_path)) except _RemovedConfigurationOptionError: raise except Exception as e: @@ -332,21 +357,19 @@ def load_with_overrides( raise FileNotFoundError(f"Configuration file not found: {config_file}") logger.info(f"Loading configuration file: {config_file}") print(f"Loading configuration file: {config_file}") - explicit_config = ConfigurationLoader.from_yaml_file(config_file) - config_data["memory_db_type"] = explicit_config.memory_db_type - config_data["initializers"] = [ - {"name": ic.name, "args": ic.args} if ic.args else ic.name - for ic in explicit_config._initializer_configs - ] - # Preserve None vs [] distinction from config file - config_data["initialization_scripts"] = explicit_config.initialization_scripts - config_data["env_files"] = explicit_config.env_files - config_data["env_akv_ref"] = explicit_config.env_akv_ref - config_data["silent"] = explicit_config.silent - if explicit_config.operator: - config_data["operator"] = explicit_config.operator - if explicit_config.operation: - config_data["operation"] = explicit_config.operation + explicit_config = cls.from_yaml_file(config_file) + explicit_data = to_init_data(explicit_config) + for field_name in explicit_config._configured_fields & (init_fields - {"extensions"}): + config_data[field_name] = explicit_data[field_name] + if explicit_config._reset_extensions: + config_data["extensions"] = {} + for field_name in explicit_config._configured_extension_fields: + if field_name in explicit_config.extensions: + config_data["extensions"][field_name] = explicit_config.extensions[field_name] + else: + config_data["extensions"].pop(field_name, None) + for field_name in explicit_config._configured_nested_extension_fields: + config_data["extensions"][field_name] = explicit_config.extensions[field_name] # 3. Apply overrides (non-None values take precedence) # Convert Sequence to list to match dataclass field types if memory_db_type is not None: @@ -370,7 +393,7 @@ def load_with_overrides( if env_akv_ref is not None: config_data["env_akv_ref"] = list(env_akv_ref) - return ConfigurationLoader.from_dict(config_data) + return cls.from_dict(config_data) @classmethod def get_default_config_path(cls) -> pathlib.Path: diff --git a/tests/unit/cli/test_config_reader.py b/tests/unit/cli/test_config_reader.py index 52776713b6..4c0cb52e98 100644 --- a/tests/unit/cli/test_config_reader.py +++ b/tests/unit/cli/test_config_reader.py @@ -11,8 +11,11 @@ from pyrit.cli import _config_reader from pyrit.cli._config_reader import ( + DEFAULT_SERVER_STARTUP_TIMEOUT, DEFAULT_SERVER_URL, ConfigError, + ServerSettings, + read_server_settings, read_server_url, validate_client_config, ) @@ -20,6 +23,7 @@ def test_default_server_url_constant(): assert DEFAULT_SERVER_URL == "http://localhost:8000" + assert DEFAULT_SERVER_STARTUP_TIMEOUT == 120.0 def test_read_server_url_returns_none_when_no_files(tmp_path): @@ -107,6 +111,38 @@ def test_read_server_url_empty_string_treated_as_missing(tmp_path): assert read_server_url(config_file=empty) is None +def test_read_server_settings_returns_defaults_when_no_files(tmp_path): + nonexistent = tmp_path / "missing.yaml" + with patch.object(_config_reader, "_DEFAULT_CONFIG_FILE", tmp_path / "missing_default.yaml"): + assert read_server_settings(config_file=nonexistent) == ServerSettings( + url=None, + startup_timeout=120.0, + ) + + +def test_read_server_settings_overlay_overrides_startup_timeout(tmp_path): + default = tmp_path / "default.yaml" + default.write_text("server:\n url: http://default:8000\n startup_timeout: 180\n", encoding="utf-8") + overlay = tmp_path / "overlay.yaml" + overlay.write_text("server:\n url: http://overlay:9000\n startup_timeout: 45\n", encoding="utf-8") + + with patch.object(_config_reader, "_DEFAULT_CONFIG_FILE", default): + assert read_server_settings(config_file=overlay) == ServerSettings( + url="http://overlay:9000", + startup_timeout=45.0, + ) + + +@pytest.mark.parametrize("startup_timeout", ["slow", True, 0, -1, float("inf")]) +def test_read_server_settings_rejects_invalid_startup_timeout(tmp_path, startup_timeout): + bad = tmp_path / "bad.yaml" + bad.write_text(f"server:\n startup_timeout: {startup_timeout}\n", encoding="utf-8") + + with patch.object(_config_reader, "_DEFAULT_CONFIG_FILE", tmp_path / "missing.yaml"): + with pytest.raises(ConfigError, match="server.startup_timeout"): + read_server_settings(config_file=bad) + + def test_validate_client_config_rejects_removed_scenario_block(tmp_path): cfg = tmp_path / "conf.yaml" cfg.write_text("scenario:\n name: test\n", encoding="utf-8") diff --git a/tests/unit/cli/test_pyrit_scan.py b/tests/unit/cli/test_pyrit_scan.py index 53d827b02c..366a95c387 100644 --- a/tests/unit/cli/test_pyrit_scan.py +++ b/tests/unit/cli/test_pyrit_scan.py @@ -179,6 +179,15 @@ def test_parse_args_with_stop_server(self): args = pyrit_scan.parse_args(["--stop-server"]) assert args.stop_server is True + def test_parse_args_with_startup_timeout(self): + args = pyrit_scan.parse_args(["--start-server", "--startup-timeout", "45.5"]) + assert args.startup_timeout == 45.5 + + @pytest.mark.parametrize("value", ["0", "-1", "inf", "nan", "slow"]) + def test_parse_args_rejects_invalid_startup_timeout(self, value): + with pytest.raises(SystemExit): + pyrit_scan.parse_args(["--start-server", "--startup-timeout", value]) + def test_main_with_invalid_args(self): result = pyrit_scan.main(["--invalid-flag"]) assert result == 2 @@ -519,7 +528,8 @@ class TestStopServerOnPort: @patch("sys.platform", "win32") @patch("subprocess.run") @patch("os.kill") - def test_stop_on_windows_finds_pid_via_netstat(self, mock_kill, mock_run): + @patch("pyrit.cli._server_launcher._wait_for_process_exit", return_value=True) + def test_stop_on_windows_finds_pid_via_netstat(self, _mock_wait, mock_kill, mock_run): from pyrit.cli import _server_launcher mock_run.return_value = MagicMock( @@ -547,7 +557,8 @@ def test_stop_on_windows_does_not_match_port_substring(self, mock_kill, mock_run @patch("sys.platform", "win32") @patch("subprocess.run") @patch("os.kill") - def test_stop_on_windows_matches_ipv6_local_address(self, mock_kill, mock_run): + @patch("pyrit.cli._server_launcher._wait_for_process_exit", return_value=True) + def test_stop_on_windows_matches_ipv6_local_address(self, _mock_wait, mock_kill, mock_run): from pyrit.cli import _server_launcher mock_run.return_value = MagicMock( @@ -559,7 +570,8 @@ def test_stop_on_windows_matches_ipv6_local_address(self, mock_kill, mock_run): @patch("sys.platform", "linux") @patch("subprocess.run") @patch("os.kill") - def test_stop_on_unix_finds_pid_via_lsof(self, mock_kill, mock_run): + @patch("pyrit.cli._server_launcher._wait_for_process_exit", return_value=True) + def test_stop_on_unix_finds_pid_via_lsof(self, _mock_wait, mock_kill, mock_run): from pyrit.cli import _server_launcher mock_run.return_value = MagicMock(stdout="5678\n") @@ -712,9 +724,15 @@ async def test_uses_cli_flag_when_provided(self): start_server=False, config_file=None, ) - with patch( - "pyrit.cli._server_launcher.ServerLauncher.probe_health_async", - new=AsyncMock(return_value=True), + with ( + patch( + "pyrit.cli._config_reader.read_server_settings", + return_value=pyrit_scan_config_reader.ServerSettings(), + ), + patch( + "pyrit.cli._server_launcher.ServerLauncher.probe_health_async", + new=AsyncMock(return_value=True), + ), ): result = await pyrit_scan._resolve_server_url_async(parsed_args=parsed) assert result == "http://override:7000" @@ -723,8 +741,8 @@ async def test_returns_none_when_unhealthy_and_no_start_server(self): parsed = Namespace(server_url=None, start_server=False, config_file=None) with ( patch( - "pyrit.cli._config_reader.read_server_url", - return_value=None, + "pyrit.cli._config_reader.read_server_settings", + return_value=pyrit_scan_config_reader.ServerSettings(), ), patch( "pyrit.cli._server_launcher.ServerLauncher.probe_health_async", @@ -735,23 +753,36 @@ async def test_returns_none_when_unhealthy_and_no_start_server(self): async def test_auto_starts_server_when_requested(self): parsed = Namespace(server_url=None, start_server=True, config_file=None) + start_async_mock = AsyncMock(return_value="http://localhost:8000") with ( - patch("pyrit.cli._config_reader.read_server_url", return_value=None), + patch( + "pyrit.cli._config_reader.read_server_settings", + return_value=pyrit_scan_config_reader.ServerSettings(), + ), patch( "pyrit.cli._server_launcher.ServerLauncher.probe_health_async", new=AsyncMock(return_value=False), ), patch( "pyrit.cli._server_launcher.ServerLauncher.start_async", - new=AsyncMock(return_value="http://localhost:8000"), + new=start_async_mock, ), ): assert await pyrit_scan._resolve_server_url_async(parsed_args=parsed) == "http://localhost:8000" + start_async_mock.assert_awaited_once_with( + host="localhost", + port=8000, + config_file=None, + startup_timeout=120.0, + ) async def test_returns_none_when_start_server_raises(self, capsys): parsed = Namespace(server_url=None, start_server=True, config_file=None) with ( - patch("pyrit.cli._config_reader.read_server_url", return_value=None), + patch( + "pyrit.cli._config_reader.read_server_settings", + return_value=pyrit_scan_config_reader.ServerSettings(), + ), patch( "pyrit.cli._server_launcher.ServerLauncher.probe_health_async", new=AsyncMock(return_value=False), @@ -768,6 +799,10 @@ async def test_start_server_refuses_remote_url(self, capsys): parsed = Namespace(server_url="http://other:9999", start_server=True, config_file=None) start_async_mock = AsyncMock() with ( + patch( + "pyrit.cli._config_reader.read_server_settings", + return_value=pyrit_scan_config_reader.ServerSettings(), + ), patch( "pyrit.cli._server_launcher.ServerLauncher.probe_health_async", new=AsyncMock(return_value=False), @@ -788,6 +823,10 @@ async def test_start_server_uses_custom_local_port(self): parsed = Namespace(server_url="http://127.0.0.1:8765", start_server=True, config_file=None) start_async_mock = AsyncMock(return_value="http://127.0.0.1:8765") with ( + patch( + "pyrit.cli._config_reader.read_server_settings", + return_value=pyrit_scan_config_reader.ServerSettings(), + ), patch( "pyrit.cli._server_launcher.ServerLauncher.probe_health_async", new=AsyncMock(return_value=False), @@ -799,7 +838,38 @@ async def test_start_server_uses_custom_local_port(self): ): result = await pyrit_scan._resolve_server_url_async(parsed_args=parsed) assert result == "http://127.0.0.1:8765" - start_async_mock.assert_awaited_once_with(host="127.0.0.1", port=8765, config_file=None) + start_async_mock.assert_awaited_once_with( + host="127.0.0.1", + port=8765, + config_file=None, + startup_timeout=120.0, + ) + + async def test_start_server_cli_timeout_overrides_config(self): + parsed = Namespace( + server_url=None, + start_server=True, + config_file=None, + startup_timeout=15.0, + ) + start_async_mock = AsyncMock(return_value="http://localhost:8000") + with ( + patch( + "pyrit.cli._config_reader.read_server_settings", + return_value=pyrit_scan_config_reader.ServerSettings(startup_timeout=90.0), + ), + patch( + "pyrit.cli._server_launcher.ServerLauncher.probe_health_async", + new=AsyncMock(return_value=False), + ), + patch( + "pyrit.cli._server_launcher.ServerLauncher.start_async", + new=start_async_mock, + ), + ): + await pyrit_scan._resolve_server_url_async(parsed_args=parsed) + + assert start_async_mock.await_args.kwargs["startup_timeout"] == 15.0 async def test_resolution_order_cli_beats_config_beats_default(self): """CLI flag > config-file value > built-in default.""" @@ -807,8 +877,8 @@ async def test_resolution_order_cli_beats_config_beats_default(self): parsed = Namespace(server_url="http://cli:1111", start_server=False, config_file=None) with ( patch( - "pyrit.cli._config_reader.read_server_url", - return_value="http://cfg:2222", + "pyrit.cli._config_reader.read_server_settings", + return_value=pyrit_scan_config_reader.ServerSettings(url="http://cfg:2222"), ), patch( "pyrit.cli._server_launcher.ServerLauncher.probe_health_async", @@ -821,8 +891,8 @@ async def test_resolution_order_cli_beats_config_beats_default(self): parsed = Namespace(server_url=None, start_server=False, config_file=None) with ( patch( - "pyrit.cli._config_reader.read_server_url", - return_value="http://cfg:2222", + "pyrit.cli._config_reader.read_server_settings", + return_value=pyrit_scan_config_reader.ServerSettings(url="http://cfg:2222"), ), patch( "pyrit.cli._server_launcher.ServerLauncher.probe_health_async", @@ -836,7 +906,10 @@ async def test_resolution_order_cli_beats_config_beats_default(self): parsed = Namespace(server_url=None, start_server=False, config_file=None) with ( - patch("pyrit.cli._config_reader.read_server_url", return_value=None), + patch( + "pyrit.cli._config_reader.read_server_settings", + return_value=pyrit_scan_config_reader.ServerSettings(), + ), patch( "pyrit.cli._server_launcher.ServerLauncher.probe_health_async", new=AsyncMock(return_value=True), @@ -1051,7 +1124,8 @@ def test_main_start_server_only_prints_url_and_returns_zero(self, mock_client_cl return_value=True, ) @patch("pyrit.cli._server_launcher.stop_server_on_port", return_value=True) - def test_main_stop_server_kills_process_and_returns_zero(self, _stop_mock, _mock_probe, capsys): + def test_main_stop_server_kills_process_and_returns_zero(self, _stop_mock, mock_probe, capsys): + mock_probe.side_effect = [True, False] result = pyrit_scan.main(["--stop-server"]) assert result == 0 assert "stopped" in capsys.readouterr().out @@ -1064,9 +1138,32 @@ def test_main_stop_server_kills_process_and_returns_zero(self, _stop_mock, _mock @patch("pyrit.cli._server_launcher.stop_server_on_port", return_value=False) def test_main_stop_server_when_process_cannot_be_identified(self, _stop_mock, _mock_probe, capsys): result = pyrit_scan.main(["--stop-server"]) - assert result == 0 + assert result == 1 out = capsys.readouterr().out - assert "could not identify" in out + assert "could not be stopped" in out + + @patch( + "pyrit.cli._server_launcher.ServerLauncher.probe_health_async", + new_callable=AsyncMock, + return_value=True, + ) + @patch("pyrit.cli._server_launcher.stop_server_on_port", return_value=True) + def test_main_stop_server_fails_when_backend_remains_healthy(self, _stop_mock, _mock_probe, capsys): + result = pyrit_scan.main(["--stop-server"]) + assert result == 1 + assert "still responding" in capsys.readouterr().out + + @patch( + "pyrit.cli._server_launcher.ServerLauncher.probe_health_async", + new_callable=AsyncMock, + ) + @patch("pyrit.cli._server_launcher.stop_server_on_port") + def test_main_stop_server_refuses_remote_url(self, stop_mock, probe_mock, capsys): + result = pyrit_scan.main(["--stop-server", "--server-url", "http://remote:8000"]) + assert result == 1 + stop_mock.assert_not_called() + probe_mock.assert_not_called() + assert "Cannot stop non-local server" in capsys.readouterr().err @patch( "pyrit.cli._server_launcher.ServerLauncher.probe_health_async", diff --git a/tests/unit/cli/test_pyrit_shell.py b/tests/unit/cli/test_pyrit_shell.py index 5063c1eb63..95a993c763 100644 --- a/tests/unit/cli/test_pyrit_shell.py +++ b/tests/unit/cli/test_pyrit_shell.py @@ -218,7 +218,32 @@ def test_do_stop_server_no_launcher(self, shell, capsys): ): s.do_stop_server("") captured = capsys.readouterr() - assert "No server found" in captured.out + assert "could not be stopped" in captured.out + + def test_do_stop_server_reports_owned_launcher_failure(self, shell, capsys): + s, _ = shell + s._launcher = MagicMock() + s._launcher.stop.return_value = False + + s.do_stop_server("") + + assert "could not be stopped" in capsys.readouterr().out + + def test_do_stop_server_refuses_remote_url(self, shell, capsys): + s, _ = shell + s._base_url = "http://remote:8765" + with ( + patch( + "pyrit.cli._server_launcher.ServerLauncher.probe_health_async", + new_callable=AsyncMock, + ) as probe_mock, + patch("pyrit.cli._server_launcher.stop_server_on_port") as stop_mock, + ): + s.do_stop_server("") + + probe_mock.assert_not_called() + stop_mock.assert_not_called() + assert "Cannot stop non-local server" in capsys.readouterr().out def test_ensure_client_already_connected(self, shell): s, _ = shell diff --git a/tests/unit/cli/test_server_launcher.py b/tests/unit/cli/test_server_launcher.py index 7c5201d0b2..bd59b6ce51 100644 --- a/tests/unit/cli/test_server_launcher.py +++ b/tests/unit/cli/test_server_launcher.py @@ -5,14 +5,26 @@ Unit tests for pyrit.cli._server_launcher.ServerLauncher. """ +import asyncio +import inspect +import json +import signal import subprocess +import sys from pathlib import Path -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, call, patch import pytest +from pyrit.cli import _server_launcher from pyrit.cli._server_launcher import ServerLauncher + +@pytest.fixture(autouse=True) +def isolate_pid_directory(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(_server_launcher, "_PID_DIRECTORY", tmp_path / "run") + + # --------------------------------------------------------------------------- # probe_health_async # --------------------------------------------------------------------------- @@ -66,6 +78,7 @@ async def test_start_async_spawns_subprocess_and_waits_for_health(): with ( patch.object(ServerLauncher, "probe_health_async", new=probe), patch("subprocess.Popen", return_value=fake_proc) as popen_mock, + patch.object(_server_launcher, "_find_pid_on_port", return_value=9876), patch("asyncio.sleep", new=AsyncMock(return_value=None)), ): url = await launcher.start_async( @@ -84,6 +97,13 @@ async def test_start_async_spawns_subprocess_and_waits_for_health(): assert "/tmp/foo.yaml" in cmd or "\\tmp\\foo.yaml" in cmd assert "--log-level" in cmd assert "INFO" in cmd + record = json.loads(_server_launcher._pid_file_path(port=8001).read_text(encoding="utf-8")) + assert record == {"host": "localhost", "port": 8001, "pid": 9876} + + +def test_start_async_defaults_to_longer_startup_timeout(): + parameter = inspect.signature(ServerLauncher.start_async).parameters["startup_timeout"] + assert parameter.default == 120.0 async def test_start_async_redirects_child_stdio_to_log_file(): @@ -98,6 +118,7 @@ async def test_start_async_redirects_child_stdio_to_log_file(): with ( patch.object(ServerLauncher, "probe_health_async", new=probe), patch("subprocess.Popen", return_value=fake_proc) as popen_mock, + patch.object(_server_launcher, "_find_pid_on_port", return_value=4321), patch("asyncio.sleep", new=AsyncMock(return_value=None)), ): await launcher.start_async(host="localhost", port=8001, startup_timeout=5) @@ -135,10 +156,72 @@ async def test_start_async_raises_when_timeout_exhausted(): with ( patch.object(ServerLauncher, "probe_health_async", new=probe), patch("subprocess.Popen", return_value=fake_proc), - patch("asyncio.sleep", new=AsyncMock(return_value=None)), + patch.object(_server_launcher, "_terminate_process_tree", return_value=True) as stop_tree_mock, ): with pytest.raises(RuntimeError, match="did not become healthy"): - await launcher.start_async(host="localhost", port=8000, startup_timeout=2) + await launcher.start_async(host="localhost", port=8000, startup_timeout=0.01) + stop_tree_mock.assert_called_once_with(process=fake_proc) + assert launcher.pid is None + assert not _server_launcher._pid_file_path(port=8000).exists() + + +async def test_start_async_cancellation_cleans_up_spawned_process(): + launcher = ServerLauncher() + fake_proc = MagicMock() + fake_proc.pid = 99 + fake_proc.poll.return_value = None + probe_started = asyncio.Event() + probe_calls = 0 + + async def probe_health_async(*, base_url: str) -> bool: + nonlocal probe_calls + probe_calls += 1 + if probe_calls == 1: + return False + probe_started.set() + await asyncio.Event().wait() + return False + + with ( + patch.object(ServerLauncher, "probe_health_async", new=AsyncMock(side_effect=probe_health_async)), + patch("subprocess.Popen", return_value=fake_proc), + patch.object(_server_launcher, "_terminate_process_tree", return_value=True) as stop_tree_mock, + ): + startup_task = asyncio.create_task(launcher.start_async(host="localhost", port=8000, startup_timeout=60)) + await probe_started.wait() + startup_task.cancel() + with pytest.raises(asyncio.CancelledError): + await startup_task + + stop_tree_mock.assert_called_once_with(process=fake_proc) + assert launcher.pid is None + assert not _server_launcher._pid_file_path(port=8000).exists() + + +async def test_start_async_timeout_caps_hanging_health_probe(): + launcher = ServerLauncher() + fake_proc = MagicMock() + fake_proc.pid = 99 + fake_proc.poll.return_value = None + probe_calls = 0 + + async def probe_health_async(*, base_url: str) -> bool: + nonlocal probe_calls + probe_calls += 1 + if probe_calls == 1: + return False + await asyncio.Event().wait() + return False + + with ( + patch.object(ServerLauncher, "probe_health_async", new=AsyncMock(side_effect=probe_health_async)), + patch("subprocess.Popen", return_value=fake_proc), + patch.object(_server_launcher, "_terminate_process_tree", return_value=True) as stop_tree_mock, + ): + with pytest.raises(RuntimeError, match="did not become healthy"): + await launcher.start_async(host="localhost", port=8000, startup_timeout=0.01) + + stop_tree_mock.assert_called_once_with(process=fake_proc) async def test_start_async_prints_log_path_on_success(capsys): @@ -151,6 +234,7 @@ async def test_start_async_prints_log_path_on_success(capsys): with ( patch.object(ServerLauncher, "probe_health_async", new=probe), patch("subprocess.Popen", return_value=fake_proc), + patch.object(_server_launcher, "_find_pid_on_port", return_value=4321), patch("asyncio.sleep", new=AsyncMock(return_value=None)), ): await launcher.start_async(host="localhost", port=8001, startup_timeout=5) @@ -192,11 +276,11 @@ async def test_start_async_echoes_log_tail_on_timeout(capsys): with ( patch.object(ServerLauncher, "probe_health_async", new=probe), patch("subprocess.Popen", return_value=fake_proc), - patch("asyncio.sleep", new=AsyncMock(return_value=None)), + patch.object(_server_launcher, "_terminate_process_tree", return_value=True), patch.object(ServerLauncher, "_read_log_tail", return_value="Traceback: boom"), ): with pytest.raises(RuntimeError, match="did not become healthy"): - await launcher.start_async(host="localhost", port=8000, startup_timeout=2) + await launcher.start_async(host="localhost", port=8000, startup_timeout=0.01) err = capsys.readouterr().err assert "Traceback: boom" in err @@ -238,32 +322,130 @@ def test_stop_terminates_process(): launcher = ServerLauncher() fake_proc = MagicMock() fake_proc.pid = 12345 + fake_proc.poll.return_value = None launcher._process = fake_proc launcher._pid = 12345 + launcher._port = 8000 - launcher.stop() + with patch.object(_server_launcher, "_terminate_process_tree", return_value=True) as stop_tree_mock: + assert launcher.stop() is True - fake_proc.terminate.assert_called_once() - fake_proc.wait.assert_called_once_with(timeout=5) + stop_tree_mock.assert_called_once_with(process=fake_proc) assert launcher.pid is None assert launcher._process is None -def test_stop_swallows_termination_errors(): +def test_stop_reports_termination_errors_and_retains_process(): launcher = ServerLauncher() fake_proc = MagicMock() fake_proc.pid = 12345 - fake_proc.terminate.side_effect = OSError("permission denied") + fake_proc.poll.return_value = None launcher._process = fake_proc launcher._pid = 12345 + launcher._port = 8000 - # Should not raise. - launcher.stop() - assert launcher._process is None - assert launcher.pid is None + with patch.object(_server_launcher, "_terminate_process_tree", return_value=False): + assert launcher.stop() is False + assert launcher._process is fake_proc + assert launcher.pid == 12345 + + +def test_terminate_process_tree_kills_unresponsive_unix_group(): + fake_proc = MagicMock() + fake_proc.pid = 12345 + fake_proc.poll.return_value = None + fake_proc.wait.side_effect = [subprocess.TimeoutExpired(cmd="backend", timeout=10), 0] + + with ( + patch("os.name", "posix"), + patch("os.killpg", create=True) as kill_group_mock, + patch("signal.SIGKILL", 9, create=True), + ): + assert _server_launcher._terminate_process_tree(process=fake_proc) is True + + assert kill_group_mock.call_args_list == [ + call(12345, signal.SIGTERM), + call(12345, 9), + ] + + +def test_terminate_process_tree_uses_windows_pid_tree(): + fake_proc = MagicMock() + fake_proc.pid = 12345 + result = MagicMock(returncode=0) + + with patch("os.name", "nt"), patch("subprocess.run", return_value=result) as run_mock: + assert _server_launcher._terminate_process_tree(process=fake_proc) is True + + assert run_mock.call_args.args[0] == ["taskkill", "/PID", "12345", "/T", "/F"] def test_stop_is_noop_when_no_process(): launcher = ServerLauncher() - launcher.stop() # Should not raise. + assert launcher.stop() is True assert launcher.pid is None + + +# --------------------------------------------------------------------------- +# persisted PID resolution +# --------------------------------------------------------------------------- + + +def test_stop_server_uses_valid_recorded_pid(): + _server_launcher._write_pid_record(host="localhost", port=8000, pid=1234) + + with ( + patch.object(_server_launcher, "_find_pid_on_port", return_value=1234), + patch.object(_server_launcher, "_wait_for_process_exit", return_value=True), + patch("os.kill") as kill_mock, + ): + assert _server_launcher.stop_server_on_port(port=8000) is True + + kill_mock.assert_called_once_with(1234, signal.SIGTERM) + assert not _server_launcher._pid_file_path(port=8000).exists() + + +def test_stop_server_discards_stale_record_and_uses_listener_pid(): + _server_launcher._write_pid_record(host="localhost", port=8000, pid=1234) + + with ( + patch.object(_server_launcher, "_find_pid_on_port", return_value=5678), + patch.object(_server_launcher, "_wait_for_process_exit", return_value=True), + patch("os.kill") as kill_mock, + ): + assert _server_launcher.stop_server_on_port(port=8000) is True + + kill_mock.assert_called_once_with(5678, signal.SIGTERM) + assert not _server_launcher._pid_file_path(port=8000).exists() + + +def test_windows_pid_lookup_prefers_get_net_tcp_connection(): + result = MagicMock(stdout="116284\n") + with patch("sys.platform", "win32"), patch("subprocess.run", return_value=result) as run_mock: + assert _server_launcher._find_pid_on_port(port=8765) == 116284 + + assert run_mock.call_args.args[0][0] == "powershell.exe" + + +def test_windows_process_exists_uses_non_destructive_lookup(): + result = MagicMock(returncode=0) + with ( + patch("sys.platform", "win32"), + patch("subprocess.run", return_value=result) as run_mock, + patch("os.kill") as kill_mock, + ): + assert _server_launcher._process_exists(pid=1234) is True + + kill_mock.assert_not_called() + assert "Get-Process" in run_mock.call_args.args[0][-1] + + +@pytest.mark.skipif(sys.platform != "win32", reason="Windows-specific process existence behavior") +def test_windows_process_exists_does_not_terminate_process(): + process = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"]) + try: + assert _server_launcher._process_exists(pid=process.pid) is True + assert process.poll() is None + finally: + process.terminate() + process.wait(timeout=5) diff --git a/tests/unit/setup/test_configuration_loader.py b/tests/unit/setup/test_configuration_loader.py index a42de74afe..020c1bca9b 100644 --- a/tests/unit/setup/test_configuration_loader.py +++ b/tests/unit/setup/test_configuration_loader.py @@ -626,6 +626,105 @@ def test_load_with_overrides_preserves_silent_from_config_file(self, mock_defaul finally: config_path.unlink() + @mock.patch("pyrit.setup.configuration_loader.DEFAULT_CONFIG_PATH") + def test_load_with_overrides_preserves_all_explicit_config_fields(self, mock_default_path, tmp_path): + mock_default_path.exists.return_value = False + config_path = tmp_path / "explicit.yaml" + config_path.write_text( + """ +max_concurrent_scenario_runs: 7 +allow_custom_initializers: true +server: + url: http://localhost:8765/ + startup_timeout: 45 +extensions: + feature_flag: enabled +""", + encoding="utf-8", + ) + + config = ConfigurationLoader.load_with_overrides(config_file=config_path) + + assert config.max_concurrent_scenario_runs == 7 + assert config.allow_custom_initializers is True + assert config.server == {"url": "http://localhost:8765/", "startup_timeout": 45} + assert config.server_config == ServerConfig(url="http://localhost:8765", startup_timeout=45.0) + assert config.extensions == {"feature_flag": "enabled"} + + def test_load_with_overrides_uses_key_presence_when_merging_layers(self, tmp_path): + default_path = tmp_path / ".pyrit_conf" + default_path.write_text( + """ +initializers: + - target +initialization_scripts: + - default.py +silent: true +operator: default-operator +operation: default-operation +allow_custom_initializers: true +server: + url: http://localhost:9000 + startup_timeout: 180 +extensions: + default_extension: enabled +""", + encoding="utf-8", + ) + overlay_path = tmp_path / "overlay.yaml" + overlay_path.write_text( + """ +initialization_scripts: [] +silent: false +operator: "" +operation: null +allow_custom_initializers: false +extensions: {} +""", + encoding="utf-8", + ) + + with mock.patch("pyrit.setup.configuration_loader.DEFAULT_CONFIG_PATH", default_path): + config = ConfigurationLoader.load_with_overrides(config_file=overlay_path) + + assert config.initializers == ["target"] + assert config.initialization_scripts == [] + assert config.silent is False + assert config.operator == "" + assert config.operation is None + assert config.allow_custom_initializers is False + assert config.server_config == ServerConfig(url="http://localhost:9000", startup_timeout=180.0) + assert config.extensions == {} + + def test_load_with_overrides_merges_explicit_extension_keys(self, tmp_path): + default_path = tmp_path / ".pyrit_conf" + default_path.write_text("alpha: 1\nbeta: 2\n", encoding="utf-8") + overlay_path = tmp_path / "overlay.yaml" + overlay_path.write_text("alpha: 3\n", encoding="utf-8") + + with mock.patch("pyrit.setup.configuration_loader.DEFAULT_CONFIG_PATH", default_path): + config = ConfigurationLoader.load_with_overrides(config_file=overlay_path) + + assert config.extensions == {"alpha": 3, "beta": 2} + + def test_load_with_overrides_preserves_extension_provenance(self, tmp_path): + default_path = tmp_path / ".pyrit_conf" + default_path.write_text("alpha: 1\nbeta: 2\n", encoding="utf-8") + overlay_path = tmp_path / "overlay.yaml" + overlay_path.write_text( + "alpha: top-level\nextensions:\n alpha: nested\n", + encoding="utf-8", + ) + + with mock.patch("pyrit.setup.configuration_loader.DEFAULT_CONFIG_PATH", default_path): + config = ConfigurationLoader.load_with_overrides(config_file=overlay_path) + + assert config.extensions == {"alpha": "nested", "beta": 2} + + def test_extensions_reject_non_string_keys(self): + with pytest.raises(ValueError, match="extensions keys must be strings"): + ConfigurationLoader.from_dict({"extensions": {1: "value"}}) + @pytest.mark.parametrize("scenario_config", [None, {"name": "scam"}]) def test_scenario_config_block_is_rejected(scenario_config): @@ -669,6 +768,15 @@ def test_server_dict_without_url_uses_default(self): config = ConfigurationLoader(server={}) assert config.server_config == ServerConfig(url="http://localhost:8000") + def test_server_startup_timeout_normalizes(self): + config = ConfigurationLoader(server={"startup_timeout": 45}) + assert config.server_config == ServerConfig(url="http://localhost:8000", startup_timeout=45.0) + + @pytest.mark.parametrize("startup_timeout", [True, 0, -1, float("inf"), "slow"]) + def test_server_invalid_startup_timeout_raises(self, startup_timeout): + with pytest.raises(ValueError, match="'startup_timeout' must be a finite number greater than 0"): + ConfigurationLoader(server={"startup_timeout": startup_timeout}) + def test_server_url_non_string_raises(self): with pytest.raises(ValueError, match="Server 'url' must be a string"): ConfigurationLoader(server={"url": 12345})