Merge pull request #1108 from erkin/equuleus

remote: T3356: Backport remote module

2 files changed, 354 insertions, 109 deletions

diff --git a/python/vyos/remote.py b/python/vyos/remote.py
index db20b173d..aa62ac60d 100644
--- a/python/vyos/remote.py
+++ b/python/vyos/remote.py
@@ -14,133 +14,315 @@
 # License along with this library.  If not, see <http://www.gnu.org/licenses/>.
 
 import os
+import shutil
 import socket
+import ssl
+import stat
 import sys
 import tempfile
-from ftplib import FTP
 import urllib.parse
-import urllib.request
 
-from vyos.util import cmd
+from ftplib import FTP
+from ftplib import FTP_TLS
+
 from paramiko import SSHClient
+from paramiko import MissingHostKeyPolicy
 
+from requests import Session
+from requests.adapters import HTTPAdapter
+from requests.packages.urllib3 import PoolManager
 
-def upload_ftp(local_path, hostname, remote_path,\
-               username='anonymous', password='', port=21, source=None):
-    with open(local_path, 'rb') as file:
-        with FTP(source_address=source) as conn:
-            conn.connect(hostname, port)
-            conn.login(username, password)
-            conn.storbinary(f'STOR {remote_path}', file)
-
-def download_ftp(local_path, hostname, remote_path,\
-                 username='anonymous', password='', port=21, source=None):
-    with open(local_path, 'wb') as file:
-        with FTP(source_address=source) as conn:
-            conn.connect(hostname, port)
-            conn.login(username, password)
-            conn.retrbinary(f'RETR {remote_path}', file.write)
-
-def upload_sftp(local_path, hostname, remote_path,\
-                username=None, password=None, port=22, source=None):
-    sock = None
-    if source:
-        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        sock.bind((source, 0))
-        sock.connect((hostname, port))
-    with SSHClient() as ssh:
-        ssh.load_system_host_keys()
-        ssh.connect(hostname, port, username, password, sock=sock)
-        with ssh.open_sftp() as sftp:
-            sftp.put(local_path, remote_path)
-    if sock:
-        sock.shutdown()
-        sock.close()
-
-def download_sftp(local_path, hostname, remote_path,\
-                  username=None, password=None, port=22, source=None):
-    sock = None
-    if source:
-        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        sock.bind((source, 0))
-        sock.connect((hostname, port))
-    with SSHClient() as ssh:
-        ssh.load_system_host_keys()
-        ssh.connect(hostname, port, username, password, sock=sock)
-        with ssh.open_sftp() as sftp:
-            sftp.get(remote_path, local_path)
-    if sock:
-        sock.shutdown()
-        sock.close()
-
-def upload_tftp(local_path, hostname, remote_path, port=69, source=None):
-    source_option = f'--interface {source}' if source else ''
-    with open(local_path, 'rb') as file:
-        cmd(f'curl {source_option} -s -T - tftp://{hostname}:{port}/{remote_path}',\
-            stderr=None, input=file.read()).encode()
-
-def download_tftp(local_path, hostname, remote_path, port=69, source=None):
-    source_option = f'--interface {source}' if source else ''
-    with open(local_path, 'wb') as file:
-        file.write(cmd(f'curl {source_option} -s tftp://{hostname}:{port}/{remote_path}',\
-                       stderr=None).encode())
-
-def download_http(urlstring, local_path):
-    with open(local_path, 'wb') as file:
-        with urllib.request.urlopen(urlstring) as response:
-            file.write(response.read())
-
-def download(local_path, urlstring, source=None):
+from vyos.util import ask_yes_no
+from vyos.util import begin
+from vyos.util import cmd
+from vyos.util import make_incremental_progressbar
+from vyos.util import make_progressbar
+from vyos.util import print_error
+from vyos.version import get_version
+
+
+CHUNK_SIZE = 8192
+
+class InteractivePolicy(MissingHostKeyPolicy):
     """
-    Dispatch the appropriate download function for the given URL and save to local path.
+    Paramiko policy for interactively querying the user on whether to proceed
+     with SSH connections to unknown hosts.
     """
