cloudru-qa-ssh-client 1.0.0__py3-none-any.whl

ssh_client/clients/asyncssh_client.py

@@ -0,0 +1,155 @@
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Unpack, cast
+
+import asyncssh
+
+from ssh_client.client import BaseAsyncSSHClient
+from ssh_client.config import ExecutionSettings
+from ssh_client.results import LiveSSHResult, SSHResult
+
+from .exceptions import CommandNonZeroExitCodeError, SSHClientNotDefinedError
+
+if TYPE_CHECKING:
+    from contextlib import AbstractAsyncContextManager
+
+    from ssh_client.types import AsyncOperation, ExecuteOptions, InterruptCallback, LiveExecuteOptions, OutputCallback
+
+
+@dataclass
+class _StreamReader:
+    process: asyncssh.SSHClientProcess[str]
+    lines: list[str]
+    callbacks: tuple[OutputCallback | None, ...]
+    interrupt: InterruptCallback | None
+
+    async def read(self, stream: asyncssh.SSHReader[str]) -> None:
+        async for raw in stream:
+            if not raw:
+                continue
+            line = raw.removesuffix("\n").removesuffix("\r")
+            self.lines.append(line)
+            for callback in self.callbacks:
+                if callback is not None:
+                    callback(line)
+            if self.interrupt is not None and self.interrupt(line):
+                self.process.send_signal("INT")
+                return
+
+
+class AsyncSSHClient(BaseAsyncSSHClient[asyncssh.SSHClientConnection]):
+    """Run asynchronous SSH commands with trusted server keys."""
+
+    async def connect(self) -> None:
+        """Open a connection after checking the server host key."""
+        if self._connected:
+            return
+        operation: AsyncOperation[None] = self._connect_attempt
+        operation = self._decorate(operation, self._connect_decorators)
+        if self.connection_retry is None:
+            await operation()
+        else:
+            await self.connection_retry.connect_async(operation)
+
+    async def execute(self, command: str | list[str], **options: Unpack[ExecuteOptions]) -> SSHResult:
+        """Execute a command and wait for connection cleanup.
+
+        :param command: Shell command or preformatted shell fragments.
+        :param options: Timeout and keyword-only status-check settings.
+        :return: Captured output and exit status.
+        """
+        settings = ExecutionSettings(**options)
+        text = " ".join(command) if isinstance(command, list) else command
+        try:
+            await self.connect()
+            self._request(text)
+
+            async def execute() -> SSHResult:
+                return await self._execute_command(text, settings)
+
+            operation: AsyncOperation[SSHResult] = execute
+            operation = self._decorate(operation, self._execute_decorators)
+            result = (
+                await operation() if self.command_retry is None else await self.command_retry.execute_async(operation)
+            )
+            self._response(text, result)
+            if settings.raise_for_status and result.exit_code != 0:
+                raise CommandNonZeroExitCodeError.from_result(result)
+            return result
+        finally:
+            await self.aclose()
+
+    async def execute_live(self, command: str | list[str], **options: Unpack[LiveExecuteOptions]) -> LiveSSHResult:
+        """Stream stdout and stderr concurrently and release the connection.
+
+        :param command: Shell command or preformatted shell fragments.
+        :param options: Timeout, status checking, output callbacks and interrupt callbacks.
+        :return: Collected output lines and exit status.
+        """
+        settings = ExecutionSettings(**options)
+        text = " ".join(command) if isinstance(command, list) else command
+        try:
+            await self.connect()
+            self._request(text)
+            result = await self._execute_live_command(text, settings)
+            self._response(text, result)
+            if settings.raise_for_status and result.exit_code != 0:
+                raise CommandNonZeroExitCodeError.from_result(result)
+            return result
+        finally:
+            await self.aclose()
+
+    async def _connect_attempt(self) -> None:
+        await self.aclose()
+        if not self.username:
+            message = "Username must be set by authentication options (with_password or with_private_key)"
+            raise ValueError(message)
+        client_keys: list[str] | tuple[()] = [self.private_key_path] if self.private_key_path is not None else ()
+        self._client = await asyncssh.connect(
+            self.host,
+            port=self.port,
+            username=self.username,
+            password=self.password,
+            client_keys=client_keys,
+            known_hosts=self.known_hosts if self.known_hosts is not None else (),
+            connect_timeout=self.connection_timeout if self.connection_timeout is not None else self.timeout,
+        )
+        self._connected = True
+
+    def _require_client(self) -> asyncssh.SSHClientConnection:
+        if self._client is None or not self._connected:
+            raise SSHClientNotDefinedError
+        return self._client
+
+    async def _execute_command(self, command: str, settings: ExecutionSettings) -> SSHResult:
+        async with asyncio.timeout(self.timeout if settings.timeout is None else settings.timeout):
+            result = await self._require_client().run(command, encoding="utf-8")
+        output = result.stdout.decode("utf-8") if isinstance(result.stdout, bytes) else result.stdout
+        errors = result.stderr.decode("utf-8") if isinstance(result.stderr, bytes) else result.stderr
+        return SSHResult(output or "", errors or "", result.returncode)
+
+    async def _execute_live_command(self, command: str, settings: ExecutionSettings) -> LiveSSHResult:
+        async with cast(
+            "AbstractAsyncContextManager[asyncssh.SSHClientProcess[str]]",
+            self._require_client().create_process(command, encoding="utf-8"),
+        ) as process:
+            output: list[str] = []
+            errors: list[str] = []
+            stdout = _StreamReader(
+                process,
+                output,
+                (settings.stdout_callback, self.live_stdout_callback),
+                settings.interrupt_stdout_callback,
+            )
+            stderr = _StreamReader(
+                process,
+                errors,
+                (settings.stderr_callback, self.live_stderr_callback),
+                settings.interrupt_stderr_callback,
+            )
+            async with asyncio.timeout(self.timeout if settings.timeout is None else settings.timeout):
+                await asyncio.gather(stdout.read(process.stdout), stderr.read(process.stderr))
+                await process.wait()
+            return LiveSSHResult(output, errors, process.returncode)

