summaryrefslogtreecommitdiff
path: root/python/vyos
diff options
context:
space:
mode:
Diffstat (limited to 'python/vyos')
-rw-r--r--python/vyos/remote.py572
-rw-r--r--python/vyos/util.py14
2 files changed, 268 insertions, 318 deletions
diff --git a/python/vyos/remote.py b/python/vyos/remote.py
index e972050b7..2419f8873 100644
--- a/python/vyos/remote.py
+++ b/python/vyos/remote.py
@@ -13,38 +13,40 @@
# You should have received a copy of the GNU Lesser General Public
# License along with this library. If not, see <http://www.gnu.org/licenses/>.
-from ftplib import FTP
import os
import shutil
import socket
+import ssl
import stat
import sys
import tempfile
import urllib.parse
-import urllib.request as urlreq
-from vyos.template import get_ip
-from vyos.template import ip_from_cidr
-from vyos.template import is_interface
-from vyos.template import is_ipv6
-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
+
from vyos.util import ask_yes_no
-from vyos.util import print_error
-from vyos.util import make_progressbar
+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
-from paramiko import SSHClient
-from paramiko import SSHException
-from paramiko import MissingHostKeyPolicy
-# This is a hardcoded path and no environment variable can change it.
-KNOWN_HOSTS_FILE = os.path.expanduser('~/.ssh/known_hosts')
+
CHUNK_SIZE = 8192
class InteractivePolicy(MissingHostKeyPolicy):
"""
- Policy for interactively querying the user on whether to proceed with
- SSH connections to unknown hosts.
+ Paramiko 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_error(f"Host '{hostname}' not found in known hosts.")
@@ -57,337 +59,270 @@ class InteractivePolicy(MissingHostKeyPolicy):
else:
raise SSHException(f"Cannot connect to unknown host '{hostname}'.")
-
-## Helper routines
-def get_authentication_variables(default_username=None, default_password=None):
+class SourceAdapter(HTTPAdapter):
"""
- Return the environment variables `$REMOTE_USERNAME` and `$REMOTE_PASSWORD` and
- return the defaults provided if environment variables are empty or nonexistent.
+ urllib3 transport adapter for setting source addresses per session.
"""
- username, password = os.getenv('REMOTE_USERNAME'), os.getenv('REMOTE_PASSWORD')
- # Fall back to defaults if the username variable doesn't exist or is an empty string.
- # Note that this is different from `os.getenv('REMOTE_USERNAME', default=default_username)`,
- # as we want the username and the password to have the same behaviour.
- if not username:
- return default_username, default_password
- else:
- return username, password
-
-def get_source_address(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)
+
+class WrappedFile:
+ def __init__(self, obj, size=None, chunk_size=CHUNK_SIZE):
+ self._obj = obj
+ self._progress = size and make_incremental_progressbar(chunk_size / size)
+ def read(self, size=-1):
+ if self._progress:
+ next(self._progress)
+ self._obj.read(size)
+ def write(self, size=-1):
+ if self._progress:
+ next(self._progress)
+ self._obj.write(size)
+ def __getattr__(self, attr):
+ return getattr(self._obj, attr)
+
+def check_storage(path, size):
"""
- Take a string vaguely indicating an origin source (interface, hostname or IP address),
- return a tuple in the format `(source_pair, address_family)` where
- `source_pair` is `(source_address, source_port)`.
+ Check whether `path` has enough storage space for a transfer of `size` bytes.
"""
- # TODO: Properly distinguish between IPv4 and IPv6.
- port = 0
- if is_interface(source):
- source = ip_from_cidr(get_ip(source)[0])
- if is_ipv6(source):
- return (source, port), socket.AF_INET6
+ 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:
- return (socket.gethostbyname(source), port), socket.AF_INET
-
-def get_port_from_url(url):
- """
- Return the port number from the given `url` named tuple, fall back to
- the default if there isn't one.
- """
- defaults = {"http": 80, "https": 443, "ftp": 21, "tftp": 69,\
- "ssh": 22, "scp": 22, "sftp": 22}
- if url.port:
- return url.port
- else:
- return defaults[url.scheme]
-
-
-## FTP routines
-def upload_ftp(local_path, hostname, remote_path,\
- username='anonymous', password='', port=21,\
- source_pair=None, progressbar=False):
- size = os.path.getsize(local_path)
- with FTP(source_address=source_pair) as conn:
- conn.connect(hostname, port)
- conn.login(username, password)
- with open(local_path, 'rb') as file:
- if progressbar and size:
+ 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}".')
+
+
+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: next(progress)
+ callback = lambda block: begin(f.write(block), next(progress))
else:
- callback = None
- conn.storbinary(f'STOR {remote_path}', file, CHUNK_SIZE, callback)
-
-def download_ftp(local_path, hostname, remote_path,\
- username='anonymous', password='', port=21,\
- source_pair=None, progressbar=False):
- with FTP(source_address=source_pair) as conn:
- conn.connect(hostname, port)
- conn.login(username, password)
- size = conn.size(remote_path)
- with open(local_path, 'wb') as file:
- # No progressbar if we can't determine the size.
- if progressbar and size:
+ 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: (file.write(block), next(progress))
+ callback = lambda block: next(progress)
else:
- callback = file.write
- conn.retrbinary(f'RETR {remote_path}', callback, CHUNK_SIZE)
-
-def get_ftp_file_size(hostname, remote_path,\
- username='anonymous', password='', port=21,\
- source_pair=None):
- with FTP(source_address=source) as conn:
- conn.connect(hostname, port)
- conn.login(username, password)
- size = conn.size(remote_path)
- if size:
- return size
- else:
- # SIZE is an extension to the FTP specification, although it's extremely common.
- raise ValueError('Failed to receive file size from FTP server. \
- Perhaps the server does not implement the SIZE command?')
-
-
-## SFTP/SCP routines
-def transfer_sftp(mode, local_path, hostname, remote_path,\
- username=None, password=None, port=22,\
- source_tuple=None, progressbar=False):
- sock = None
- if source_tuple:
- (source_address, source_port), address_family = source_tuple
- sock = socket.socket(address_family, socket.SOCK_STREAM)
- sock.bind((source_address, source_port))
- sock.connect((hostname, port))
- callback = make_progressbar() if progressbar else None
- with SSHClient() as ssh:
+ 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()
- if os.path.exists(KNOWN_HOSTS_FILE):
- ssh.load_host_keys(KNOWN_HOSTS_FILE)
+ # 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())
- ssh.connect(hostname, port, username, password, sock=sock)
- with ssh.open_sftp() as sftp:
- if mode == 'upload':
+ # `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) as r:
+ # Abort early if the destination is inaccessible.
+ r.raise_for_status()
+ # Check for the prospective file size.
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:
- 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)
-
-def download_sftp(*args, **kwargs):
- transfer_sftp('download', *args, **kwargs)
-
-def get_sftp_file_size(*args, **kwargs):
- return transfer_sftp('size', None, *args, **kwargs)
-
-
-## TFTP routines
-def upload_tftp(local_path, hostname, remote_path, port=69, source=None, progressbar=False):
- source_option = f'--interface {source}' if source else ''
- progress_flag = '--progress-bar' if progressbar else '-s'
- with open(local_path, 'rb') as file:
- cmd(f'curl {source_option} {progress_flag} -T - tftp://{hostname}:{port}/{remote_path}',\
- stderr=None, input=file.read()).encode()
-
-def download_tftp(local_path, hostname, remote_path, port=69, source=None, progressbar=False):
- source_option = f'--interface {source}' if source else ''
- # Not really applicable but we pass it for the sake of uniformity.
- progress_flag = '--progress-bar' if progressbar else '-s'
- with open(local_path, 'wb') as file:
- file.write(cmd(f'curl {source_option} {progress_flag} tftp://{hostname}:{port}/{remote_path}',\
- stderr=None).encode())
-
-# get_tftp_file_size() is unimplemented because there is no way to obtain a file's size through TFTP,
-# as TFTP does not specify a SIZE command.
-
-
-## HTTP(S) routines
-def install_request_opener(urlstring, username, password):
- """
- Take `username` and `password` strings and install the appropriate
- password manager to `urllib.request.urlopen()` for the given `urlstring`.
- """
- manager = urlreq.HTTPPasswordMgrWithDefaultRealm()
- manager.add_password(None, urlstring, username, password)
- urlreq.install_opener(urlreq.build_opener(urlreq.HTTPBasicAuthHandler(manager)))
-
-# upload_http() is unimplemented.
-
-def download_http(local_path, urlstring, username=None, password=None, progressbar=False):
- """
- Download the file from from `urlstring` to `local_path`.
- Optionally takes `username` and `password` for authentication.
- """
- request = urlreq.Request(urlstring, headers={'User-Agent': 'VyOS/' + get_version()})
- if username:
- install_request_opener(urlstring, username, password)
- with open(local_path, 'wb') as file, urlreq.urlopen(request) as response:
- size = response.getheader('Content-Length')
- if progressbar and size:
- progress = make_incremental_progressbar(CHUNK_SIZE / int(size))
- next(progress)
- for chunk in iter(lambda: response.read(CHUNK_SIZE), b''):
- file.write(chunk)
- next(progress)
- next(progress)
- # If we can't determine the size or if a progress bar wasn't requested,
- # we can let `shutil` take care of the copying.
- else:
- shutil.copyfileobj(response, file)
-
-def get_http_file_size(urlstring, username=None, password=None):
- """
- Return the size of the file from `urlstring` in terms of number of bytes.
- Optionally takes `username` and `password` for authentication.
- """
- request = urlreq.Request(urlstring, headers={'User-Agent': 'VyOS/' + get_version()})
- if username:
- install_request_opener(urlstring, username, password)
- with urlreq.urlopen(request) as response:
- size = response.getheader('Content-Length')
- if size:
- return int(size)
- # The server didn't send 'Content-Length' in the response headers.
- else:
- raise ValueError('Failed to receive file size from HTTP server.')
-
-
-## Dynamic dispatchers
-def download(local_path, urlstring, source=None, progressbar=False):
+ 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(self.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):
+ size = os.path.getsize(location) if self.progressbar else None
+ # Keep in mind that `data` can be a file-like or iterable object.
+ with self._establish() as s, file(location, 'rb') as f:
+ s.post(self.urlstring, data=WrappedFile(f, size))
+
+
+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):
"""
- Dispatch the appropriate download function for the given `urlstring` and save to `local_path`.
- Optionally takes a `source` address or interface (not valid for HTTP(S)).
- Supports HTTP, HTTPS, FTP, SFTP, SCP (through SFTP) and TFTP.
- Reads `$REMOTE_USERNAME` and `$REMOTE_PASSWORD` environment variables.
+ Dynamically dispatch the appropriate protocol class.
"""
- url = urllib.parse.urlparse(urlstring)
- username, password = get_authentication_variables(url.username, url.password)
- port = get_port_from_url(url)
-
- if url.scheme == 'http' or url.scheme == 'https':
- if source:
- print_error('Warning: Custom source address not supported for HTTP connections.')
- download_http(local_path, urlstring, username, password, progressbar)
- elif url.scheme == 'ftp':
- source = get_source_address(source)[0] if source else None
- username = username if username else 'anonymous'
- download_ftp(local_path, url.hostname, url.path, username, password, port, source, progressbar)
- elif url.scheme == 'sftp' or url.scheme == 'scp':
- source = get_source_address(source) if source else None
- download_sftp(local_path, url.hostname, url.path, username, password, port, source, progressbar)
- elif url.scheme == 'tftp':
- download_tftp(local_path, url.hostname, url.path, port, source, progressbar)
- else:
- raise ValueError(f'Unsupported URL scheme: {url.scheme}')
+ 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 upload(local_path, urlstring, source=None, progressbar=False):
- """
- Dispatch the appropriate upload function for the given URL and upload from local path.
- Optionally takes a `source` address.
- Supports FTP, SFTP, SCP (through SFTP) and TFTP.
- Reads `$REMOTE_USERNAME` and `$REMOTE_PASSWORD` environment variables.
- """
- url = urllib.parse.urlparse(urlstring)
- username, password = get_authentication_variables(url.username, url.password)
- port = get_port_from_url(url)
-
- if url.scheme == 'ftp':
- username = username if username else 'anonymous'
- source = get_source_address(source)[0] if source else None
- upload_ftp(local_path, url.hostname, url.path, username, password, port, source, progressbar)
- elif url.scheme == 'sftp' or url.scheme == 'scp':
- source = get_source_address(source) if source else None
- upload_sftp(local_path, url.hostname, url.path, username, password, port, source, progressbar)
- elif url.scheme == 'tftp':
- upload_tftp(local_path, url.hostname, url.path, port, source, progressbar)
- else:
- raise ValueError(f'Unsupported URL scheme: {url.scheme}')
+def download(local_path, urlstring, *args, **kwargs):
+ urlc(urlstring, *args, **kwargs).download(local_path)
-def get_remote_file_size(urlstring, source=None):
- """
- Dispatch the appropriate function to return the size of the remote file from `urlstring`
- in terms of number of bytes.
- Optionally takes a `source` address (not valid for HTTP(S)).
- Supports HTTP, HTTPS, FTP and SFTP (through SFTP).
- Reads `$REMOTE_USERNAME` and `$REMOTE_PASSWORD` environment variables.
- """
- url = urllib.parse.urlparse(urlstring)
- username, password = get_authentication_variables(url.username, url.password)
- port = get_port_from_url(url)
-
- if url.scheme == 'http' or url.scheme == 'https':
- if source:
- print_error('Warning: Custom source address not supported for HTTP connections.')
- return get_http_file_size(urlstring, username, password)
- elif url.scheme == 'ftp':
- source = get_source_address(source)[0] if source else None
- username = username if username else 'anonymous'
- return get_ftp_file_size(url.hostname, url.path, username, password, port, source)
- elif url.scheme == 'sftp' or url.scheme == 'scp':
- source = get_source_address(source) if source else None
- return get_sftp_file_size(url.hostname, url.path, username, password, port, source)
- else:
- raise ValueError(f'Unsupported URL scheme: {url.scheme}')
+def upload(local_path, urlstring, *args, **kwargs):
+ urlc(urlstring, *args, **kwargs).upload(local_path)
-def get_remote_config(urlstring, source=None):
+def get_remote_config(urlstring, source_host='', source_port=0):
"""
- Download remote (config) file from `urlstring` and return the contents as a string.
- Args:
- remote file URI:
- tftp://<host>[:<port>]/<file>
- http[s]://<host>[:<port>]/<file>
- [scp|sftp|ftp]://[<user>[:<passwd>]@]<host>[:port]/<file>
- source address (optional):
- <interface>
- <IP address>
+ Quietly download a file and return it as a string.
"""
temp = tempfile.NamedTemporaryFile(delete=False).name
try:
- download(temp, urlstring, source)
- 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=None):
+def friendly_download(local_path, urlstring, source_host='', source_port=0):
"""
- Download from `urlstring` to `local_path` in an informative way.
- Checks the storage space before attempting download.
- Intended to be called from interactive, user-facing scripts.
+ Download with a progress bar, reassuring messages and free space checks.
"""
- destination_directory = os.path.dirname(local_path)
try:
- free_space = shutil.disk_usage(destination_directory).free
- try:
- file_size = get_remote_file_size(urlstring, 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:
- # Can't do a storage check in this case, so we bravely continue.
- file_size = 0
- print_error('Could not determine the file size in advance.')
- else:
- print_error('Downloading...')
- download(local_path, urlstring, source, progressbar=file_size > 1024 * 1024)
+ print_error('Downloading...')
+ download(local_path, urlstring, True, True, source_host, source_port)
except KeyboardInterrupt:
print_error('Download aborted by user.')
sys.exit(1)
@@ -401,3 +336,4 @@ def friendly_download(local_path, urlstring, source=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 9aa1f98d2..d8e83ab8d 100644
--- a/python/vyos/util.py
+++ b/python/vyos/util.py
@@ -856,6 +856,20 @@ def make_incremental_progressbar(increment: float):
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.