T3356: remote: Add friendly download procedure for user-facing scripts

1 files changed, 48 insertions, 13 deletions

diff --git a/python/vyos/remote.py b/python/vyos/remote.py
index 52d234d4a..aef4dda7e 100644
--- a/python/vyos/remote.py
+++ b/python/vyos/remote.py
@@ -15,6 +15,7 @@
 
 from ftplib import FTP
 import os
+import shutil
 import socket
 import sys
 import tempfile
@@ -28,14 +29,23 @@ from paramiko import SSHClient, SSHException, MissingHostKeyPolicy
 
 known_hosts_file = os.path.expanduser('~/.ssh/known_hosts')
 
+def print_error(str):
+    """
+    Used for warnings and out-of-band messages to avoid mangling precious
+    stdout output.
+    """
+    sys.stderr.write(str)
+    sys.stderr.write('\n')
+    sys.stderr.flush()
+
 class InteractivePolicy(MissingHostKeyPolicy):
     """
     Policy for interactively querying the user on whether to proceed with
     SSH connections to unknown hosts.
     """
     def missing_host_key(self, client, hostname, key):
-        print(f"Host '{hostname}' not found in known hosts.")
-        print('Fingerprint: ' + key.get_fingerprint().hex())
+        print_error(f"Host '{hostname}' not found in known hosts.")
+        print_error('Fingerprint: ' + key.get_fingerprint().hex())
         if ask_yes_no('Do you wish to continue?'):
             if client._host_keys_filename and ask_yes_no('Do you wish to permanently add this host/key pair to known hosts?'):
                 client._host_keys.add(hostname, key.get_name(), key)
@@ -186,7 +196,7 @@ def download(local_path, urlstring, authentication=None, source=None):
 
     if url.scheme == 'http' or url.scheme == 'https':
         if source:
-            print('Warning: Custom source address not supported for HTTP connections.', file=sys.stderr)
+            print_error('Warning: Custom source address not supported for HTTP connections.')
         download_http(urlstring, local_path, username, password)
     elif url.scheme == 'ftp':
         username = username if username else 'anonymous'
@@ -247,16 +257,6 @@ def get_remote_file_size(urlstring, authentication=None, source=None):
     else:
         raise ValueError(f'Unsupported URL scheme: {url.scheme}')
 
-def get_remote_file_size_maybe(urlstring, authentication=None, source=None):
-    """
-    Passes arguments to `get_remote_file_size()` but returns 0 if it fails.
-    Intended to be used in shell scripts only.
-    """
-    try:
-        return get_remote_file_size(urlstring, authentication, source)
-    except ValueError:
-        return 0
-
 def get_remote_config(urlstring, authentication=None, source=None):
     """
     Download remote (config) file from `urlstring` and return the contents as a string.
@@ -282,3 +282,38 @@ def get_remote_config(urlstring, authentication=None, source=None):
             return file.read()
     finally:
         os.remove(temp)
+
+def friendly_download(local_path, urlstring, authentication=None, source=None):
+    """
+    Intended to be called from interactive, user-facing scripts.
+    """
+    destination_directory = os.path.dirname(local_path)
+    free_space = shutil.disk_usage(destination_directory).free
+    try:
+        try:
+            file_size = get_remote_file_size(urlstring, authentication, source)
+            if file_size < 1024 * 1024:
+                print_error(f'The file is {file_size / 1024.0:.3f} KiB.')
+            else:
+                print_error(f'The file is {file_size / (1024.0 * 1024.0):.3f} MiB.')
+            if file_size > free_space:
+                raise OSError(f'Not enough disk space available in "{destination_directory}".')
+        except ValueError:
+            print_error('Could not determine the file size in advance.')
+        else:
+            # TODO: Progress bar
+            print_error('Downloading...')
+            download(local_path, urlstring, authentication, source)
+    except KeyboardInterrupt:
+        print_error('Download aborted by user.')
+        sys.exit(1)
+    except:
+        import traceback
+        # There are a myriad different reasons a download could fail.
+        # SSH errors, FTP errors, I/O errors, HTTP errors (403, 404...)
+        # We omit the scary stack trace but print the error nevertheless.
+        print_error(f'Failed to download {urlstring}.')
+        traceback.print_exception(*sys.exc_info()[:2], None)
+        sys.exit(1)
+    else:
+        print_error('Download complete.')