ssh_client/clients/exceptions.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from ssh_client.results import LiveSSHResult, SSHResult
+
+
+class SSHClientNotDefinedError(Exception):
+    """Report an operation without an authenticated transport."""
+
+    def __init__(self, message: str = "SSH client is not defined") -> None:
+        """Keep the transport failure reason.
+
+        :param message: Failure description.
+        """
+        self.message = message
+        super().__init__(message)
+
+
+class CommandNonZeroExitCodeError(Exception):
+    """Report a command which returned a failing exit status."""
+
+    @classmethod
+    def from_result(cls, result: SSHResult | LiveSSHResult) -> CommandNonZeroExitCodeError:
+        """Describe the failed command result.
+
+        :param result: Failed command output and exit status.
+        :return: Exception describing the command failure.
+        """
+        stderr = result.stderr if isinstance(result, SSHResult) else "\n".join(result.stderr_lines)
+        return cls(f"Command failed with exit code {result.exit_code}: {stderr}")

ssh_client/clients/paramiko_client.py

@@ -0,0 +1,190 @@
+from __future__ import annotations
+
+import time
+from codecs import getincrementaldecoder
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Unpack
+
+import paramiko
+
+from ssh_client.client import BaseSSHClient
+from ssh_client.config import ExecutionSettings
+from ssh_client.results import LiveSSHResult, SSHResult
+
+from .exceptions import CommandNonZeroExitCodeError, SSHClientNotDefinedError
+
+if TYPE_CHECKING:
+    from codecs import IncrementalDecoder
+
+    from ssh_client.types import ExecuteOptions, LiveExecuteOptions, OutputCallback
+
+
+@dataclass
+class _LiveOutput:
+    settings: ExecutionSettings
+    callbacks: tuple[OutputCallback | None, ...]
+    lines: list[str] = field(default_factory=list)
+    pending: str = ""
+    decoder: IncrementalDecoder = field(default_factory=getincrementaldecoder("utf-8"))
+
+    def receive(self, data: bytes, *, final: bool = False) -> bool:
+        self.pending += self.decoder.decode(data, final=final)
+        while "\n" in self.pending:
+            line, self.pending = self.pending.split("\n", 1)
+            if self._line(line.removesuffix("\r")):
+                return True
+        return bool(final and self.pending and self._line(self.pending))
+
+    def _line(self, line: str) -> bool:
+        self.lines.append(line)
+        for callback in self.callbacks:
+            if callback is not None:
+                callback(line)
+        interrupt = self.settings.interrupt_stdout_callback
+        return interrupt is not None and interrupt(line)
+
+
+class SyncSSHClient(BaseSSHClient[paramiko.SSHClient]):
+    """Run SSH commands through Paramiko with trusted server keys."""
+
+    def connect(self) -> None:
+        """Open a connection after checking the server host key."""
+        if self._connected:
+            return
+        operation = self._decorate(self._connect_attempt, self._connect_decorators)
+        if self.connection_retry is None:
+            operation()
+        else:
+            self.connection_retry.connect(operation)
+
+    def execute(self, command: str | list[str], **options: Unpack[ExecuteOptions]) -> SSHResult:
+        """Execute a command and close its connection.
+
+        :param command: Shell command or preformatted shell fragments.
+        :param options: Timeout and keyword-only status-check settings.
+        :return: Captured output and exit status.
+        """
+        settings = ExecutionSettings(**options)
+        text = " ".join(command) if isinstance(command, list) else command
+        try:
+            self.connect()
+            self._request(text)
+
+            def execute() -> SSHResult:
+                return self._execute_command(text, settings)
+
+            operation = self._decorate(execute, self._execute_decorators)
+            result = operation() if self.command_retry is None else self.command_retry.execute(operation)
+            self._response(text, result)
+            if settings.raise_for_status and result.exit_code != 0:
+                raise CommandNonZeroExitCodeError.from_result(result)
+            return result
+        finally:
+            self.close()
+
+    def execute_live(self, command: str | list[str], **options: Unpack[LiveExecuteOptions]) -> LiveSSHResult:
+        """Stream command lines through callbacks and close the connection.
+
+        Paramiko uses a terminal so that interrupting output also stops the remote process;
+        stderr is therefore merged into stdout.
+
+        :param command: Shell command or preformatted shell fragments.
+        :param options: Timeout, status checking, stdout callbacks and interrupt callback.
+        :return: Collected output lines and exit status.
+        """
+        settings = ExecutionSettings(**options)
+        text = " ".join(command) if isinstance(command, list) else command
+        try:
+            self.connect()
+            self._request(text)
+            result = self._execute_live_command(text, settings)
+            self._response(text, result)
+            if settings.raise_for_status and result.exit_code != 0:
+                raise CommandNonZeroExitCodeError.from_result(result)
+            return result
+        finally:
+            self.close()
+
+    def _connect_attempt(self) -> None:
+        self.close()
+        if not self.username:
+            message = "Username must be set by authentication options (with_password or with_private_key)"
+            raise ValueError(message)
+        self._client = paramiko.SSHClient()
+        if self.known_hosts is None:
+            self._client.load_system_host_keys()
+        else:
+            self._client.load_host_keys(self.known_hosts)
+        try:
+            self._client.connect(
+                hostname=self.host,
+                port=self.port,
+                username=self.username,
+                password=self.password,
+                key_filename=self.private_key_path,
+                timeout=self.connection_timeout if self.connection_timeout is not None else self.timeout,
+                allow_agent=self.password is None and self.private_key_path is None,
+                look_for_keys=self.password is None and self.private_key_path is None,
+            )
+        except BaseException:
+            self.close()
+            raise
+        self._connected = True
+
+    def _require_client(self) -> paramiko.SSHClient:
+        if self._client is None or not self._connected:
+            raise SSHClientNotDefinedError
+        return self._client
+
+    def _execute_command(self, command: str, settings: ExecutionSettings) -> SSHResult:
+        transport = self._require_client()
+        stdin, stdout, stderr = transport.exec_command(command, timeout=settings.timeout or self.timeout)
+        stdin.close()
+        try:
+            return self._collect_output(stdout.channel, settings)
+        finally:
+            stdout.close()
+            stderr.close()
+
+    def _collect_output(self, channel: paramiko.Channel, settings: ExecutionSettings) -> SSHResult:
+        output = bytearray()
+        errors = bytearray()
+        deadline = time.monotonic() + (self.timeout if settings.timeout is None else settings.timeout)
+        while True:
+            if channel.recv_ready():
+                output.extend(channel.recv(65536))
+            if channel.recv_stderr_ready():
+                errors.extend(channel.recv_stderr(65536))
+            if channel.exit_status_ready() and not channel.recv_ready() and not channel.recv_stderr_ready():
+                return SSHResult(output.decode("utf-8"), errors.decode("utf-8"), channel.recv_exit_status())
+            self._wait_for_output(deadline)
+
+    def _execute_live_command(self, command: str, settings: ExecutionSettings) -> LiveSSHResult:
+        stdin, stdout, stderr = self._require_client().exec_command(
+            command,
+            timeout=settings.timeout or self.timeout,
+            get_pty=True,
+        )
+        stdin.close()
+        output = _LiveOutput(settings, (settings.stdout_callback, self.live_stdout_callback))
+        deadline = time.monotonic() + (self.timeout if settings.timeout is None else settings.timeout)
+        try:
+            while True:
+                if stdout.channel.recv_ready() and output.receive(stdout.channel.recv(65536)):
+                    stdout.channel.send(b"\x03")
+                    stdout.channel.close()
+                    return LiveSSHResult(output.lines, [], -1)
+                if stdout.channel.exit_status_ready() and not stdout.channel.recv_ready():
+                    output.receive(b"", final=True)
+                    return LiveSSHResult(output.lines, [], stdout.channel.recv_exit_status())
+                self._wait_for_output(deadline)
+        finally:
+            stdout.close()
+            stderr.close()
+
+    @staticmethod
+    def _wait_for_output(deadline: float) -> None:
+        if time.monotonic() >= deadline:
+            message = "SSH command timed out"
+            raise TimeoutError(message)
+        time.sleep(0.01)

