From 3053cf2b1f0b96e4236a02ca0e672d9660b58824 Mon Sep 17 00:00:00 2001
From: Ruslan Kuprieiev <ruslan@iterative.ai>
Date: Tue, 8 Jan 2019 09:58:02 +0200
Subject: [PATCH] ssh: support specifying port in the url

Fixes #1473

Signed-off-by: Ruslan Kuprieiev <ruslan@iterative.ai>
---
 dvc/cli.py               |  3 ++-
 dvc/output/ssh.py        | 12 +++++++----
 dvc/remote/ssh.py        | 25 ++++++++++++++--------
 tests/unit/remote/ssh.py | 45 ++++++++++++++++++++++++++++++++++++++++
 4 files changed, 71 insertions(+), 14 deletions(-)

diff --git a/dvc/cli.py b/dvc/cli.py
index 605fa8957..949a79bc5 100644
--- a/dvc/cli.py
+++ b/dvc/cli.py
@@ -714,7 +714,8 @@ def parse_args(argv=None):
                         help='Name.')
     remote_add_parser.add_argument(
                         'url',
-                        help='URL.')
+                        help='URL. See full list of supported urls at '
+                             'https://dvc.org/doc/commands-reference/remote')
     remote_add_parser.add_argument(
                         '-d',
                         '--default',
diff --git a/dvc/output/ssh.py b/dvc/output/ssh.py
index f26a48bed..6953ae50c 100644
--- a/dvc/output/ssh.py
+++ b/dvc/output/ssh.py
@@ -1,3 +1,4 @@
+import getpass
 import posixpath
 
 from dvc.utils.compat import urlparse
@@ -21,15 +22,18 @@ class OutputSSH(OutputBase):
                                         remote=remote,
                                         cache=cache,
                                         metric=metric)
-        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')
+        parsed = urlparse(path)
+        host = remote.host if remote else parsed.hostname
+        port = remote.port if remote else (parsed.port
+                                           or RemoteSSH.DEFAULT_PORT)
+        user = remote.user if remote else (parsed.username
+                                           or getpass.getuser())
 
         if remote:
             path = posixpath.join(remote.prefix,
                                   urlparse(path).path.lstrip('/'))
         else:
-            path = self.match(self.url).group('path')
+            path = parsed.path
 
         self.path_info = {'scheme': 'ssh',
                           'host': host,
diff --git a/dvc/remote/ssh.py b/dvc/remote/ssh.py
index 04fdf8c3f..2758ef903 100644
--- a/dvc/remote/ssh.py
+++ b/dvc/remote/ssh.py
@@ -9,10 +9,12 @@ except ImportError:
 
 import dvc.prompt as prompt
 import dvc.logger as logger
-from dvc.progress import progress
-from dvc.remote.base import RemoteBase, RemoteBaseCmdError
+
 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'):
@@ -44,8 +46,9 @@ class RemoteSSHCmdError(RemoteBaseCmdError):
 class RemoteSSH(RemoteBase):
     scheme = 'ssh'
 
-    # NOTE: temporarily only absolute paths are allowed
-    REGEX = r'^ssh://((?P<user>.*)@)?(?P<host>[^/]*):(?P<path>/.*)$'
+    # 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}
 
@@ -58,13 +61,17 @@ class RemoteSSH(RemoteBase):
 
     def __init__(self, project, config):
         super(RemoteSSH, self).__init__(project, config)
-        self.url = config.get(Config.SECTION_REMOTE_URL, '/')
-        self.host = self.group('host')
+        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 self.group('user')
+                     or parsed.username
                      or getpass.getuser())
-        self.prefix = self.group('path')
-        self.port = config.get(Config.SECTION_REMOTE_PORT, self.DEFAULT_PORT)
+        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)
diff --git a/tests/unit/remote/ssh.py b/tests/unit/remote/ssh.py
index 0eed57b3a..13aafea48 100644
--- a/tests/unit/remote/ssh.py
+++ b/tests/unit/remote/ssh.py
@@ -22,3 +22,48 @@ class TestRemoteSSH(TestCase):
         config['user'] = user
         remote = RemoteSSH(None, config)
         self.assertEqual(remote.user, user)
+
+    def test_url(self):
+        user = 'test'
+        host = '123.45.67.89'
+        port = 1234
+        path = '/path/to/dir'
+
+        # URL ssh://[user@]host.xz[:port]/path
+        url = 'ssh://{}@{}:{}{}'.format(user, host, port, path)
+        config = {'url': url}
+
+        remote = RemoteSSH(None, config)
+
+        self.assertEqual(remote.user, user)
+        self.assertEqual(remote.host, host)
+        self.assertEqual(remote.port, port)
+        self.assertEqual(remote.prefix, path)
+
+        # SCP-like URL ssh://[user@]host.xz:path
+        url = 'ssh://{}@{}:{}'.format(user, host, path)
+        config = {'url': url}
+
+        remote = RemoteSSH(None, config)
+
+        self.assertEqual(remote.user, user)
+        self.assertEqual(remote.host, host)
+        self.assertEqual(remote.port, remote.DEFAULT_PORT)
+        self.assertEqual(remote.prefix, path)
+
+    def test_port(self):
+        url = 'ssh://127.0.0.1/path/to/dir'
+        config = {'url': url}
+        remote = RemoteSSH(None, config)
+        self.assertEqual(remote.port, remote.DEFAULT_PORT)
+
+        port = 1234
+        url = 'ssh://127.0.0.1:{}/path/to/dir'.format(port)
+        config = {'url': url}
+        remote = RemoteSSH(None, config)
+        self.assertEqual(remote.port, port)
+
+        port = 4321
+        config['port'] = port
+        remote = RemoteSSH(None, config)
+        self.assertEqual(remote.port, port)