diff options
Diffstat (limited to 'python/vyos/remote.py')
| -rw-r--r-- | python/vyos/remote.py | 85 |
1 files changed, 63 insertions, 22 deletions
diff --git a/python/vyos/remote.py b/python/vyos/remote.py index 66044fa52..cf731c881 100644 --- a/python/vyos/remote.py +++ b/python/vyos/remote.py @@ -25,22 +25,21 @@ import urllib.parse from ftplib import FTP from ftplib import FTP_TLS -from paramiko import SSHClient +from paramiko import SSHClient, SSHException 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 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.utils.io import ask_yes_no +from vyos.utils.io import make_incremental_progressbar +from vyos.utils.io import make_progressbar +from vyos.utils.io import print_error +from vyos.utils.misc import begin +from vyos.utils.process import cmd from vyos.version import get_version - CHUNK_SIZE = 8192 class InteractivePolicy(MissingHostKeyPolicy): @@ -51,7 +50,7 @@ class InteractivePolicy(MissingHostKeyPolicy): 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 sys.stdout.isatty() and 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) @@ -97,7 +96,13 @@ def check_storage(path, size): class FtpC: - def __init__(self, url, progressbar=False, check_space=False, source_host='', source_port=0): + def __init__(self, + url, + progressbar=False, + check_space=False, + source_host='', + source_port=0, + timeout=10): self.secure = url.scheme == 'ftps' self.hostname = url.hostname self.path = url.path @@ -107,12 +112,15 @@ class FtpC: self.source = (source_host, source_port) self.progressbar = progressbar self.check_space = check_space + self.timeout = timeout def _establish(self): if self.secure: - return FTP_TLS(source_address=self.source, context=ssl.create_default_context()) + return FTP_TLS(source_address=self.source, + context=ssl.create_default_context(), + timeout=self.timeout) else: - return FTP(source_address=self.source) + return FTP(source_address=self.source, timeout=self.timeout) def download(self, location: str): # Open the file upfront before establishing connection. @@ -151,7 +159,13 @@ class FtpC: class SshC: known_hosts = os.path.expanduser('~/.ssh/known_hosts') - def __init__(self, url, progressbar=False, check_space=False, source_host='', source_port=0): + def __init__(self, + url, + progressbar=False, + check_space=False, + source_host='', + source_port=0, + timeout=10.0): self.hostname = url.hostname self.path = url.path self.username = url.username or os.getenv('REMOTE_USERNAME') @@ -160,6 +174,7 @@ class SshC: self.source = (source_host, source_port) self.progressbar = progressbar self.check_space = check_space + self.timeout = timeout def _establish(self): ssh = SSHClient() @@ -170,7 +185,7 @@ class SshC: 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) + sock = socket.create_connection((self.hostname, self.port), self.timeout, self.source) ssh.connect(self.hostname, self.port, self.username, self.password, sock=sock) return ssh @@ -199,13 +214,20 @@ class SshC: class HttpC: - def __init__(self, url, progressbar=False, check_space=False, source_host='', source_port=0): + def __init__(self, + url, + progressbar=False, + check_space=False, + source_host='', + source_port=0, + timeout=10.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') + self.timeout = timeout def _establish(self): session = Session() @@ -221,8 +243,11 @@ class HttpC: # 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: + with s.head(self.urlstring, + allow_redirects=True, + timeout=self.timeout) as r: # Abort early if the destination is inaccessible. + print('pre-3') r.raise_for_status() # If the request got redirected, keep the last URL we ended up with. final_urlstring = r.url @@ -236,7 +261,8 @@ class HttpC: size = None if self.check_space: check_storage(location, size) - with s.get(final_urlstring, stream=True) as r, open(location, 'wb') as f: + with s.get(final_urlstring, stream=True, + timeout=self.timeout) as r, open(location, 'wb') as f: if self.progressbar and size: progress = make_incremental_progressbar(CHUNK_SIZE / size) next(progress) @@ -250,7 +276,10 @@ class HttpC: 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) + s.post(self.urlstring, + data=f, + allow_redirects=True, + timeout=self.timeout) class TftpC: @@ -259,10 +288,16 @@ class TftpC: # 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): + def __init__(self, + url, + progressbar=False, + check_space=False, + source_host=None, + source_port=0, + timeout=10): 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.command = f'curl {source_option} {progress_flag} --connect-timeout {timeout}' self.urlstring = urllib.parse.urlunsplit(url) def download(self, location: str): @@ -287,10 +322,16 @@ def urlc(urlstring, *args, **kwargs): raise ValueError(f'Unsupported URL scheme: "{url.scheme}"') def download(local_path, urlstring, *args, **kwargs): - urlc(urlstring, *args, **kwargs).download(local_path) + try: + urlc(urlstring, *args, **kwargs).download(local_path) + except Exception as err: + print_error(f'Unable to download "{urlstring}": {err}') def upload(local_path, urlstring, *args, **kwargs): - urlc(urlstring, *args, **kwargs).upload(local_path) + try: + urlc(urlstring, *args, **kwargs).upload(local_path) + except Exception as err: + print_error(f'Unable to upload "{urlstring}": {err}') def get_remote_config(urlstring, source_host='', source_port=0): """ |