mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-09 07:18:27 +01:00
348 lines
12 KiB
Python
348 lines
12 KiB
Python
import os
|
|
import getpass
|
|
import posixpath
|
|
|
|
try:
|
|
import paramiko
|
|
except ImportError:
|
|
paramiko = None
|
|
|
|
import dvc.prompt as prompt
|
|
import dvc.logger as logger
|
|
|
|
from dvc.config import Config
|
|
from dvc.progress import progress
|
|
from dvc.utils.compat import urlparse
|
|
from dvc.exceptions import DvcException
|
|
from dvc.remote.base import RemoteBase, RemoteBaseCmdError
|
|
|
|
|
|
def sizeof_fmt(num, suffix='B'):
|
|
""" Convert number of bytes to human-readable string """
|
|
for unit in ['', 'K', 'M', 'G', 'T', 'P', 'E', 'Z']:
|
|
if abs(num) < 1024.0:
|
|
return "%3.1f%s%s" % (num, unit, suffix)
|
|
num /= 1024.0
|
|
return "%.1f%s%s" % (num, 'Y', suffix)
|
|
|
|
|
|
def percent_cb(name, complete, total):
|
|
""" Callback for updating target progress """
|
|
logger.debug('{}: {} transferred out of {}'.format(name,
|
|
sizeof_fmt(complete),
|
|
sizeof_fmt(total)))
|
|
progress.update_target(name, complete, total)
|
|
|
|
|
|
def create_cb(name):
|
|
""" Create callback function for multipart object """
|
|
return (lambda cur, tot: percent_cb(name, cur, tot))
|
|
|
|
|
|
class RemoteSSHCmdError(RemoteBaseCmdError):
|
|
pass
|
|
|
|
|
|
class RemoteSSH(RemoteBase):
|
|
scheme = 'ssh'
|
|
|
|
# NOTE: we support both URL-like (ssh://[user@]host.xz[:port]/path) and
|
|
# SCP-like (ssh://[user@]host.xz:path) urls.
|
|
REGEX = r'^ssh://.*$'
|
|
|
|
REQUIRES = {'paramiko': paramiko}
|
|
|
|
JOBS = 4
|
|
|
|
PARAM_CHECKSUM = 'md5'
|
|
|
|
DEFAULT_PORT = 22
|
|
TIMEOUT = 1800
|
|
|
|
def __init__(self, project, config):
|
|
super(RemoteSSH, self).__init__(project, config)
|
|
self.url = config.get(Config.SECTION_REMOTE_URL, 'ssh://')
|
|
|
|
parsed = urlparse(self.url)
|
|
self.host = parsed.hostname
|
|
self.user = (config.get(Config.SECTION_REMOTE_USER)
|
|
or parsed.username
|
|
or getpass.getuser())
|
|
self.prefix = parsed.path
|
|
self.port = (config.get(Config.SECTION_REMOTE_PORT)
|
|
or parsed.port
|
|
or self.DEFAULT_PORT)
|
|
self.keyfile = config.get(Config.SECTION_REMOTE_KEY_FILE, None)
|
|
self.timeout = config.get(Config.SECTION_REMOTE_TIMEOUT, self.TIMEOUT)
|
|
self.password = config.get(Config.SECTION_REMOTE_PASSWORD, None)
|
|
self.ask_password = config.get(Config.SECTION_REMOTE_ASK_PASSWORD,
|
|
False)
|
|
|
|
self.path_info = {'scheme': 'ssh',
|
|
'host': self.host,
|
|
'user': self.user,
|
|
'port': self.port}
|
|
|
|
def ssh(self, host=None, user=None, port=None):
|
|
msg = "Establishing ssh connection with '{}' " \
|
|
"through port '{}' as user '{}'"
|
|
logger.debug(msg.format(host, port, user))
|
|
|
|
ssh = paramiko.SSHClient()
|
|
|
|
ssh.load_system_host_keys()
|
|
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
|
|
|
if self.ask_password and self.password is None:
|
|
msg = ("Enter a private key passphrase or a password for "
|
|
"host '{}' port '{}' user '{}'").format(host, port, user)
|
|
self.password = prompt.password(msg)
|
|
|
|
ssh.connect(host,
|
|
username=user,
|
|
port=port,
|
|
key_filename=self.keyfile,
|
|
timeout=self.timeout,
|
|
password=self.password)
|
|
|
|
return ssh
|
|
|
|
def exists(self, path_info):
|
|
assert not isinstance(path_info, list)
|
|
assert path_info['scheme'] == 'ssh'
|
|
|
|
with self.ssh(path_info['host'],
|
|
path_info['user'],
|
|
path_info['port']) as ssh:
|
|
try:
|
|
self._exec(ssh, 'test -e {}'.format(path_info['path']))
|
|
exists = True
|
|
except RemoteSSHCmdError:
|
|
exists = False
|
|
|
|
return exists
|
|
|
|
def _exec(self, ssh, cmd):
|
|
stdin, stdout, stderr = ssh.exec_command(cmd)
|
|
channel = stdout.channel
|
|
|
|
stdin.close()
|
|
channel.shutdown_write()
|
|
|
|
stdout_chunks = []
|
|
stderr_chunks = []
|
|
while (not channel.closed
|
|
or channel.recv_ready()
|
|
or channel.recv_stderr_ready()):
|
|
import select
|
|
got_chunk = False
|
|
readq, _, _ = select.select([stdout.channel], [], [], self.timeout)
|
|
for c in readq:
|
|
if c.recv_ready():
|
|
stdout_chunks.append(stdout.channel.recv(len(c.in_buffer)))
|
|
got_chunk = True
|
|
|
|
if c.recv_stderr_ready():
|
|
stderr_len = len(c.in_stderr_buffer)
|
|
s = stderr.channel.recv_stderr(stderr_len)
|
|
stderr_chunks.append(s)
|
|
got_chunk = True
|
|
|
|
if not got_chunk \
|
|
and stdout.channel.exit_status_ready() \
|
|
and not stderr.channel.recv_stderr_ready() \
|
|
and not stdout.channel.recv_ready():
|
|
stdout.channel.shutdown_read()
|
|
stdout.channel.close()
|
|
break
|
|
|
|
stdout.close()
|
|
stderr.close()
|
|
|
|
ret = stdout.channel.recv_exit_status()
|
|
if ret != 0:
|
|
err = b''.join(stderr_chunks).decode('utf-8')
|
|
raise RemoteSSHCmdError(cmd, ret, err)
|
|
|
|
return b''.join(stdout_chunks).decode('utf-8')
|
|
|
|
def md5(self, path_info):
|
|
if path_info['scheme'] != 'ssh':
|
|
raise NotImplementedError
|
|
|
|
ssh = self.ssh(host=path_info['host'],
|
|
user=path_info['user'],
|
|
port=path_info['port'])
|
|
|
|
# Use different md5 commands depending on os
|
|
stdout = self._exec(ssh, 'uname').strip()
|
|
if stdout == 'Darwin':
|
|
md5cmd = 'md5'
|
|
index = -1
|
|
elif stdout == 'Linux':
|
|
md5cmd = 'md5sum'
|
|
index = 0
|
|
else:
|
|
msg = "'{}' is not supported as a remote".format(stdout)
|
|
raise DvcException(msg)
|
|
|
|
stdout = self._exec(ssh, '{} {}'.format(md5cmd, path_info['path']))
|
|
md5 = stdout.split()[index]
|
|
ssh.close()
|
|
|
|
assert len(md5) == 32
|
|
|
|
return md5
|
|
|
|
def copy(self, from_info, to_info, ssh=None):
|
|
if from_info['scheme'] != 'ssh' or to_info['scheme'] != 'ssh':
|
|
raise NotImplementedError
|
|
|
|
assert from_info['host'] == to_info['host']
|
|
assert from_info['user'] == to_info['user']
|
|
|
|
s = ssh if ssh else self.ssh(host=from_info['host'],
|
|
user=from_info['user'],
|
|
port=from_info['port'])
|
|
|
|
dname = posixpath.dirname(to_info['path'])
|
|
self._exec(s, 'mkdir -p {}'.format(dname))
|
|
self._exec(s, 'cp {} {}'.format(from_info['path'], to_info['path']))
|
|
|
|
if not ssh:
|
|
s.close()
|
|
|
|
def save_info(self, path_info):
|
|
if path_info['scheme'] != 'ssh':
|
|
raise NotImplementedError
|
|
|
|
return {self.PARAM_CHECKSUM: self.md5(path_info)}
|
|
|
|
def save(self, path_info):
|
|
if path_info['scheme'] != 'ssh':
|
|
raise NotImplementedError
|
|
|
|
md5 = self.md5(path_info)
|
|
dest = path_info.copy()
|
|
dest['path'] = self.checksum_to_path(md5)
|
|
|
|
self.copy(path_info, dest)
|
|
|
|
return {self.PARAM_CHECKSUM: md5}
|
|
|
|
def remove(self, path_info):
|
|
if path_info['scheme'] != 'ssh':
|
|
raise NotImplementedError
|
|
|
|
logger.debug('Removing ssh://{}@{}/{}'.format(path_info['user'],
|
|
path_info['host'],
|
|
path_info['path']))
|
|
|
|
ssh = self.ssh(host=path_info['host'],
|
|
user=path_info['user'],
|
|
port=path_info['port'])
|
|
ssh.open_sftp().remove(path_info['path'])
|
|
ssh.close()
|
|
|
|
def download(self,
|
|
from_infos,
|
|
to_infos,
|
|
no_progress_bar=False,
|
|
names=None):
|
|
names = self._verify_path_args(from_infos, to_infos, names)
|
|
|
|
ssh = self.ssh(host=from_infos[0]['host'],
|
|
user=from_infos[0]['user'],
|
|
port=from_infos[0]['port'])
|
|
|
|
for to_info, from_info, name in zip(to_infos, from_infos, names):
|
|
if from_info['scheme'] != 'ssh':
|
|
raise NotImplementedError
|
|
|
|
if to_info['scheme'] == 'ssh':
|
|
assert from_info['host'] == to_info['host']
|
|
assert from_info['port'] == to_info['port']
|
|
assert from_info['user'] == to_info['user']
|
|
self.copy(from_info, to_info, ssh=ssh)
|
|
continue
|
|
|
|
if to_info['scheme'] != 'local':
|
|
raise NotImplementedError
|
|
|
|
msg = "Downloading '{}/{}' to '{}'".format(from_info['host'],
|
|
from_info['path'],
|
|
to_info['path'])
|
|
logger.debug(msg)
|
|
|
|
if not name:
|
|
name = os.path.basename(to_info['path'])
|
|
|
|
self._makedirs(to_info['path'])
|
|
tmp_file = self.tmp_file(to_info['path'])
|
|
try:
|
|
ssh.open_sftp().get(from_info['path'],
|
|
tmp_file,
|
|
callback=create_cb(name))
|
|
except Exception:
|
|
msg = "failed to download '{}/{}' to '{}'"
|
|
logger.error(msg.format(from_info['host'],
|
|
from_info['path'],
|
|
to_info['path']))
|
|
continue
|
|
|
|
os.rename(tmp_file, to_info['path'])
|
|
progress.finish_target(name)
|
|
|
|
ssh.close()
|
|
|
|
def upload(self, from_infos, to_infos, names=None):
|
|
names = self._verify_path_args(to_infos, from_infos, names)
|
|
|
|
ssh = self.ssh(host=to_infos[0]['host'],
|
|
user=to_infos[0]['user'],
|
|
port=to_infos[0]['port'])
|
|
sftp = ssh.open_sftp()
|
|
|
|
for from_info, to_info, name in zip(from_infos, to_infos, names):
|
|
if to_info['scheme'] != 'ssh':
|
|
raise NotImplementedError
|
|
|
|
if from_info['scheme'] != 'local':
|
|
raise NotImplementedError
|
|
|
|
logger.debug("Uploading '{}' to '{}/{}'".format(from_info['path'],
|
|
to_info['host'],
|
|
to_info['path']))
|
|
|
|
if not name:
|
|
name = os.path.basename(from_info['path'])
|
|
|
|
dname = posixpath.dirname(to_info['path'])
|
|
self._exec(ssh, 'mkdir -p {}'.format(dname))
|
|
|
|
try:
|
|
sftp.put(from_info['path'],
|
|
to_info['path'],
|
|
callback=create_cb(name))
|
|
except Exception:
|
|
msg = "failed to upload '{}' to '{}/{}'"
|
|
logger.error(msg.format(from_info['path'],
|
|
to_info['host'],
|
|
to_info['path']))
|
|
continue
|
|
|
|
progress.finish_target(name)
|
|
|
|
sftp.close()
|
|
ssh.close()
|
|
|
|
def list_cache_paths(self):
|
|
ssh = self.ssh(host=self.host,
|
|
user=self.user,
|
|
port=self.port)
|
|
cmd = 'find {} -type f -follow -print'.format(self.prefix)
|
|
stdout = self._exec(ssh, cmd)
|
|
flist = stdout.split()
|
|
ssh.close()
|
|
return flist
|