local-ci-cd-system/src/ci_orchestrator/commands/wait.py

"""``wait-ready`` sub-command.

Replaces ``scripts/Wait-VMReady.ps1``. Polls a guest VM until WinRM
(Windows) or SSH (Linux) is responsive. Designed to be invoked from the
PowerShell shim ``scripts/Wait-VMReady.ps1`` which translates PS-style
PascalCase parameters (``-VMPath``, ``-IPAddress``, ``-Transport``,
``-SshKeyPath``, ``-SshUser``) into click kebab-case options.
"""

from __future__ import annotations

import sys
import time
from typing import TYPE_CHECKING

import click

from ci_orchestrator.backends import load_backend
from ci_orchestrator.backends.errors import BackendError
from ci_orchestrator.backends.protocol import VmHandle
from ci_orchestrator.config import load_config
from ci_orchestrator.credentials import KeyringCredentialStore
from ci_orchestrator.transport.errors import TransportError
from ci_orchestrator.transport.ssh import SshTransport
from ci_orchestrator.transport.winrm import WinRmTransport

if TYPE_CHECKING:  # pragma: no cover
    from ci_orchestrator.backends.protocol import VmBackend


# Exit codes (kept in sync with the PowerShell original where possible):
#   0 ready, 2 timeout-running, 3 timeout-ip, 4 timeout-transport
_EXIT_TIMEOUT_RUNNING = 2
_EXIT_TIMEOUT_IP = 3
_EXIT_TIMEOUT_TRANSPORT = 4


def _normalise_guest_os(value: str) -> str:
    """Map Transport/GuestOS aliases to ``windows``/``linux``."""
    v = value.strip().lower()
    if v in {"winrm", "windows", "win"}:
        return "windows"
    if v in {"ssh", "linux", "lnx"}:
        return "linux"
    raise click.BadParameter(f"Unknown guest-os/transport value: {value!r}")


def _wait_running(
    backend: VmBackend, handle: VmHandle, deadline: float, poll: float
) -> bool:
    while time.monotonic() < deadline:
        try:
            if backend.is_running(handle):
                return True
        except BackendError as exc:
            click.echo(f"[wait-ready] backend probe error: {exc}", err=True)
        time.sleep(poll)
    return False


def _wait_for_ip(
    backend: VmBackend, handle: VmHandle, deadline: float, poll: float
) -> str | None:
    while time.monotonic() < deadline:
        try:
            ip = backend.get_ip(handle)
        except BackendError:
            ip = None
        if ip:
            return ip
        time.sleep(poll)
    return None


@click.command("wait-ready")
@click.option(
    "--vmx",
    "--vm-path",
    "vmx",
    required=True,
    help="Path to the guest VMX file.",
)
@click.option(
    "--ip-address",
    "ip_address",
    default=None,
    help="Guest IP address. If omitted the backend is polled for it.",
)
@click.option(
    "--guest-os",
    "--transport",
    "guest_os",
    default="windows",
    show_default=True,
    callback=lambda _ctx, _p, v: _normalise_guest_os(v),
    help="windows/WinRM or linux/SSH (case-insensitive).",
)
@click.option("--timeout", "--timeout-seconds", "timeout", type=float, default=300.0, show_default=True)
@click.option(
    "--poll-interval",
    "--poll-interval-seconds",
    "poll_interval",
    type=float,
    default=5.0,
    show_default=True,
)
@click.option("--vmrun-path", "vmrun_path", default=None, help="Override vmrun.exe path.")
@click.option(
    "--credential-target",
    default=None,
    help="Override credential target name (defaults to config.guest_cred_target).",
)
@click.option("--ssh-user", default="ci_build", show_default=True)
@click.option(
    "--ssh-key-path",
    default=None,
    help="SSH private key (defaults to config.ssh_key_path).",
)
@click.option(
    "--skip-ping",
    is_flag=True,
    default=False,
    help="Accepted for PS compatibility; ignored (Python uses transport probe directly).",
)
def wait_ready(
    vmx: str,
    ip_address: str | None,
    guest_os: str,
    timeout: float,
    poll_interval: float,
    vmrun_path: str | None,
    credential_target: str | None,
    ssh_user: str,
    ssh_key_path: str | None,
    skip_ping: bool,
) -> None:
    """Poll a guest until WinRM (Windows) or SSH (Linux) is responsive."""
    del skip_ping  # accepted for PS compatibility, no-op
    config = load_config()
    if vmrun_path is not None:
        # Re-load backend with override; cheap, no real connections.
        from ci_orchestrator.backends.workstation import WorkstationVmrunBackend

        backend: VmBackend = WorkstationVmrunBackend(vmrun_path=vmrun_path)
    else:
        backend = load_backend(config)
    handle = VmHandle(identifier=vmx)
    deadline = time.monotonic() + timeout

    click.echo(f"[wait-ready] waiting for backend to report VM running: {vmx}")
    if not _wait_running(backend, handle, deadline, poll_interval):
        click.echo("[wait-ready] timeout waiting for VM to be running.", err=True)
        sys.exit(_EXIT_TIMEOUT_RUNNING)

    if ip_address is None:
        ip = _wait_for_ip(backend, handle, deadline, poll_interval)
        if not ip:
            click.echo("[wait-ready] timeout waiting for guest IP.", err=True)
            sys.exit(_EXIT_TIMEOUT_IP)
        click.echo(f"[wait-ready] guest IP: {ip}")
    else:
        ip = ip_address
        click.echo(f"[wait-ready] using provided IP: {ip}")

    if guest_os == "windows":
        store = KeyringCredentialStore()
        target = credential_target or config.guest_cred_target
        cred = store.get(target)
        while time.monotonic() < deadline:
            try:
                with WinRmTransport(ip, cred.username, cred.password) as t:
                    if t.is_ready():
                        click.echo("[wait-ready] WinRM ready.")
                        return
            except TransportError as exc:
                click.echo(f"[wait-ready] WinRM probe failed: {exc}", err=True)
            time.sleep(poll_interval)
    else:
        key_path: str | None = ssh_key_path
        if key_path is None and config.ssh_key_path is not None:
            key_path = str(config.ssh_key_path)
        while time.monotonic() < deadline:
            try:
                with SshTransport(ip, username=ssh_user, key_path=key_path) as t:
                    if t.is_ready():
                        click.echo("[wait-ready] SSH ready.")
                        return
            except TransportError as exc:
                click.echo(f"[wait-ready] SSH probe failed: {exc}", err=True)
            time.sleep(poll_interval)

    click.echo("[wait-ready] timeout waiting for transport to be ready.", err=True)
    sys.exit(_EXIT_TIMEOUT_TRANSPORT)


__all__ = [
    "KeyringCredentialStore",
    "SshTransport",
    "WinRmTransport",
    "load_backend",
    "wait_ready",
]