dvc/tests/remotes/ssh.py

import locale
import os
import subprocess

import mockssh
import pytest
from funcy import cached_property
from mockssh.server import Handler

from dvc.path_info import URLInfo

from .base import Base
from .local import Local

TEST_SSH_USER = "user"
TEST_SSH_KEY_PATH = os.path.join(
    os.path.abspath(os.path.dirname(__file__)), f"{TEST_SSH_USER}.key"
)


# See: https://github.com/carletes/mock-ssh-server/issues/22
class SSHMockHandler(Handler):
    def handle_client(self, channel):
        try:
            command = self.command_queues[channel.chanid].get(block=True)
            self.log.debug("Executing %s", command)
            if isinstance(command, bytes):
                command = command.decode()
            p = subprocess.Popen(
                command,
                shell=True,
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
            )
            stdout, stderr = p.communicate()
            channel.sendall(stdout)
            channel.sendall_stderr(stderr)
            channel.send_exit_status(p.returncode)
        except Exception:  # pylint: disable=broad-except
            self.log.error(
                "Error handling client (channel: %s)", channel, exc_info=True
            )
        finally:
            try:
                channel.close()
            except EOFError:
                self.log.debug("Tried to close already closed channel")


class SSHMockServer(mockssh.Server):
    handler_cls = SSHMockHandler


class SSHMocked(Base, URLInfo):
    @staticmethod
    def get_url(user, port):  # pylint: disable=arguments-differ
        path = Local.get_storagepath()
        if os.name == "nt":
            # NOTE: On Windows Local.get_storagepath() will return an
            # ntpath that looks something like `C:\some\path`, which is not
            # compatible with SFTP paths [1], so we need to convert it to
            # a proper posixpath.
            # To do that, we should construct a posixpath that would be
            # relative to the server's root.
            # Our URL format requires absolute paths, so the
            # resulting path would look like `/some/path`.
            #
            # [1]https://tools.ietf.org/html/draft-ietf-secsh-filexfer-13#section-6
            drive, path = os.path.splitdrive(path)

            # Hackish way to make sure SSH server runs on same drive as tests
            # and temporary directories are.
            # Context: https://github.com/iterative/dvc/pull/4660
            test_drive, _ = os.path.splitdrive(os.getcwd())
            if drive.lower() != test_drive.lower():
                raise Exception("Did you forget to use `tmp_dir?`")

            path = path.replace("\\", "/")
        url = f"ssh://{user}@127.0.0.1:{port}{path}"
        return url

    @cached_property
    def config(self):
        return {"url": self.url, "keyfile": TEST_SSH_KEY_PATH}

    @cached_property
    def _ssh(self):
        from sshfs import SSHFileSystem

        return SSHFileSystem(
            host=self.host,
            port=self.port,
            username=TEST_SSH_USER,
            client_keys=[TEST_SSH_KEY_PATH],
        )

    def is_file(self):
        return self._ssh.isfile(self.path)

    def is_dir(self):
        return self._ssh.isdir(self.path)

    def exists(self):
        return self._ssh.exists(self.path)

    def mkdir(self, mode=0o777, parents=False, exist_ok=False):
        assert mode == 0o777
        assert parents

        self._ssh.makedirs(self.path, exist_ok=exist_ok)

    def write_bytes(self, contents):
        assert isinstance(contents, bytes)
        with self._ssh.open(self.path, "wb") as fobj:
            fobj.write(contents)

    def read_bytes(self):
        with self._ssh.open(self.path, "rb") as fobj:
            return fobj.read()

    def read_text(self, encoding=None, errors=None):
        if not encoding:
            encoding = locale.getpreferredencoding(False)
        assert errors is None
        return self.read_bytes().decode(encoding)


@pytest.fixture
def ssh_server(test_config):
    test_config.requires("ssh")
    users = {TEST_SSH_USER: TEST_SSH_KEY_PATH}
    with SSHMockServer(users) as s:
        yield s


@pytest.fixture
def ssh_connection(ssh_server):
    from sshfs import SSHFileSystem

    yield SSHFileSystem(
        host=ssh_server.host,
        port=ssh_server.port,
        username=TEST_SSH_USER,
        client_files=[TEST_SSH_KEY_PATH],
    )


@pytest.fixture
def ssh(ssh_server, monkeypatch):
    from dvc.fs.ssh import SSHFileSystem

    # NOTE: see http://github.com/iterative/dvc/pull/3501
    monkeypatch.setattr(SSHFileSystem, "CAN_TRAVERSE", False)

    return SSHMocked(SSHMocked.get_url(TEST_SSH_USER, ssh_server.port))