connectors/ssh: add flag to use SCP as file transfer protocol

DonDebonair · web-flow · commit 70f045015532 · 2025-08-25T13:18:22.000+01:00
diff --git a/pyinfra/connectors/scp/__init__.py b/pyinfra/connectors/scp/__init__.py
@@ -0,0 +1 @@
+from .client import SCPClient  # noqa: F401
diff --git a/pyinfra/connectors/scp/client.py b/pyinfra/connectors/scp/client.py
@@ -0,0 +1,204 @@
+from __future__ import annotations
+
+import ntpath
+import os
+from pathlib import PurePath
+from shlex import quote
+from socket import timeout as SocketTimeoutError
+from typing import IO, AnyStr
+
+from paramiko import Channel
+from paramiko.transport import Transport
+
+SCP_COMMAND = b"scp"
+
+
+# Unicode conversion functions; assume UTF-8
+def asbytes(s: bytes | str | PurePath) -> bytes:
+    """Turns unicode into bytes, if needed.
+
+    Assumes UTF-8.
+    """
+    if isinstance(s, bytes):
+        return s
+    elif isinstance(s, PurePath):
+        return bytes(s)
+    else:
+        return s.encode("utf-8")
+
+
+def asunicode(s: bytes | str) -> str:
+    """Turns bytes into unicode, if needed.
+
+    Uses UTF-8.
+    """
+    if isinstance(s, bytes):
+        return s.decode("utf-8", "replace")
+    else:
+        return s
+
+
+class SCPClient:
+    """
+    An scp1 implementation, compatible with openssh scp.
+    Raises SCPException for all transport related errors. Local filesystem
+    and OS errors pass through.
+
+    Main public methods are .putfo and .getfo
+    """
+
+    def __init__(
+        self,
+        transport: Transport,
+        buff_size: int = 16384,
+        socket_timeout: float = 10.0,
+    ):
+        self.transport = transport
+        self.buff_size = buff_size
+        self.socket_timeout = socket_timeout
+        self._channel: Channel | None = None
+        self.scp_command = SCP_COMMAND
+
+    @property
+    def channel(self) -> Channel:
+        """Return an open Channel, (re)opening if needed."""
+        if self._channel is None or self._channel.closed:
+            self._channel = self.transport.open_session()
+        return self._channel
+
+    def __enter__(self):
+        _ = self.channel  # triggers opening if not already open
+        return self
+
+    def __exit__(self, type, value, traceback):
+        self.close()
+
+    def putfo(
+        self,
+        fl: IO[AnyStr],
+        remote_path: str | bytes,
+        mode: str | bytes = "0644",
+        size: int | None = None,
+    ) -> None:
+        if size is None:
+            pos = fl.tell()
+            fl.seek(0, os.SEEK_END)  # Seek to end
+            size = fl.tell() - pos
+            fl.seek(pos, os.SEEK_SET)  # Seek back
+
+        self.channel.settimeout(self.socket_timeout)
+        self.channel.exec_command(
+            self.scp_command + b" -t " + asbytes(quote(asunicode(remote_path)))
+        )
+        self._recv_confirm()
+        self._send_file(fl, remote_path, mode, size=size)
+        self.close()
+
+    def getfo(self, remote_path: str, fl: IO):
+        remote_path_sanitized = quote(remote_path)
+        if os.name == "nt":
+            remote_file_name = ntpath.basename(remote_path_sanitized)
+        else:
+            remote_file_name = os.path.basename(remote_path_sanitized)
+        self.channel.settimeout(self.socket_timeout)
+        self.channel.exec_command(self.scp_command + b" -f " + asbytes(remote_path_sanitized))
+        self._recv_all(fl, remote_file_name)
+        self.close()
+        return fl
+
+    def close(self):
+        """close scp channel"""
+        if self._channel is not None:
+            self._channel.close()
+            self._channel = None
+
+    def _send_file(self, fl, name, mode, size):
+        basename = asbytes(os.path.basename(name))
+        # The protocol can't handle \n in the filename.
+        # Quote them as the control sequence \^J for now,
+        # which is how openssh handles it.
+        self.channel.sendall(
+            ("C%s %d " % (mode, size)).encode("ascii") + basename.replace(b"\n", b"\\^J") + b"\n"
+        )
+        self._recv_confirm()
+        file_pos = 0
+        buff_size = self.buff_size
+        chan = self.channel
+        while file_pos < size:
+            chan.sendall(fl.read(buff_size))
+            file_pos = fl.tell()
+        chan.sendall(b"\x00")
+        self._recv_confirm()
+
+    def _recv_confirm(self):
+        # read scp response
+        msg = b""
+        try:
+            msg = self.channel.recv(512)
+        except SocketTimeoutError:
+            raise SCPException("Timeout waiting for scp response")
+        # slice off the first byte, so this compare will work in py2 and py3
+        if msg and msg[0:1] == b"\x00":
+            return
+        elif msg and msg[0:1] == b"\x01":
+            raise SCPException(asunicode(msg[1:]))
+        elif self.channel.recv_stderr_ready():
+            msg = self.channel.recv_stderr(512)
+            raise SCPException(asunicode(msg))
+        elif not msg:
+            raise SCPException("No response from server")
+        else:
+            raise SCPException("Invalid response from server", msg)
+
+    def _recv_all(self, fh: IO, remote_file_name: str) -> None:
+        # loop over scp commands, and receive as necessary
+        commands = (b"C",)
+        while not self.channel.closed:
+            # wait for command as long as we're open
+            self.channel.sendall(b"\x00")
+            msg = self.channel.recv(1024)
+            if not msg:  # chan closed while receiving
+                break
+            assert msg[-1:] == b"\n"
+            msg = msg[:-1]
+            code = msg[0:1]
+            if code not in commands:
+                raise SCPException(asunicode(msg[1:]))
+            self._recv_file(msg[1:], fh, remote_file_name)
+
+    def _recv_file(self, cmd: bytes, fh: IO, remote_file_name: str) -> None:
+        chan = self.channel
+        parts = cmd.strip().split(b" ", 2)
+
+        try:
+            size = int(parts[1])
+        except (ValueError, IndexError):
+            chan.send(b"\x01")
+            chan.close()
+            raise SCPException("Bad file format")
+
+        buff_size = self.buff_size
+        pos = 0
+        chan.send(b"\x00")
+        try:
+            while pos < size:
+                # we have to make sure we don't read the final byte
+                if size - pos <= buff_size:
+                    buff_size = size - pos
+                data = chan.recv(buff_size)
+                if not data:
+                    raise SCPException("Underlying channel was closed")
+                fh.write(data)
+                pos = fh.tell()
+            msg = chan.recv(512)
+            if msg and msg[0:1] != b"\x00":
+                raise SCPException(asunicode(msg[1:]))
+        except SocketTimeoutError:
+            chan.close()
+            raise SCPException("Error receiving, socket.timeout")
+
+
+class SCPException(Exception):
+    """SCP exception class"""
+
+    pass
diff --git a/pyinfra/connectors/ssh.py b/pyinfra/connectors/ssh.py
@@ -5,7 +5,7 @@
 from shutil import which
 from socket import error as socket_error, gaierror
 from time import sleep