ssh_client/config.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .types import InterruptCallback, OutputCallback
+    from .utils.retries import RetryPolicy
+
+
+@dataclass
+class SSHClientConfig:
+    """Store connection and command settings."""
+
+    username: str | None = None
+    password: str | None = None
+    private_key_path: str | None = None
+    known_hosts: str | None = None
+    timeout: float = 30
+    connection_timeout: float | None = None
+    connection_retry: RetryPolicy | None = None
+    command_retry: RetryPolicy | None = None
+    live_stdout_callback: OutputCallback | None = None
+    live_stderr_callback: OutputCallback | None = None
+
+
+@dataclass
+class ExecutionSettings:
+    """Group command settings passed to execution helpers."""
+
+    timeout: float | None = None
+    raise_for_status: bool = False
+    stdout_callback: OutputCallback | None = None
+    stderr_callback: OutputCallback | None = None
+    interrupt_stdout_callback: InterruptCallback | None = None
+    interrupt_stderr_callback: InterruptCallback | None = None

ssh_client/event_hooks.py

@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+import re
+from typing import TYPE_CHECKING
+
+from .results import LiveSSHResult
+from .utils.masking import mask_request_data, mask_response_data
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+    from logging import Logger
+
+    from .results import SSHResult
+    from .types import RequestHook, ResponseHook, StructuredLogger
+
+type MaskPattern = str | re.Pattern[str]
+
+
+def _outputs(result: SSHResult | LiveSSHResult, patterns: Sequence[MaskPattern]) -> tuple[str, str]:
+    if isinstance(result, LiveSSHResult):
+        return mask_response_data("\n".join(result.stdout_lines), "\n".join(result.stderr_lines), patterns)
+    return mask_response_data(result.stdout, result.stderr, patterns)
+
+
+def create_request_logger_hook(logger: Logger, mask_patterns: Sequence[MaskPattern] | None = None) -> RequestHook:
+    """Create a sanitized request log callback.
+
+    :param logger: Standard logger.
+    :param mask_patterns: Literal secrets or regular expressions.
+    :return: Request observer.
+    """
+
+    def log_request(*, host: str, port: int, username: str | None, command: str) -> None:
+        logger.info(
+            "SSH request: host=%s, port=%s, username=%s, command=%s",
+            host,
+            port,
+            username,
+            mask_request_data(command, mask_patterns or ()),
+        )
+
+    return log_request
+
+
+def create_response_logger_hook(logger: Logger, mask_patterns: Sequence[MaskPattern] | None = None) -> ResponseHook:
+    """Create a sanitized command result log callback.
+
+    :param logger: Standard logger.
+    :param mask_patterns: Literal secrets or regular expressions.
+    :return: Response observer.
+    """
+
+    def log_response(*, host: str, command: str, result: SSHResult | LiveSSHResult) -> None:
+        output, errors = _outputs(result, mask_patterns or ())
+        logger.info(
+            "SSH response: host=%s, exit_code=%s, stdout=%s, stderr=%s, command=%s",
+            host,
+            result.exit_code,
+            output,
+            errors,
+            mask_request_data(command, mask_patterns or ()),
+        )
+
+    return log_response
+
+
+def create_request_structlogger_hook(
+    logger: StructuredLogger,
+    mask_patterns: Sequence[MaskPattern] | None = None,
+) -> RequestHook:
+    """Create a sanitized structured request observer.
+
+    :param logger: Structured logger.
+    :param mask_patterns: Literal secrets or regular expressions.
+    :return: Request observer.
+    """
+
+    def log_request(*, host: str, port: int, username: str | None, command: str) -> None:
+        logger.info(
+            "ssh request sent",
+            host=host,
+            port=port,
+            username=username,
+            command=mask_request_data(command, mask_patterns or ()),
+        )
+
+    return log_request
+
+
+def create_response_structlogger_hook(
+    logger: StructuredLogger,
+    mask_patterns: Sequence[MaskPattern] | None = None,
+) -> ResponseHook:
+    """Create a sanitized structured command result observer.
+
+    :param logger: Structured logger.
+    :param mask_patterns: Literal secrets or regular expressions.
+    :return: Response observer.
+    """
+
+    def log_response(*, host: str, command: str, result: SSHResult | LiveSSHResult) -> None:
+        output, errors = _outputs(result, mask_patterns or ())
+        logger.info(
+            "ssh response received",
+            host=host,
+            exit_code=result.exit_code,
+            stdout=output,
+            stderr=errors,
+            command=mask_request_data(command, mask_patterns or ()),
+        )
+
+    return log_response

