From be167b110dabb1f7f6db7351d828bba1e54e358a Mon Sep 17 00:00:00 2001 From: erkin Date: Sun, 20 Jun 2021 17:05:14 +0300 Subject: T3563: remote: Allow `paramiko.transport.Transport().close()` to automatically close the socket --- python/vyos/remote.py | 51 +++++++++++++++++++++++---------------------------- 1 file changed, 23 insertions(+), 28 deletions(-) (limited to 'python') 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) -- cgit v1.2.3