From 1283689ed444d4150b597707a63a0bf10c7115be Mon Sep 17 00:00:00 2001 From: Ruslan Kuprieiev Date: Mon, 27 Aug 2018 09:09:49 +0300 Subject: [PATCH 1/2] ssh: add ability to configure port Fixes #1060 Signed-off-by: Ruslan Kuprieiev --- dvc/config.py | 2 ++ dvc/dependency/ssh.py | 2 ++ dvc/remote/ssh.py | 43 ++++++++++++++++++++++++++++++++----------- 3 files changed, 36 insertions(+), 11 deletions(-) diff --git a/dvc/config.py b/dvc/config.py index 352065fe3..bf8fd6b31 100644 --- a/dvc/config.py +++ b/dvc/config.py @@ -132,6 +132,7 @@ class Config(object): SECTION_REMOTE_FMT = 'remote "{}"' SECTION_REMOTE_URL = 'url' SECTION_REMOTE_USER = 'user' + SECTION_REMOTE_PORT = 'port' SECTION_REMOTE_SCHEMA = { SECTION_REMOTE_URL: And(supported_url, error="Unsupported URL"), Optional(SECTION_AWS_REGION): str, @@ -141,6 +142,7 @@ class Config(object): Optional(SECTION_GCP_PROJECTNAME): str, Optional(SECTION_CACHE_TYPE): SECTION_CACHE_TYPE_SCHEMA, Optional(SECTION_REMOTE_USER): str, + Optional(SECTION_REMOTE_PORT): int, } SCHEMA = { diff --git a/dvc/dependency/ssh.py b/dvc/dependency/ssh.py index a6d555062..7ff346354 100644 --- a/dvc/dependency/ssh.py +++ b/dvc/dependency/ssh.py @@ -18,6 +18,7 @@ class DependencySSH(DependencyBase): self.remote = remote if remote else RemoteSSH(stage.project, {}) host = remote.host if remote else self.group('host') + port = remote.port if remote else RemoteSSH.DEFAULT_PORT user = remote.user if remote else self.group('user') if remote: path = posixpath.join(remote.prefix, @@ -27,6 +28,7 @@ class DependencySSH(DependencyBase): self.path_info = {'scheme': 'ssh', 'host': host, + 'port': port, 'user': user, 'path': path} diff --git a/dvc/remote/ssh.py b/dvc/remote/ssh.py index 10236918d..7ef79c363 100644 --- a/dvc/remote/ssh.py +++ b/dvc/remote/ssh.py @@ -46,6 +46,8 @@ class RemoteSSH(RemoteBase): REQUIRES = {'paramiko': paramiko} PARAM_MD5 = 'md5' + DEFAULT_PORT = 22 + def __init__(self, project, config): self.project = project self.url = config.get(Config.SECTION_REMOTE_URL, '/') @@ -55,30 +57,35 @@ class RemoteSSH(RemoteBase): self.user = config.get(Config.SECTION_REMOTE_USER, getpass.getuser()) self.prefix = self.group('path') + self.port = config.get(Config.SECTION_REMOTE_PORT, self.DEFAULT_PORT) def md5s_to_path_infos(self, md5s): return [{'scheme': 'ssh', 'host': self.host, 'user': self.user, + 'port': self.port, 'path': posixpath.join(self.prefix, md5[0:2], md5[2:])} for md5 in md5s] - def ssh(self, host=None, user=None): - msg = "Establishing ssh connection with '{}' as user '{}'" - Logger.debug(msg.format(host, user)) + 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()) - ssh.connect(host, username=user) + ssh.connect(host, username=user, port=port) return ssh def exists(self, path_infos): ret = [] - ssh = self.ssh(host=self.host, user=self.user) + 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) plist = stdout.split() @@ -103,7 +110,9 @@ class RemoteSSH(RemoteBase): if path_info['scheme'] != 'ssh': raise NotImplementedError - ssh = self.ssh(host=path_info['host'], user=path_info['user']) + 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() @@ -133,7 +142,8 @@ class RemoteSSH(RemoteBase): assert from_info['user'] == to_info['user'] s = ssh if ssh else self.ssh(host=from_info['host'], - user=from_info['user']) + user=from_info['user'], + port=from_info['port']) dname = posixpath.dirname(to_info['path']) self._exec(s, 'mkdir -p {}'.format(dname)) @@ -159,6 +169,7 @@ class RemoteSSH(RemoteBase): cache = {} cache['scheme'] = 'ssh' cache['host'] = self.host + cache['port'] = self.port cache['user'] = self.user cache['path'] = posixpath.join(self.prefix, md5[0:2], md5[2:]) @@ -236,7 +247,9 @@ class RemoteSSH(RemoteBase): path_info['host'], path_info['path'])) - ssh = self.ssh(host=path_info['host'], user=path_info['user']) + ssh = self.ssh(host=path_info['host'], + user=path_info['user'], + port=path_info['port']) ssh.open_sftp().remove(path_info['path']) ssh.close() @@ -247,7 +260,9 @@ class RemoteSSH(RemoteBase): 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']) + 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': @@ -255,6 +270,7 @@ class RemoteSSH(RemoteBase): 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.cp(from_info, to_info, ssh=ssh) continue @@ -291,7 +307,9 @@ class RemoteSSH(RemoteBase): 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']) + 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): @@ -332,7 +350,9 @@ class RemoteSSH(RemoteBase): return posixpath.dirname(relpath) + posixpath.basename(relpath) def _all_md5s(self): - ssh = self.ssh(host=self.host, user=self.user) + 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() @@ -351,6 +371,7 @@ class RemoteSSH(RemoteBase): path_info = {'scheme': 'ssh', 'user': self.user, 'host': self.host, + 'port': self.port, 'path': posixpath.join(self.prefix, md5[0:2], md5[2:])} self.remove(path_info) From c46bba68894e18000bc6e730cc47abc5834aebc8 Mon Sep 17 00:00:00 2001 From: Ruslan Kuprieiev Date: Mon, 27 Aug 2018 12:54:47 +0300 Subject: [PATCH 2/2] ssh: configure private key file Fixes #1059 Signed-off-by: Ruslan Kuprieiev --- dvc/config.py | 2 ++ dvc/remote/ssh.py | 6 +++++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/dvc/config.py b/dvc/config.py index bf8fd6b31..5f2e2d8cb 100644 --- a/dvc/config.py +++ b/dvc/config.py @@ -133,6 +133,7 @@ class Config(object): SECTION_REMOTE_URL = 'url' SECTION_REMOTE_USER = 'user' SECTION_REMOTE_PORT = 'port' + SECTION_REMOTE_KEY_FILE = 'keyfile' SECTION_REMOTE_SCHEMA = { SECTION_REMOTE_URL: And(supported_url, error="Unsupported URL"), Optional(SECTION_AWS_REGION): str, @@ -143,6 +144,7 @@ class Config(object): Optional(SECTION_CACHE_TYPE): SECTION_CACHE_TYPE_SCHEMA, Optional(SECTION_REMOTE_USER): str, Optional(SECTION_REMOTE_PORT): int, + Optional(SECTION_REMOTE_KEY_FILE): str, } SCHEMA = { diff --git a/dvc/remote/ssh.py b/dvc/remote/ssh.py index 7ef79c363..b6b8afdbb 100644 --- a/dvc/remote/ssh.py +++ b/dvc/remote/ssh.py @@ -58,6 +58,7 @@ class RemoteSSH(RemoteBase): getpass.getuser()) self.prefix = self.group('path') self.port = config.get(Config.SECTION_REMOTE_PORT, self.DEFAULT_PORT) + self.keyfile = config.get(Config.SECTION_REMOTE_KEY_FILE, None) def md5s_to_path_infos(self, md5s): return [{'scheme': 'ssh', @@ -77,7 +78,10 @@ class RemoteSSH(RemoteBase): ssh.load_system_host_keys() ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy()) - ssh.connect(host, username=user, port=port) + ssh.connect(host, + username=user, + port=port, + key_filename=self.keyfile) return ssh