mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-10 07:18:27 +01:00
271 lines
8.5 KiB
Python
271 lines
8.5 KiB
Python
import getpass
|
|
import io
|
|
import logging
|
|
import os
|
|
import threading
|
|
from contextlib import closing, contextmanager
|
|
from urllib.parse import urlparse
|
|
|
|
from funcy import first, memoize, silent, wrap_with
|
|
|
|
import dvc.prompt as prompt
|
|
from dvc.hash_info import HashInfo
|
|
from dvc.scheme import Schemes
|
|
|
|
from ..base import BaseTree
|
|
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 SSHTree(BaseTree):
|
|
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, repo, config):
|
|
super().__init__(repo, config)
|
|
url = config.get("url")
|
|
if url:
|
|
parsed = urlparse(url)
|
|
user_ssh_config = self._load_user_ssh_config(parsed.hostname)
|
|
|
|
host = user_ssh_config.get("hostname", parsed.hostname)
|
|
user = (
|
|
config.get("user")
|
|
or parsed.username
|
|
or user_ssh_config.get("user")
|
|
or getpass.getuser()
|
|
)
|
|
port = (
|
|
config.get("port")
|
|
or parsed.port
|
|
or self._try_get_ssh_config_port(user_ssh_config)
|
|
or self.DEFAULT_PORT
|
|
)
|
|
self.path_info = self.PATH_CLS.from_parts(
|
|
scheme=self.scheme,
|
|
host=host,
|
|
user=user,
|
|
port=port,
|
|
path=parsed.path,
|
|
)
|
|
else:
|
|
self.path_info = None
|
|
user_ssh_config = {}
|
|
|
|
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 ssh_config_filename():
|
|
return os.path.expanduser(os.path.join("~", ".ssh", "config"))
|
|
|
|
@staticmethod
|
|
def _load_user_ssh_config(hostname):
|
|
import paramiko
|
|
|
|
user_config_file = SSHTree.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, path_info=None):
|
|
if path_info is None:
|
|
path_info = self.path_info
|
|
|
|
# NOTE: we use the same password regardless of the server :(
|
|
if self.ask_password and self.password is None:
|
|
host, user, port = path_info.host, path_info.user, path_info.port
|
|
self.password = ask_password(host, user, port)
|
|
|
|
def ssh(self, path_info):
|
|
self.ensure_credentials(path_info)
|
|
|
|
from .connection import SSHConnection
|
|
|
|
return get_connection(
|
|
SSHConnection,
|
|
path_info.host,
|
|
username=path_info.user,
|
|
port=path_info.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):
|
|
assert mode in {"r", "rt", "rb", "wb"}
|
|
|
|
with self.ssh(path_info) 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, use_dvcignore=True):
|
|
with self.ssh(path_info) as ssh:
|
|
return ssh.exists(path_info.path)
|
|
|
|
def isdir(self, path_info):
|
|
with self.ssh(path_info) as ssh:
|
|
return ssh.isdir(path_info.path)
|
|
|
|
def isfile(self, path_info):
|
|
with self.ssh(path_info) as ssh:
|
|
return ssh.isfile(path_info.path)
|
|
|
|
def walk_files(self, path_info, **kwargs):
|
|
with self.ssh(path_info) 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(path_info) as ssh:
|
|
ssh.remove(path_info.path)
|
|
|
|
def makedirs(self, path_info):
|
|
with self.ssh(path_info) as ssh:
|
|
ssh.makedirs(path_info.path)
|
|
|
|
def move(self, from_info, to_info, mode=None):
|
|
assert mode is None
|
|
if from_info.scheme != self.scheme or to_info.scheme != self.scheme:
|
|
raise NotImplementedError
|
|
|
|
with self.ssh(from_info) 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(from_info) 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(from_info) 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(to_info) 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(from_info) 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(from_info) as ssh:
|
|
ssh.reflink(from_info.path, to_info.path)
|
|
|
|
def get_file_hash(self, path_info):
|
|
if path_info.scheme != self.scheme:
|
|
raise NotImplementedError
|
|
|
|
with self.ssh(path_info) as ssh:
|
|
hash_info = HashInfo(self.PARAM_CHECKSUM, ssh.md5(path_info.path),)
|
|
|
|
if hash_info:
|
|
hash_info.size = ssh.getsize(path_info.path)
|
|
|
|
return hash_info
|
|
|
|
def getsize(self, path_info):
|
|
with self.ssh(path_info) as ssh:
|
|
return ssh.getsize(path_info.path)
|
|
|
|
def _download(self, from_info, to_file, name=None, no_progress_bar=False):
|
|
with self.ssh(from_info) 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):
|
|
with self.ssh(to_info) as ssh:
|
|
ssh.upload(
|
|
from_file,
|
|
to_info.path,
|
|
progress_title=name,
|
|
no_progress_bar=no_progress_bar,
|
|
)
|