local-ci-cd-system/src/ci_orchestrator/transport/winrm.py

"""WinRM transport via ``pypsrp``.

Replaces the ``New-PSSession`` / ``Invoke-Command`` calls used in the
PowerShell scripts. CI build VMs use a self-signed TLS certificate on
WinRM port 5986, so SSL validation is disabled by default — appropriate
for an isolated lab network (see ``scripts/_Common.psm1``).
"""

from __future__ import annotations

import contextlib
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any

from ci_orchestrator.transport.errors import TransportCommandError, TransportConnectError

if TYPE_CHECKING:  # pragma: no cover
    from pypsrp.client import Client


@dataclass(frozen=True)
class WinRmResult:
    """Result of a remote PowerShell invocation."""

    stdout: str
    stderr: str
    returncode: int

    @property
    def ok(self) -> bool:
        return self.returncode == 0


class WinRmTransport:
    """Thin wrapper around :class:`pypsrp.client.Client`.

    Parameters
    ----------
    host:
        IP or DNS name of the guest.
    username, password:
        Credentials. Pass via :class:`~ci_orchestrator.credentials.Credential`
        in caller code; do not embed plaintext.
    port:
        WinRM HTTPS port (default 5986).
    cert_validation:
        When ``False`` (default), self-signed certs are accepted — required
        for the ephemeral build VMs in this lab.
    """

    def __init__(
        self,
        host: str,
        username: str,
        password: str,
        *,
        port: int = 5986,
        ssl: bool = True,
        cert_validation: bool = False,
        connection_timeout: int = 30,
    ) -> None:
        self._host = host
        self._username = username
        self._password = password
        self._port = port
        self._ssl = ssl
        self._cert_validation = cert_validation
        self._connection_timeout = connection_timeout
        self._client: Client | None = None

    # ----------------------------------------------------------- connection
    def _connect(self) -> Client:
        if self._client is not None:
            return self._client
        try:
            from pypsrp.client import Client  # local import: optional dep at runtime
        except ImportError as exc:  # pragma: no cover - dep declared in pyproject
            raise TransportConnectError(
                "pypsrp is not installed; run `pip install pypsrp`."
            ) from exc
        try:
            self._client = Client(
                self._host,
                username=self._username,
                password=self._password,
                port=self._port,
                ssl=self._ssl,
                cert_validation=self._cert_validation,
                connection_timeout=self._connection_timeout,
            )
        except Exception as exc:
            raise TransportConnectError(
                f"WinRM connection to {self._host}:{self._port} failed: {exc}"
            ) from exc
        return self._client

    def close(self) -> None:
        if self._client is not None:
            with contextlib.suppress(Exception):
                self._client.wsman.close()
            self._client = None

    def __enter__(self) -> WinRmTransport:
        self._connect()
        return self

    def __exit__(self, *_exc: object) -> None:
        self.close()

    # --------------------------------------------------------------- ops
    def run(self, script: str, *, check: bool = True) -> WinRmResult:
        """Run a PowerShell script block in the guest."""
        client = self._connect()
        try:
            stdout, stderr, returncode = client.execute_ps(script)
        except Exception as exc:
            raise TransportConnectError(f"WinRM execute_ps failed: {exc}") from exc
        # pypsrp returns stderr as a PSDataStreams object (newer versions) or a
        # list of error records (older versions); normalise to a plain string.
        stderr_any: Any = stderr
        if isinstance(stderr_any, str):
            stderr_str = stderr_any
        elif hasattr(stderr_any, "error"):
            # PSDataStreams — extract .error records.
            error_records = getattr(stderr_any, "error", None) or []
            stderr_str = "\n".join(str(e) for e in error_records)
        else:
            try:
                stderr_str = "\n".join(str(s) for s in stderr_any)
            except TypeError:
                stderr_str = str(stderr_any)
        result = WinRmResult(stdout=stdout, stderr=stderr_str, returncode=int(returncode))
        if check and not result.ok:
            raise TransportCommandError(result.returncode, result.stdout, result.stderr)
        return result

    def copy(self, local_path: str | Path, remote_path: str) -> None:
        """Upload a file to the guest via PSRP."""
        client = self._connect()
        try:
            client.copy(str(local_path), remote_path)
        except Exception as exc:
            raise TransportConnectError(f"WinRM copy failed: {exc}") from exc

    def fetch(self, remote_path: str, local_path: str | Path) -> None:
        """Download a file from the guest via PSRP."""
        client = self._connect()
        try:
            client.fetch(remote_path, str(local_path))
        except Exception as exc:
            raise TransportConnectError(f"WinRM fetch failed: {exc}") from exc

    def is_ready(self) -> bool:
        """Lightweight readiness probe: run ``$true`` and check exit code."""
        try:
            result = self.run("$true | Out-Null", check=False)
        except (TransportConnectError, TransportCommandError):
            return False
        return result.ok