ssh_client/examples/basic_usage.py

@@ -0,0 +1,83 @@
+"""Command examples with host key verification and connection cleanup."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import TYPE_CHECKING, cast
+
+import structlog
+
+from ssh_client import AsyncSSHClient, SyncSSHClient, options
+
+if TYPE_CHECKING:
+    from ssh_client.types import StructuredLogger
+
+logger = logging.getLogger(__name__)
+
+
+def sync_example() -> None:
+    """Run a command with an SSH agent key and the default known_hosts."""
+    with SyncSSHClient("localhost", 22, options.with_private_key("qa-user", "/path/to/private/key")) as client:
+        result = client.execute("uname -a", timeout=10, raise_for_status=True)
+        logger.info("Remote system: %s", result.stdout)
+
+
+async def async_example() -> None:
+    """Run an asynchronous command with an explicit list of trusted servers."""
+    async with AsyncSSHClient(
+        "localhost",
+        22,
+        options.with_private_key("qa-user", "/path/to/private/key"),
+        options.with_known_hosts("/path/to/known_hosts"),
+        options.with_logger(logger),
+    ) as client:
+        result = await client.execute("uname -a", timeout=10, raise_for_status=True)
+        logger.info("Remote system: %s", result.stdout)
+
+
+def advanced_sync_example() -> None:
+    """Pass command lines to a handler and limit reconnects."""
+    with SyncSSHClient(
+        "localhost",
+        22,
+        options.with_private_key("qa-user", "/path/to/private/key"),
+        options.with_live_execute_callback(lambda line: logger.info("Output: %s", line)),
+        options.retry_if_connection_fails(3, 1, reraise_exception=True),
+    ) as client:
+        result = client.execute_live("printf 'first\\nsecond\\n'", timeout=10, raise_for_status=True)
+        logger.info("Command exit status: %s", result.exit_code)
+
+
+async def advanced_async_example() -> None:
+    """Override settings for one command using a separate client copy."""
+    async with AsyncSSHClient(
+        "localhost",
+        22,
+        options.with_private_key("qa-user", "/path/to/private/key"),
+    ) as client:
+        async with client.with_options(options.with_request_timeout(10)) as temporary:
+            result = await temporary.execute_live("printf 'one\\ntwo\\n'", raise_for_status=True)
+        logger.info("Output: %s", result.stdout_lines)
+
+
+def structlog_example() -> None:
+    """Log a command through structlog with the specified text masked."""
+    structured = cast("StructuredLogger", structlog.get_logger())
+    with SyncSSHClient(
+        "localhost",
+        22,
+        options.with_private_key("qa-user", "/path/to/private/key"),
+        options.with_structlogger(structured, mask_patterns=["sensitive-value"]),
+    ) as client:
+        result = client.execute("printf 'sensitive-value'", raise_for_status=True)
+        logger.info("Command exit status: %s", result.exit_code)
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+    sync_example()
+    asyncio.run(async_example())
+    advanced_sync_example()
+    asyncio.run(advanced_async_example())
+    structlog_example()