-    url = urllib.parse.urlparse(urlstring)
-    if url.scheme == 'http' or url.scheme == 'https':
-        if source:
-            print("Warning: Custom source address not supported for HTTP connections.", file=sys.stderr)
-        download_http(urlstring, local_path)
-    elif url.scheme == 'ftp':
-        username = url.username if url.username else 'anonymous'
-        download_ftp(local_path, url.hostname, url.path, username, url.password, source=source)
-    elif url.scheme == 'sftp' or url.scheme == 'scp':
-        download_sftp(local_path, url.hostname, url.path, url.username, url.password, source=source)
-    elif url.scheme == 'tftp':
-        download_tftp(local_path, url.hostname, url.path, source=source)
-    else:
-        raise ValueError(f'Unsupported URL scheme: {url.scheme}')
+    def missing_host_key(self, client, hostname, key):
+        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)
+                client.save_host_keys(client._host_keys_filename)
+        else:
+            raise SSHException(f"Cannot connect to unknown host '{hostname}'.")
 
-def upload(local_path, urlstring, source=None):
+class SourceAdapter(HTTPAdapter):
     """
-    Dispatch the appropriate upload function for the given URL and upload from local path.
+    urllib3 transport adapter for setting source addresses per session.
     """
-    url = urllib.parse.urlparse(urlstring)
-    if url.scheme == 'ftp':
-        username = url.username if url.username else 'anonymous'
-        upload_ftp(local_path, url.hostname, url.path, username, url.password, source=source)
-    elif url.scheme == 'sftp' or url.scheme == 'scp':
-        upload_sftp(local_path, url.hostname, url.path, url.username, url.password, source=source)
-    elif url.scheme == 'tftp':
-        upload_tftp(local_path, url.hostname, url.path, source=source)
+    def __init__(self, source_pair, *args, **kwargs):
+        # A source pair is a tuple of a source host string and source port respectively.
+        # Supply '' and 0 respectively for default values.
+        self._source_pair = source_pair
+        super(SourceAdapter, self).__init__(*args, **kwargs)
+
+    def init_poolmanager(self, connections, maxsize, block=False):
+        self.poolmanager = PoolManager(
+            num_pools=connections, maxsize=maxsize,
+            block=block, source_address=self._source_pair)
+
+
+def check_storage(path, size):
+    """
+    Check whether `path` has enough storage space for a transfer of `size` bytes.
+    """
+    path = os.path.abspath(os.path.expanduser(path))
+    directory = path if os.path.isdir(path) else (os.path.dirname(os.path.expanduser(path)) or os.getcwd())
+    # `size` can be None or 0 to indicate unknown size.
+    if not size:
+        print_error('Warning: Cannot determine size of remote file.')
+        print_error('Bravely continuing regardless.')
+        return
+
+    if size < 1024 * 1024:
+        print_error(f'The file is {size / 1024.0:.3f} KiB.')
     else:
-        raise ValueError(f'Unsupported URL scheme: {url.scheme}')
+        print_error(f'The file is {size / (1024.0 * 1024.0):.3f} MiB.')
+
+    # Will throw `FileNotFoundError' if `directory' is absent.
+    if size > shutil.disk_usage(directory).free:
+        raise OSError(f'Not enough disk space available in "{directory}".')
+
 
