From 682737e4a20cf28d0e79b486df7e284153f87f99 Mon Sep 17 00:00:00 2001
From: Ruslan Kuprieiev <ruslan@iterative.ai>
Date: Wed, 19 Sep 2018 05:50:06 +0300
Subject: [PATCH] ssh: add 'password' and 'ask_password' options

Fixes #1124

Signed-off-by: Ruslan Kuprieiev <ruslan@iterative.ai>
---
 dvc/config.py       |  4 ++++
 dvc/progress.py     |  6 +++++-
 dvc/prompt.py       | 15 +++++++++++++++
 dvc/remote/local.py |  6 +++---
 dvc/remote/ssh.py   | 10 +++++++++-
 5 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/dvc/config.py b/dvc/config.py
index d8abad665..6a1fff22b 100644
--- a/dvc/config.py
+++ b/dvc/config.py
@@ -143,6 +143,8 @@ class Config(object):
     SECTION_REMOTE_PORT = 'port'
     SECTION_REMOTE_KEY_FILE = 'keyfile'
     SECTION_REMOTE_TIMEOUT = 'timeout'
+    SECTION_REMOTE_PASSWORD = 'password'
+    SECTION_REMOTE_ASK_PASSWORD = 'ask_password'
     SECTION_REMOTE_SCHEMA = {
         SECTION_REMOTE_URL: And(supported_url, error="Unsupported URL"),
         Optional(SECTION_AWS_REGION): str,
@@ -155,6 +157,8 @@ class Config(object):
         Optional(SECTION_REMOTE_PORT): Use(int),
         Optional(SECTION_REMOTE_KEY_FILE): str,
         Optional(SECTION_REMOTE_TIMEOUT): Use(int),
+        Optional(SECTION_REMOTE_PASSWORD): str,
+        Optional(SECTION_REMOTE_ASK_PASSWORD): And(str, is_bool, Use(to_bool)),
     }
 
     SECTION_STATE = 'state'
diff --git a/dvc/progress.py b/dvc/progress.py
index c6445d85b..27b73fc4a 100644
--- a/dvc/progress.py
+++ b/dvc/progress.py
@@ -16,6 +16,10 @@ class Progress(object):
         self._n_total = total
         self._n_finished = 0
 
+    @property
+    def is_finished(self):
+        return self._n_total == self._n_finished
+
     def _clearln(self):
         print('\r\x1b[K', end='')
 
@@ -63,7 +67,7 @@ class Progress(object):
             progress = int((100 * current)/total) if current < total else 100
             percent = str(progress) + "% "
 
-        if self._n_total:
+        if self._n_total > 1:
             num = "({}/{}): ".format(self._n_finished + 1, self._n_total)
         else:
             num = ""
diff --git a/dvc/prompt.py b/dvc/prompt.py
index bbece3ce6..1144d4316 100644
--- a/dvc/prompt.py
+++ b/dvc/prompt.py
@@ -1,5 +1,7 @@
 import sys
 
+from dvc.progress import progress
+
 try:
     # NOTE: in Python3 raw_input() was renamed to input()
     input = raw_input
@@ -23,3 +25,16 @@ class Prompt(object):
             answer = input('Enter \'yes\' or \'no\'.\n').lower()
 
         return answer[0] == "y"
+
+    def prompt_password(self, msg):  # pragma: no cover
+        import getpass
+
+        if not sys.stdout.isatty():
+            return None
+
+        msg = 'Enter password for {}:\n'.format(msg)
+
+        if not progress.is_finished:
+            msg = u'\n' + msg
+
+        return getpass.getpass(msg)
diff --git a/dvc/remote/local.py b/dvc/remote/local.py
index a30bad54e..0f6e40a83 100644
--- a/dvc/remote/local.py
+++ b/dvc/remote/local.py
@@ -491,7 +491,7 @@ class RemoteLOCAL(RemoteBase):
         Logger.info("Preparing to pull data from {}".format(remote.url))
         title = "Collecting information"
 
-        progress.set_n_total(0)
+        progress.set_n_total(1)
         progress.update_target(title, 0, 100)
 
         checksum_infos, missing = self._collect(checksum_infos)
@@ -527,7 +527,7 @@ class RemoteLOCAL(RemoteBase):
                  show_checksums=False):
         title = "Collecting information"
 
-        progress.set_n_total(0)
+        progress.set_n_total(1)
         progress.update_target(title, 0, 100)
 
         grouped = zip(*self._group(checksum_infos,
@@ -599,7 +599,7 @@ class RemoteLOCAL(RemoteBase):
         Logger.info("Preparing to push data to {}".format(remote.url))
         title = "Collecting information"
 
-        progress.set_n_total(0)
+        progress.set_n_total(1)
         progress.update_target(title, 0, 100)
 
         checksum_infos = self._collect(checksum_infos)[0]
diff --git a/dvc/remote/ssh.py b/dvc/remote/ssh.py
index 39d8061ba..1ec1ee93a 100644
--- a/dvc/remote/ssh.py
+++ b/dvc/remote/ssh.py
@@ -61,6 +61,9 @@ class RemoteSSH(RemoteBase):
         self.port = config.get(Config.SECTION_REMOTE_PORT, 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)
 
     def md5s_to_path_infos(self, md5s):
         return [{'scheme': 'ssh',
@@ -80,11 +83,16 @@ class RemoteSSH(RemoteBase):
         ssh.load_system_host_keys()
         ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
 
+        if self.ask_password and self.password is None:
+            msg = "host '{}' user '{}' port '{}'".format(host, port, user)
+            self.password = self.project.prompt.prompt_password(msg)
+
         ssh.connect(host,
                     username=user,
                     port=port,
                     key_filename=self.keyfile,
-                    timeout=self.timeout)
+                    timeout=self.timeout,
+                    password=self.password)
 
         return ssh