From 600514cba95b819ddf7f1fdde7fba2fa8ba21942 Mon Sep 17 00:00:00 2001
From: "Mr. Outis" <mroutis@protonmail.com>
Date: Wed, 20 Feb 2019 23:42:25 -0600
Subject: [PATCH] ssh: upload to a tmp file to prevent data corruption

Fix #1570
---
 dvc/remote/azure.py          |  3 ++-
 dvc/remote/base.py           |  7 -------
 dvc/remote/gs.py             |  3 ++-
 dvc/remote/local.py          |  4 ++--
 dvc/remote/s3.py             |  3 ++-
 dvc/remote/ssh/connection.py | 12 +++++++-----
 dvc/utils/__init__.py        |  7 +++++++
 7 files changed, 22 insertions(+), 17 deletions(-)

diff --git a/dvc/remote/azure.py b/dvc/remote/azure.py
index 43a978d02..3853a9d67 100644
--- a/dvc/remote/azure.py
+++ b/dvc/remote/azure.py
@@ -10,6 +10,7 @@ except ImportError:
     BlockBlobService = None
 
 import dvc.logger as logger
+from dvc.utils import tmp_fname
 from dvc.utils.compat import urlparse, makedirs
 from dvc.progress import progress
 from dvc.config import Config
@@ -170,7 +171,7 @@ class RemoteAzure(RemoteBase):
                 )
             )
 
-            tmp_file = self.tmp_file(to_info["path"])
+            tmp_file = tmp_fname(to_info["path"])
             if not name:
                 name = os.path.basename(to_info["path"])
 
diff --git a/dvc/remote/base.py b/dvc/remote/base.py
index 9952ce824..982e26c79 100644
--- a/dvc/remote/base.py
+++ b/dvc/remote/base.py
@@ -112,13 +112,6 @@ class RemoteBase(object):
             return None
         return m.group(name)
 
-    @staticmethod
-    def tmp_file(fname):
-        """ Temporary name for a partial download """
-        import uuid
-
-        return fname + "." + str(uuid.uuid4())
-
     def save_info(self, path_info):
         raise NotImplementedError
 
diff --git a/dvc/remote/gs.py b/dvc/remote/gs.py
index aa7d8ffd6..79737ec9e 100644
--- a/dvc/remote/gs.py
+++ b/dvc/remote/gs.py
@@ -9,6 +9,7 @@ except ImportError:
     storage = None
 
 import dvc.logger as logger
+from dvc.utils import tmp_fname
 from dvc.utils.compat import urlparse, makedirs
 from dvc.remote.base import RemoteBase
 from dvc.config import Config
@@ -183,7 +184,7 @@ class RemoteGS(RemoteBase):
             )
             logger.debug(msg)
 
-            tmp_file = self.tmp_file(to_info["path"])
+            tmp_file = tmp_fname(to_info["path"])
             if not name:
                 name = os.path.basename(to_info["path"])
 
diff --git a/dvc/remote/local.py b/dvc/remote/local.py
index bc698041b..ae29921ab 100644
--- a/dvc/remote/local.py
+++ b/dvc/remote/local.py
@@ -14,7 +14,7 @@ from operator import itemgetter
 import dvc.logger as logger
 from dvc.system import System
 from dvc.remote.base import RemoteBase, STATUS_MAP, STATUS_NEW, STATUS_DELETED
-from dvc.utils import remove, move, copyfile, dict_md5, to_chunks
+from dvc.utils import remove, move, copyfile, dict_md5, to_chunks, tmp_fname
 from dvc.utils import LARGE_DIR_SIZE
 from dvc.config import Config
 from dvc.exceptions import DvcException
@@ -522,7 +522,7 @@ class RemoteLOCAL(RemoteBase):
                 name = os.path.basename(to_info["path"])
 
             makedirs(os.path.dirname(to_info["path"]), exist_ok=True)
-            tmp_file = self.tmp_file(to_info["path"])
+            tmp_file = tmp_fname(to_info["path"])
             try:
                 copyfile(
                     from_info["path"],
diff --git a/dvc/remote/s3.py b/dvc/remote/s3.py
index e38369d52..a9216ec62 100644
--- a/dvc/remote/s3.py
+++ b/dvc/remote/s3.py
@@ -9,6 +9,7 @@ except ImportError:
     boto3 = None
 
 import dvc.logger as logger
+from dvc.utils import tmp_fname
 from dvc.utils.compat import urlparse, makedirs
 from dvc.progress import progress
 from dvc.config import Config
@@ -222,7 +223,7 @@ class RemoteS3(RemoteBase):
             )
             logger.debug(msg)
 
-            tmp_file = self.tmp_file(to_info["path"])
+            tmp_file = tmp_fname(to_info["path"])
             if not name:
                 name = os.path.basename(to_info["path"])
 
diff --git a/dvc/remote/ssh/connection.py b/dvc/remote/ssh/connection.py
index 10f714f8c..cc28692cd 100644
--- a/dvc/remote/ssh/connection.py
+++ b/dvc/remote/ssh/connection.py
@@ -7,6 +7,7 @@ except ImportError:
     paramiko = None
 
 import dvc.logger as logger
+from dvc.utils import tmp_fname
 from dvc.utils.compat import makedirs
 from dvc.progress import progress
 from dvc.exceptions import DvcException
@@ -145,12 +146,10 @@ class SSHConnection:
         self._sftp.remove(path)
 
     def download(self, src, dest, no_progress_bar=False, progress_title=None):
-        from uuid import uuid4
-
         self._sftp_connect()
 
         makedirs(posixpath.dirname(dest), exist_ok=True)
-        tmp_file = dest + "." + str(uuid4())
+        tmp_file = tmp_fname(dest)
 
         if no_progress_bar:
             self._sftp.get(src, tmp_file)
@@ -167,16 +166,19 @@ class SSHConnection:
         self._sftp_connect()
 
         self.makedirs(posixpath.dirname(dest))
+        tmp_file = tmp_fname(dest)
 
         if no_progress_bar:
-            self._sftp.put(src, dest)
+            self._sftp.put(src, tmp_file)
         else:
             if not progress_title:
                 progress_title = os.path.basename(dest)
 
-            self._sftp.put(src, dest, callback=create_cb(progress_title))
+            self._sftp.put(src, tmp_file, callback=create_cb(progress_title))
             progress.finish_target(progress_title)
 
+        self._sftp.rename(tmp_file, dest)
+
     def execute(self, cmd):
         stdin, stdout, stderr = self._ssh.exec_command(cmd)
         channel = stdout.channel
diff --git a/dvc/utils/__init__.py b/dvc/utils/__init__.py
index bb1fa0321..093e291db 100644
--- a/dvc/utils/__init__.py
+++ b/dvc/utils/__init__.py
@@ -212,3 +212,10 @@ def convert_to_unicode(data):
         return type(data)(map(convert_to_unicode, data))
     else:
         return data
+
+
+def tmp_fname(fname):
+    """ Temporary name for a partial download """
+    from uuid import uuid4
+
+    return fname + "." + str(uuid4()) + ".tmp"