-def get_remote_config(urlstring):
+class FtpC:
+    def __init__(self, url, progressbar=False, check_space=False, source_host='', source_port=0):
+        self.secure = url.scheme == 'ftps'
+        self.hostname = url.hostname
+        self.path = url.path
+        self.username = url.username or os.getenv('REMOTE_USERNAME', 'anonymous')
+        self.password = url.password or os.getenv('REMOTE_PASSWORD', '')
+        self.port = url.port or 21
+        self.source = (source_host, source_port)
+        self.progressbar = progressbar
+        self.check_space = check_space
+
+    def _establish(self):
+        if self.secure:
+            return FTP_TLS(source_address=self.source, context=ssl.create_default_context())
+        else:
+            return FTP(source_address=self.source)
+
+    def download(self, location: str):
+        # Open the file upfront before establishing connection.
+        with open(location, 'wb') as f, self._establish() as conn:
+            conn.connect(self.hostname, self.port)
+            conn.login(self.username, self.password)
+            # Set secure connection over TLS.
+            if self.secure:
+                conn.prot_p()
+            # Almost all FTP servers support the `SIZE' command.
+            if self.check_space:
+                check_storage(path, conn.size(self.path))
+            # No progressbar if we can't determine the size or if the file is too small.
+            if self.progressbar and size and size > CHUNK_SIZE:
+                progress = make_incremental_progressbar(CHUNK_SIZE / size)
+                next(progress)
+                callback = lambda block: begin(f.write(block), next(progress))
+            else:
+                callback = f.write
+            conn.retrbinary('RETR ' + self.path, callback, CHUNK_SIZE)
+
+    def upload(self, location: str):
+        size = os.path.getsize(location)
+        with open(location, 'rb') as f, self._establish() as conn:
+            conn.connect(self.hostname, self.port)
+            conn.login(self.username, self.password)
+            if self.secure:
+                conn.prot_p()
+            if self.progressbar and size and size > CHUNK_SIZE:
+                progress = make_incremental_progressbar(CHUNK_SIZE / size)
+                next(progress)
+                callback = lambda block: next(progress)
+            else:
+                callback = None
+            conn.storbinary('STOR ' + self.path, f, CHUNK_SIZE, callback)
+
+class SshC:
+    known_hosts = os.path.expanduser('~/.ssh/known_hosts')
+    def __init__(self, url, progressbar=False, check_space=False, source_host='', source_port=0):
+        self.hostname = url.hostname
+        self.path = url.path
+        self.username = url.username or os.getenv('REMOTE_USERNAME')
+        self.password = url.password or os.getenv('REMOTE_PASSWORD')
+        self.port = url.port or 22
+        self.source = (source_host, source_port)
+        self.progressbar = progressbar
+        self.check_space = check_space
+
+    def _establish(self):
+        ssh = SSHClient()
+        ssh.load_system_host_keys()
+        # Try to load from a user-local known hosts file if one exists.
+        if os.path.exists(self.known_hosts):
+            ssh.load_host_keys(self.known_hosts)
+        ssh.set_missing_host_key_policy(InteractivePolicy())
+        # `socket.create_connection()` automatically picks a NIC and an IPv4/IPv6 address family
+        #  for us on dual-stack systems.
+        sock = socket.create_connection((self.hostname, self.port), socket.getdefaulttimeout(), self.source)
+        ssh.connect(self.hostname, self.port, self.username, self.password, sock=sock)
+        return ssh
+
+    def download(self, location: str):
+        callback = make_progressbar() if self.progressbar else None
+        with self._establish() as ssh, ssh.open_sftp() as sftp:
+            if self.check_space:
+                check_storage(location, sftp.stat(self.path).st_size)
+            sftp.get(self.path, location, callback=callback)
+
+    def upload(self, location: str):
+        callback = make_progressbar() if self.progressbar else None
+        with self._establish() as ssh, ssh.open_sftp() as sftp:
+            try:
+                # If the remote path is a directory, use the original filename.
+                if stat.S_ISDIR(sftp.stat(self.path).st_mode):
+                    path = os.path.join(self.path, os.path.basename(location))
+                # A file exists at this destination. We're simply going to clobber it.
+                else:
+                    path = self.path
+            # This path doesn't point at any existing file. We can freely use this filename.
+            except IOError:
+                path = self.path
+            finally:
+                sftp.put(location, path, callback=callback)
+
+
+class HttpC:
+    def __init__(self, url, progressbar=False, check_space=False, source_host='', source_port=0):
+        self.urlstring = urllib.parse.urlunsplit(url)
+        self.progressbar = progressbar
+        self.check_space = check_space
+        self.source_pair = (source_host, source_port)
+        self.username = url.username or os.getenv('REMOTE_USERNAME')
+        self.password = url.password or os.getenv('REMOTE_PASSWORD')
+
+    def _establish(self):
+        session = Session()
+        session.mount(self.urlstring, SourceAdapter(self.source_pair))
+        session.headers.update({'User-Agent': 'VyOS/' + get_version()})
+        if self.username:
+            session.auth = self.username, self.password
+        return session
+
+    def download(self, location: str):
+        with self._establish() as s:
+            # We ask for uncompressed downloads so that we don't have to deal with decoding.
+            # Not only would it potentially mess up with the progress bar but
+            # `shutil.copyfileobj(request.raw, file)` does not handle automatic decoding.
+            s.headers.update({'Accept-Encoding': 'identity'})
+            with s.head(self.urlstring, allow_redirects=True) as r:
+                # Abort early if the destination is inaccessible.
+                r.raise_for_status()
+                # If the request got redirected, keep the last URL we ended up with.
+                final_urlstring = r.url
+                if r.history:
+                    print_error('Redirecting to ' + final_urlstring)
+                # Check for the prospective file size.
+                try:
+                    size = int(r.headers['Content-Length'])
+                # In case the server does not supply the header.
+                except KeyError:
+                    size = None
+            if self.check_space:
+                check_storage(location, size)
+            with s.get(final_urlstring, stream=True) as r, open(location, 'wb') as f:
+                if self.progressbar and size:
+                    progress = make_incremental_progressbar(CHUNK_SIZE / size)
+                    next(progress)
+                    for chunk in iter(lambda: begin(next(progress), r.raw.read(CHUNK_SIZE)), b''):
+                        f.write(chunk)
+                else:
+                    # We'll try to stream the download directly with `copyfileobj()` so that large
+                    #  files (like entire VyOS images) don't occupy much memory.
+                    shutil.copyfileobj(r.raw, f)
+
+    def upload(self, location: str):
+        # Does not yet support progressbars.
+        with self._establish() as s, open(location, 'rb') as f:
+            s.post(self.urlstring, data=f, allow_redirects=True)
+
+
+class TftpC:
+    # We simply allow `curl` to take over because
+    # 1. TFTP is rather simple.
+    # 2. Since there's no concept authentication, we don't need to deal with keys/passwords.
+    # 3. It would be a waste to import, audit and maintain a third-party library for TFTP.
+    # 4. I'd rather not implement the entire protocol here, no matter how simple it is.
+    def __init__(self, url, progressbar=False, check_space=False, source_host=None, source_port=0):
+        source_option = f'--interface {source_host} --local-port {source_port}' if source_host else ''
+        progress_flag = '--progress-bar' if progressbar else '-s'
+        self.command = f'curl {source_option} {progress_flag}'
+        self.urlstring = urllib.parse.urlunsplit(url)
+
+    def download(self, location: str):
+        with open(location, 'wb') as f:
+            f.write(cmd(f'{self.command} "{self.urlstring}"').encode())
+
+    def upload(self, location: str):
+        with open(location, 'rb') as f:
+            cmd(f'{self.command} -T - "{self.urlstring}"', input=f.read())
+
+
+def urlc(urlstring, *args, **kwargs):
+    """
+    Dynamically dispatch the appropriate protocol class.
+    """
+    url_classes = {'http': HttpC, 'https': HttpC, 'ftp': FtpC, 'ftps': FtpC, \
+                   'sftp': SshC, 'ssh': SshC, 'scp': SshC, 'tftp': TftpC}
+    url = urllib.parse.urlsplit(urlstring)
+    try:
+        return url_classes[url.scheme](url, *args, **kwargs)
+    except KeyError:
+        raise ValueError(f'Unsupported URL scheme: "{url.scheme}"')
+
+def download(local_path, urlstring, *args, **kwargs):
+    urlc(urlstring, *args, **kwargs).download(local_path)
+
+def upload(local_path, urlstring, *args, **kwargs):
+    urlc(urlstring, *args, **kwargs).upload(local_path)
+
+def get_remote_config(urlstring, source_host='', source_port=0):
     """
-    Download remote (config) file and return the contents.
-        Args:
-            remote file URI:
-                scp://<user>[:<passwd>]@<host>/<file>
-                sftp://<user>[:<passwd>]@<host>/<file>
-                http://<host>/<file>
-                https://<host>/<file>
-                ftp://[<user>[:<passwd>]@]<host>/<file>
-                tftp://<host>/<file>
+    Quietly download a file and return it as a string.
     """