-from typing import TYPE_CHECKING, Any, Iterable, Optional, Tuple
+from typing import IO, TYPE_CHECKING, Any, Iterable, Optional, Protocol, Tuple
 
 import click
 from paramiko import AuthenticationException, BadHostKeyException, SFTPClient, SSHException
@@ -17,6 +17,7 @@
 from pyinfra.api.util import get_file_io, memoize
 
 from .base import BaseConnector, DataMeta
+from .scp import SCPClient
 from .ssh_util import get_private_key, raise_connect_error
 from .sshuserclient import SSHClient
 from .util import (
@@ -53,6 +54,7 @@ class ConnectorData(TypedDict):
     ssh_connect_retries: int
     ssh_connect_retry_min_delay: float
     ssh_connect_retry_max_delay: float
+    ssh_file_transfer_protocol: str
 
 
 connector_data_meta: dict[str, DataMeta] = {
@@ -92,9 +94,27 @@ class ConnectorData(TypedDict):
         "Upper bound for random delay between retries",
         0.5,
     ),
+    "ssh_file_transfer_protocol": DataMeta(
+        "Protocol to use for file transfers. Can be ``sftp`` or ``scp``.",
+        "sftp",
+    ),
 }
 
 
+class FileTransferClient(Protocol):
+    def getfo(self, remote_filename: str, fl: IO) -> Any | None:
+        """
+        Get a file from the remote host, writing to the provided file-like object.
+        """
+        ...
+
+    def putfo(self, fl: IO, remote_filename: str) -> Any | None:
+        """
+        Put a file to the remote host, reading from the provided file-like object.
+        """
+        ...
+
+
 class SSHConnector(BaseConnector):
     """
     Connect to hosts over SSH. This is the default connector and all targets default
@@ -268,7 +288,7 @@ def _connect(self) -> None:
 
     @override
     def disconnect(self) -> None:
-        self.get_sftp_connection.cache.clear()
+        self.get_file_transfer_connection.cache.clear()
 
     @override
     def run_shell_command(
@@ -353,23 +373,35 @@ def execute_command() -> Tuple[int, CommandOutput]:
         return status, combined_output
 
     @memoize
-    def get_sftp_connection(self):
+    def get_file_transfer_connection(self) -> FileTransferClient | None:
         assert self.client is not None
         transport = self.client.get_transport()
         assert transport is not None, "No transport"
         try:
-            return SFTPClient.from_transport(transport)
+            if self.data["ssh_file_transfer_protocol"] == "sftp":
+                logger.debug("Using SFTP for file transfer")
+                return SFTPClient.from_transport(transport)
+            elif self.data["ssh_file_transfer_protocol"] == "scp":
+                logger.debug("Using SCP for file transfer")
+                return SCPClient(transport)
+            else:
+                raise ConnectError(
+                    "Unsupported file transfer protocol: {0}".format(
+                        self.data["ssh_file_transfer_protocol"],
+                    ),
+                )
         except SSHException as e:
+
             raise ConnectError(
                 (
                     "Unable to establish SFTP connection. Check that the SFTP subsystem "
                     "for the SSH service at {0} is enabled."
                 ).format(self.host),
             ) from e
 
-    def _get_file(self, remote_filename: str, filename_or_io):
+    def _get_file(self, remote_filename: str, filename_or_io: str | IO):
         with get_file_io(filename_or_io, "wb") as file_io:
-            sftp = self.get_sftp_connection()
+            sftp = self.get_file_transfer_connection()
             sftp.getfo(remote_filename, file_io)
 
     @override
@@ -448,7 +480,7 @@ def _put_file(self, filename_or_io, remote_location):
         while attempts < 3:
             try:
                 with get_file_io(filename_or_io) as file_io:
-                    sftp = self.get_sftp_connection()
+                    sftp = self.get_file_transfer_connection()
                     sftp.putfo(file_io, remote_location)
                 return
             except OSError as e:

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+from .client import SCPClient # noqa: F401`