dvc/dvc/fs/ssh/__init__.py

import getpass
import io
import logging
import os
import shutil
import threading
from contextlib import closing, contextmanager

from funcy import first, memoize, silent, wrap_with

from dvc import prompt
from dvc.hash_info import HashInfo
from dvc.scheme import Schemes

from ..base import BaseFileSystem
from ..pool import get_connection

logger = logging.getLogger(__name__)


@wrap_with(threading.Lock())
@memoize
def ask_password(host, user, port):
    return prompt.password(
        "Enter a private key passphrase or a password for "
        "host '{host}' port '{port}' user '{user}'".format(
            host=host, port=port, user=user
        )
    )


class SSHFileSystem(BaseFileSystem):  # pylint:disable=abstract-method
    scheme = Schemes.SSH
    REQUIRES = {"paramiko": "paramiko"}
    _JOBS = 4

    PARAM_CHECKSUM = "md5"
    # At any given time some of the connections will go over network and
    # paramiko stuff, so we would ideally have it double of server processors.
    # We use conservative setting of 4 instead to not exhaust max sessions.
    CHECKSUM_JOBS = 4
    DEFAULT_CACHE_TYPES = ["copy"]
    TRAVERSE_PREFIX_LEN = 2

    DEFAULT_PORT = 22
    TIMEOUT = 1800

    def __init__(self, **config):
        super().__init__(**config)
        user_ssh_config = self._load_user_ssh_config(config["host"])

        self.host = user_ssh_config.get("hostname", config["host"])
        self.user = (
            config.get("user")
            or user_ssh_config.get("user")
            or getpass.getuser()
        )
        self.port = (
            config.get("port")
            or self._try_get_ssh_config_port(user_ssh_config)
            or self.DEFAULT_PORT
        )

        self.keyfile = config.get(
            "keyfile"
        ) or self._try_get_ssh_config_keyfile(user_ssh_config)
        self.timeout = config.get("timeout", self.TIMEOUT)
        self.password = config.get("password", None)
        self.ask_password = config.get("ask_password", False)
        self.gss_auth = config.get("gss_auth", False)
        proxy_command = user_ssh_config.get("proxycommand", False)
        if proxy_command:
            import paramiko

            self.sock = paramiko.ProxyCommand(proxy_command)
        else:
            self.sock = None
        self.allow_agent = config.get("allow_agent", True)

    @staticmethod
    def _get_kwargs_from_urls(urlpath):
        from fsspec.implementations.sftp import SFTPFileSystem

        # pylint:disable=protected-access
        kwargs = SFTPFileSystem._get_kwargs_from_urls(urlpath)
        if "username" in kwargs:
            kwargs["user"] = kwargs.pop("username")
        return kwargs

    @staticmethod
    def ssh_config_filename():
        return os.path.expanduser(os.path.join("~", ".ssh", "config"))

    @staticmethod
    def _load_user_ssh_config(hostname):
        import paramiko

        user_config_file = SSHFileSystem.ssh_config_filename()
        user_ssh_config = {}
        if hostname and os.path.exists(user_config_file):
            ssh_config = paramiko.SSHConfig()
            with open(user_config_file) as f:
                # For whatever reason parsing directly from f is unreliable
                f_copy = io.StringIO(f.read())
                ssh_config.parse(f_copy)
            user_ssh_config = ssh_config.lookup(hostname)
        return user_ssh_config

    @staticmethod
    def _try_get_ssh_config_port(user_ssh_config):
        return silent(int)(user_ssh_config.get("port"))

    @staticmethod
    def _try_get_ssh_config_keyfile(user_ssh_config):
        return first(user_ssh_config.get("identityfile") or ())

    def ensure_credentials(self):
        # NOTE: we use the same password regardless of the server :(
        if self.ask_password and self.password is None:
            self.password = ask_password(self.host, self.user, self.port)

    def ssh(self):
        self.ensure_credentials()

        from .connection import SSHConnection

        return get_connection(
            SSHConnection,
            self.host,
            username=self.user,
            port=self.port,
            key_filename=self.keyfile,
            timeout=self.timeout,
            password=self.password,
            gss_auth=self.gss_auth,
            sock=self.sock,
            allow_agent=self.allow_agent,
        )

    @contextmanager
    def open(self, path_info, mode="r", encoding=None, **kwargs):
        assert mode in {"r", "rt", "rb", "wb"}

        with self.ssh() as ssh, closing(
            ssh.sftp.open(path_info.path, mode)
        ) as fd:
            if "b" in mode:
                yield fd
            else:
                yield io.TextIOWrapper(fd, encoding=encoding)

    def exists(self, path_info) -> bool:
        with self.ssh() as ssh:
            return ssh.exists(path_info.path)

    def isdir(self, path_info):
        with self.ssh() as ssh:
            return ssh.isdir(path_info.path)

    def isfile(self, path_info):
        with self.ssh() as ssh:
            return ssh.isfile(path_info.path)

    def walk_files(self, path_info, **kwargs):
        with self.ssh() as ssh:
            for fname in ssh.walk_files(path_info.path):
                yield path_info.replace(path=fname)

    def remove(self, path_info):
        if path_info.scheme != self.scheme:
            raise NotImplementedError

        with self.ssh() as ssh:
            ssh.remove(path_info.path)

    def makedirs(self, path_info, **kwargs):
        with self.ssh() as ssh:
            ssh.makedirs(path_info.path)

    def move(self, from_info, to_info):
        if from_info.scheme != self.scheme or to_info.scheme != self.scheme:
            raise NotImplementedError

        with self.ssh() as ssh:
            ssh.move(from_info.path, to_info.path)

    def copy(self, from_info, to_info):
        if not from_info.scheme == to_info.scheme == self.scheme:
            raise NotImplementedError

        with self.ssh() as ssh:
            ssh.atomic_copy(from_info.path, to_info.path)

    def symlink(self, from_info, to_info):
        if not from_info.scheme == to_info.scheme == self.scheme:
            raise NotImplementedError

        with self.ssh() as ssh:
            ssh.symlink(from_info.path, to_info.path)

    def hardlink(self, from_info, to_info):
        if not from_info.scheme == to_info.scheme == self.scheme:
            raise NotImplementedError

        # See dvc/remote/local/__init__.py - hardlink()
        if self.getsize(from_info) == 0:

            with self.ssh() as ssh:
                ssh.sftp.open(to_info.path, "w").close()

            logger.debug(
                "Created empty file: {src} -> {dest}".format(
                    src=str(from_info), dest=str(to_info)
                )
            )
            return

        with self.ssh() as ssh:
            ssh.hardlink(from_info.path, to_info.path)

    def reflink(self, from_info, to_info):
        if from_info.scheme != self.scheme or to_info.scheme != self.scheme:
            raise NotImplementedError

        with self.ssh() as ssh:
            ssh.reflink(from_info.path, to_info.path)

    def md5(self, path_info):
        with self.ssh() as ssh:
            return HashInfo(
                "md5",
                ssh.md5(path_info.path),
                size=ssh.getsize(path_info.path),
            )

    def info(self, path_info):
        with self.ssh() as ssh:
            return ssh.info(path_info.path)

    def _upload_fobj(self, fobj, to_info, **kwargs):
        self.makedirs(to_info.parent)
        with self.open(to_info, mode="wb") as fdest:
            shutil.copyfileobj(fobj, fdest)

    def _download(self, from_info, to_file, name=None, no_progress_bar=False):
        with self.ssh() as ssh:
            ssh.download(
                from_info.path,
                to_file,
                progress_title=name,
                no_progress_bar=no_progress_bar,
            )

    def _upload(
        self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs
    ):
        with self.ssh() as ssh:
            ssh.upload(
                from_file,
                to_info.path,
                progress_title=name,
                no_progress_bar=no_progress_bar,
            )