-    url = urllib.parse.urlparse(urlstring)
     temp = tempfile.NamedTemporaryFile(delete=False).name
     try:
-        download(temp, urlstring)
-        with open(temp, 'r') as file:
-            return file.read()
+        download(temp, urlstring, False, False, source_host, source_port)
+        with open(temp, 'r') as f:
+            return f.read()
     finally:
         os.remove(temp)
+
+def friendly_download(local_path, urlstring, source_host='', source_port=0):
+    """
+    Download with a progress bar, reassuring messages and free space checks.
+    """
+    try:
+        print_error('Downloading...')
+        download(local_path, urlstring, True, True, source_host, source_port)
+    except KeyboardInterrupt:
+        print_error('\nDownload 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.')
+        sys.exit(0)
diff --git a/python/vyos/util.py b/python/vyos/util.py
index dcb332e5c..f2f302559 100644
--- a/python/vyos/util.py
+++ b/python/vyos/util.py
@@ -704,6 +704,69 @@ def get_interface_config(interface):
     tmp = loads(cmd(f'ip -d -j link show {interface}'))[0]
     return tmp
 
+def print_error(str='', end='\n'):
+    """
+    Print `str` to stderr, terminated with `end`.
+    Used for warnings and out-of-band messages to avoid mangling precious
+     stdout output.
+    """
+    sys.stderr.write(str)
+    sys.stderr.write(end)
+    sys.stderr.flush()
+
+def make_progressbar():
+    """
+    Make a procedure that takes two arguments `done` and `total` and prints a
+     progressbar based on the ratio thereof, whose length is determined by the
+     width of the terminal.
+    """
+    import shutil, math
+    col, _ = shutil.get_terminal_size()
+    col = max(col - 15, 20)
+    def print_progressbar(done, total):
+        if done <= total:
+            increment = total / col
+            length = math.ceil(done / increment)
+            percentage = str(math.ceil(100 * done / total)).rjust(3)
+            print_error(f'[{length * "#"}{(col - length) * "_"}] {percentage}%', '\r')
+            # Print a newline so that the subsequent prints don't overwrite the full bar.
+        if done == total:
+            print_error()
+    return print_progressbar
+
+def make_incremental_progressbar(increment: float):
+    """
+    Make a generator that displays a progressbar that grows monotonically with
+     every iteration.
+    First call displays it at 0% and every subsequent iteration displays it
+     at `increment` increments where 0.0 < `increment` < 1.0.
+    Intended for FTP and HTTP transfers with stateless callbacks.
+    """
+    print_progressbar = make_progressbar()
+    total = 0.0
+    while total < 1.0:
+        print_progressbar(total, 1.0)
+        yield
+        total += increment
+    print_progressbar(1, 1)
+    # Ignore further calls.
+    while True:
+        yield
+
+def begin(*args):
+    """
+    Evaluate arguments in order and return the result of the *last* argument.
+    For combining multiple expressions in one statement. Useful for lambdas.
+    """
+    return args[-1]
+
+def begin0(*args):
+    """
+    Evaluate arguments in order and return the result of the *first* argument.
+    For combining multiple expressions in one statement. Useful for lambdas.
+    """
+    return args[0]
+
 def is_systemd_service_active(service):
     """ Test is a specified systemd service is activated.
     Returns True if service is active, false otherwise.