T3563: remote: Allow `paramiko.transport.Transport().close()` to automatically close the socket

1 files changed, 23 insertions, 28 deletions

diff --git a/python/vyos/remote.py b/python/vyos/remote.py
index 05f739dc8..80a4e7528 100644
--- a/python/vyos/remote.py
+++ b/python/vyos/remote.py
@@ -144,35 +144,30 @@ def transfer_sftp(mode, local_path, hostname, remote_path,\
         sock.bind((source, 0))
         sock.connect((hostname, port))
     callback = make_progressbar() if progressbar else None
-    try:
-        with SSHClient() as ssh:
-            ssh.load_system_host_keys()
-            if os.path.exists(KNOWN_HOSTS_FILE):
-                ssh.load_host_keys(KNOWN_HOSTS_FILE)
-            ssh.set_missing_host_key_policy(InteractivePolicy())
-            ssh.connect(hostname, port, username, password, sock=sock)
-            with ssh.open_sftp() as sftp:
-                if mode == 'upload':
-                    try:
-                        # If the remote path is a directory, use the original filename.
-                        if stat.S_ISDIR(sftp.stat(remote_path).st_mode):
-                            path = os.path.join(remote_path, os.path.basename(local_path))
-                        # A file exists at this destination. We're simply going to clobber it.
-                        else:
-                            path = remote_path
-                    # This path doesn't point at any existing file. We can freely use this filename.
-                    except IOError:
+    with SSHClient() as ssh:
+        ssh.load_system_host_keys()
+        if os.path.exists(KNOWN_HOSTS_FILE):
+            ssh.load_host_keys(KNOWN_HOSTS_FILE)
+        ssh.set_missing_host_key_policy(InteractivePolicy())
+        ssh.connect(hostname, port, username, password, sock=sock)
+        with ssh.open_sftp() as sftp:
+            if mode == 'upload':
+                try:
+                    # If the remote path is a directory, use the original filename.
+                    if stat.S_ISDIR(sftp.stat(remote_path).st_mode):
+                        path = os.path.join(remote_path, os.path.basename(local_path))
+                    # A file exists at this destination. We're simply going to clobber it.
+                    else:
                         path = remote_path
-                    finally:
-                        sftp.put(local_path, path, callback=callback)
-                elif mode == 'download':
-                    sftp.get(remote_path, local_path, callback=callback)
-                elif mode == 'size':
-                    return sftp.stat(remote_path).st_size
-    finally:
-        if sock:
-            sock.shutdown(socket.SHUT_RDWR)
-            sock.close()
+                # This path doesn't point at any existing file. We can freely use this filename.
+                except IOError:
+                    path = remote_path
+                finally:
+                    sftp.put(local_path, path, callback=callback)
+            elif mode == 'download':
+                sftp.get(remote_path, local_path, callback=callback)
+            elif mode == 'size':
+                return sftp.stat(remote_path).st_size
 
 def upload_sftp(*args, **kwargs):
     transfer_sftp('upload', *args, **kwargs)