diff --git a/aiida/transports/plugins/ssh.py b/aiida/transports/plugins/ssh.py index c6a37f76a3..43a101f808 100644 --- a/aiida/transports/plugins/ssh.py +++ b/aiida/transports/plugins/ssh.py @@ -12,9 +12,11 @@ import glob import io import os +import re from stat import S_ISDIR, S_ISREG import click +import paramiko from aiida.cmdline.params import options from aiida.cmdline.params.types.path import AbsolutePathOrEmptyParamType @@ -33,7 +35,6 @@ def parse_sshconfig(computername): :param computername: the computer name for which we want the configuration. """ - import paramiko config = paramiko.SSHConfig() try: with open(os.path.expanduser('~/.ssh/config'), encoding='utf8') as fhandle: @@ -118,12 +119,29 @@ class SshTransport(Transport): # pylint: disable=too-many-public-methods 'non_interactive_default': True } ), + ( + 'proxy_jump', { + 'prompt': + 'SSH proxy jump', + 'help': + 'SSH proxy jump for tunneling through other SSH hosts.' + ' Use a comma-separated list of hosts of the form [user@]host[:port].' + ' If user or port are not specified for a host, the user & port values from the target host are used.' + ' This option must be provided explicitly and is not parsed from the SSH config file when left empty.', + 'non_interactive_default': + True + } + ), # Managed 'manually' in connect ( 'proxy_command', { - 'prompt': 'SSH proxy command', - 'help': 'SSH proxy command for tunneling through a proxy server.' + 'prompt': + 'SSH proxy command', + 'help': + 'SSH proxy command for tunneling through a proxy server.' + ' For tunneling through another SSH host, consider using the "SSH proxy jump" option instead!' ' Leave empty to parse the proxy command from the SSH config file.', - 'non_interactive_default': True + 'non_interactive_default': + True } ), # Managed 'manually' in connect ( @@ -309,6 +327,13 @@ def _get_proxy_command_suggestion_string(cls, computer): return ' '.join(new_pieces) + @classmethod + def _get_proxy_jump_suggestion_string(cls, _): + """ + Return an empty suggestion since Paramiko does not parse ProxyJump from the SSH config. + """ + return '' + @classmethod def _get_compress_suggestion_string(cls, computer): # pylint: disable=unused-argument """ @@ -377,11 +402,11 @@ def __init__(self, *args, **kwargs): function (as port, username, password, ...); taken from the accepted paramiko.SSHClient.connect() params. """ - import paramiko super().__init__(*args, **kwargs) self._sftp = None self._proxy = None + self._proxies = [] self._machine = kwargs.pop('machine') @@ -410,7 +435,7 @@ def __init__(self, *args, **kwargs): except KeyError: pass - def open(self): + def open(self): # pylint: disable=too-many-branches,too-many-statements """ Open a SSHClient to the machine possibly using the parameters given in the __init__. @@ -420,6 +445,7 @@ def open(self): :raise aiida.common.InvalidOperation: if the channel is already open """ + from paramiko.ssh_exception import SSHException from aiida.common.exceptions import InvalidOperation from aiida.transports.util import _DetachedProxyCommand @@ -429,9 +455,65 @@ def open(self): connection_arguments = self._connect_args.copy() if 'key_filename' in connection_arguments and not connection_arguments['key_filename']: connection_arguments.pop('key_filename') - proxystring = connection_arguments.pop('proxy_command', None) - if proxystring: - self._proxy = _DetachedProxyCommand(proxystring) + + proxyjumpstring = connection_arguments.pop('proxy_jump', None) + proxycmdstring = connection_arguments.pop('proxy_command', None) + + if proxyjumpstring and proxycmdstring: + raise ValueError('The SSH proxy jump and SSH proxy command options can not be used together') + + if proxyjumpstring: + matcher = re.compile(r'^(?:(?P[^@]+)@)?(?P[^@:]+)(?::(?P\d+))?\s*$') + try: + # don't use a generator here to have everything evaluated + proxies = [matcher.match(s).groupdict() for s in proxyjumpstring.split(',')] + except AttributeError: + raise ValueError('The given configuration for the SSH proxy jump option could not be parsed') + + # proxy_jump supports a list of jump hosts, each jump host is another Paramiko SSH connection + # but when opening a forward channel on a connection, we have to give the next hop. + # So we go through adjacent pairs and by adding the final target to the list we make it universal. + for proxy, target in zip( + proxies, proxies[1:] + [{ + 'host': self._machine, + 'port': connection_arguments.get('port', 22), + }] + ): + proxy_connargs = connection_arguments.copy() + + if proxy['username']: + proxy_connargs['username'] = proxy['username'] + if proxy['port']: + proxy_connargs['port'] = int(proxy['port']) + if not target['port']: # the target port for the channel can not be None + target['port'] = connection_arguments.get('port', 22) + + proxy_client = paramiko.SSHClient() + if self._load_system_host_keys: + proxy_client.load_system_host_keys() + if self._missing_key_policy == 'RejectPolicy': + proxy_client.set_missing_host_key_policy(paramiko.RejectPolicy()) + elif self._missing_key_policy == 'WarningPolicy': + proxy_client.set_missing_host_key_policy(paramiko.WarningPolicy()) + elif self._missing_key_policy == 'AutoAddPolicy': + proxy_client.set_missing_host_key_policy(paramiko.AutoAddPolicy()) + + try: + proxy_client.connect(proxy['host'], **proxy_connargs) + except Exception as exc: + self.logger.error( + f"Error connecting to proxy '{proxy['host']}' through SSH: [{self.__class__.__name__}] {exc}, " + f'connect_args were: {proxy_connargs}' + ) + self._close_proxies() # close all since we're going to start anew on the next open() (if any) + raise + connection_arguments['sock'] = proxy_client.get_transport().open_channel( + 'direct-tcpip', (target['host'], target['port']), ('', 0) + ) + self._proxies.append(proxy_client) + + if proxycmdstring: + self._proxy = _DetachedProxyCommand(proxycmdstring) connection_arguments['sock'] = self._proxy try: @@ -441,22 +523,14 @@ def open(self): f"Error connecting to '{self._machine}' through SSH: " + f'[{self.__class__.__name__}] {exc}, ' + f'connect_args were: {self._connect_args}' ) + self._close_proxies() raise - # Open also a File transport client. SFTP by default, pure SSH in ssh_only - self.open_file_transport() - - return self - - def open_file_transport(self): - """ - Open the SFTP channel, and handle error by directing customer to try another transport - """ - from aiida.common.exceptions import InvalidOperation - from paramiko.ssh_exception import SSHException + # Open the SFTP channel, and handle error by directing customer to try another transport try: self._sftp = self._client.open_sftp() except SSHException: + self._close_proxies() raise InvalidOperation( 'Error in ssh transport plugin. This may be due to the remote computer not supporting SFTP. ' 'Try setting it up with the aiida.transports:ssh_only transport from the aiida-sshonly plugin instead.' @@ -467,6 +541,21 @@ def open_file_transport(self): # Set the current directory to a explicit path, and not to None self._sftp.chdir(self._sftp.normalize('.')) + return self + + def _close_proxies(self): + """Close all proxy connections (proxy_jump and proxy_command)""" + + # Paramiko only closes the channel when closing the main connection, but not the connection itself. + while self._proxies: + self._proxies.pop().close() + + if self._proxy: + # Paramiko should close this automatically when closing the channel, + # but since the process is started in __init__this might not happen correctly. + self._proxy.close() + self._proxy = None + def close(self): """ Close the SFTP channel, and the SSHClient. @@ -482,6 +571,8 @@ def close(self): self._sftp.close() self._client.close() + self._close_proxies() + self._is_open = False @property @@ -1156,7 +1247,6 @@ def _local_listdir(path, pattern=None): """ if not pattern: return os.listdir(path) - import re if path.startswith('/'): # always this is the case in the local case base_dir = path else: @@ -1177,7 +1267,6 @@ def listdir(self, path='.', pattern=None): """ if not pattern: return self.sftp.listdir(path) - import re if path.startswith('/'): base_dir = path else: @@ -1400,13 +1489,16 @@ def gotocomputer_command(self, remotedir): if 'username' in self._connect_args: further_params.append(f"-l {escape_for_bash(self._connect_args['username'])}") - if 'port' in self._connect_args and self._connect_args['port']: + if self._connect_args.get('port'): further_params.append(f"-p {self._connect_args['port']}") - if 'key_filename' in self._connect_args and self._connect_args['key_filename']: + if self._connect_args.get('key_filename'): further_params.append(f"-i {escape_for_bash(self._connect_args['key_filename'])}") - if 'proxy_command' in self._connect_args and self._connect_args['proxy_command']: + if self._connect_args.get('proxy_jump'): + further_params.append(f"-o ProxyJump={escape_for_bash(self._connect_args['proxy_jump'])}") + + if self._connect_args.get('proxy_command'): further_params.append(f"-o ProxyCommand={escape_for_bash(self._connect_args['proxy_command'])}") further_params_str = ' '.join(further_params) diff --git a/docs/source/howto/ssh.rst b/docs/source/howto/ssh.rst index 326d2a26af..6bf7b09e42 100644 --- a/docs/source/howto/ssh.rst +++ b/docs/source/howto/ssh.rst @@ -23,7 +23,7 @@ Very briefly, first create a new private/public keypair (``aiida``/``aiida.pub`` $ ssh-keygen -t rsa -b 4096 -f ~/.ssh/aiida -Copy the public key to the remote machine, normally this will add the public key to the rmote machine's ``~/.ssh/authorized_keys``: +Copy the public key to the remote machine, normally this will add the public key to the remote machine's ``~/.ssh/authorized_keys``: .. code-block:: console @@ -39,7 +39,7 @@ Add the following lines to your ``~/.ssh/config`` file (or create it, if it does .. note:: - If your cluster needs you to connect to another computer *PROXY* first, you can use the ``proxy_command`` feature of ssh, see :ref:`how-to:ssh:proxy`. + If your cluster needs you to connect to another computer *PROXY* first, you can use the ``ProxyJump`` or ``ProxyCommand`` feature of SSH, see :ref:`how-to:ssh:proxy`. You should now be able to access the remote computer (without the need to type a password) *via*: @@ -185,47 +185,79 @@ Connecting to a remote computer *via* a proxy server ==================================================== Some compute clusters require you to connect to an intermediate server *PROXY*, from which you can then connect to the cluster *TARGET* on which you run your calculations. -This section explains how to use the ``proxy_command`` feature of ``ssh`` in order to make this jump automatically. +This section explains how to use the ``ProxyJump`` or ``ProxyCommand`` feature of ``ssh`` in order to make this jump automatically. .. tip:: - This method can also be used to automatically tunnel into virtual private networks, if you have an account on a proxy/jumphost server with access to the network. - + This method can also be used to avoid having to start a virtual private network (VPN) client if you have an SSH account on a proxy/jumphost server which is accessible from your current network **and** from which you can access the *TARGET* machine directly. SSH configuration ^^^^^^^^^^^^^^^^^ -Edit the ``~/.ssh/config`` file on the computer on which you installed AiiDA (or create it if missing) and add the following lines:: +To decide whether to use the ``ProxyJump`` (recommended) or the ``ProxyCommand`` directive, please check the version of your SSH client first with ``ssh -V``. +The ``ProxyJump`` directive has been added in version 7.3 of OpenSSH, hence if you are using an older version of SSH (on your machine or the *PROXY*) you have to use the older ``ProxyCommand``. + +To setup the proxy configuration with ``ProxyJump``, edit the ``~/.ssh/config`` file on the computer on which you installed AiiDA (or create it if missing) +and add the following lines:: Host SHORTNAME_TARGET Hostname FULLHOSTNAME_TARGET User USER_TARGET IdentityFile ~/.ssh/aiida - ProxyCommand ssh -W %h:%p USER_PROXY@FULLHOSTNAME_PROXY + ProxyJump USER_PROXY@FULLHOSTNAME_PROXY + + Host FULLHOSTNAME_PROXY + IdentityFile ~/.ssh/aiida + +Replace the ``..._TARGET`` and ``..._PROXY`` variables with the host/user names of the respective servers. + +.. dropdown:: :fa:`plus-circle` Alternative setup with ``ProxyCommand`` + + To setup the proxy configuration with ``ProxyCommand`` **instead**, edit the ``~/.ssh/config`` file on the computer on which you installed AiiDA (or create it if missing) + and add the following lines:: + + Host SHORTNAME_TARGET + Hostname FULLHOSTNAME_TARGET + User USER_TARGET + IdentityFile ~/.ssh/aiida + ProxyCommand ssh -W %h:%p USER_PROXY@FULLHOSTNAME_PROXY -replacing the ``..._TARGET`` and ``..._PROXY`` variables with the host/user names of the respective servers. + Host FULLHOSTNAME_PROXY + IdentityFile ~/.ssh/aiida -This should allow you to directly connect to the *TARGET* server using + Replace the ``..._TARGET`` and ``..._PROXY`` variables with the host/user names of the respective servers. + +In both cases, this should allow you to directly connect to the *TARGET* server using .. code-block:: console $ ssh SHORTNAME_TARGET -For a *passwordless* connection, you need to follow the instructions :ref:`how-to:ssh:passwordless` *twice*: once for the connection from your computer to the *PROXY* server, and once for the connection from the *PROXY* server to the *TARGET* server. -.. dropdown:: Specifying an SSH key for the proxy - If you need to specify a separate SSH key for the proxy, provide it *after* the ``-W`` directive, e.g.:: +.. note :: + + If the user directory is not shared between the *PROXY* and the *TARGET* (in most supercomputing facilities your user directory is shared between the machines), you need to follow the :ref:`instructions for a passwordless connection ` *twice*: once for the connection from your computer to the *PROXY* server, and once for the connection from the *PROXY* server to the *TARGET* server (e.g. the public key must be listed in the ``~/.ssh/authorized_keys`` file of both the *PROXY* and the *TARGET* server). - ssh -W fidis.epfl.ch:22 -i /home/ubuntu/.ssh/proxy user@proxy.epfl.ch AiiDA configuration ^^^^^^^^^^^^^^^^^^^ -When :ref:`configuring the computer in AiiDA `, AiiDA will automatically parse the required information from your ``~/.ssh/config`` file. +When :ref:`configuring the computer in AiiDA `, AiiDA will automatically parse most of required information from your ``~/.ssh/config`` file. A notable exception to this is the ``proxy_jump`` directive, which **must** be specified manually. + +Simply copy & paste the same instructions as you have used for ``ProxyJump`` in your ``~/.ssh/config`` to the input for ``proxy_jump``: + +.. code-block:: console + + $ verdi computer configure ssh SHORTNAME_TARGET + ... + Allow ssh agent [True]: + SSH proxy jump []: USER_PROXY@FULLHOSTNAME_PROXY + +.. note:: A chain of proxies can be specified as a comma-separated list. If you need to specify a different username, you can so with ``USER_PROXY@...``. If no username is specified for the proxy the same username as for the *TARGET* is used. -.. dropdown:: Specifying the proxy_command manually +.. important:: Specifying the ``proxy_command`` manually When specifying or updating the ``proxy_command`` option via ``verdi computer configure ssh``, please **do not use placeholders** ``%h`` and ``%p`` but provide the *actual* hostname and port. AiiDA replaces them only when parsing from the ``~/.ssh/config`` file. diff --git a/tests/transports/test_ssh.py b/tests/transports/test_ssh.py index 98dcf8a113..8b2043f878 100644 --- a/tests/transports/test_ssh.py +++ b/tests/transports/test_ssh.py @@ -44,6 +44,55 @@ def test_auto_add_policy(): with SshTransport(machine='localhost', timeout=30, load_system_host_keys=True, key_policy='AutoAddPolicy'): pass + @staticmethod + def test_proxy_jump(): + """Test the connection with a proxy jump or several""" + with SshTransport( + machine='localhost', + proxy_jump='localhost', + timeout=30, + load_system_host_keys=True, + key_policy='AutoAddPolicy' + ): + pass + + # kind of pointless, but should work and to check that proxy chaining works + with SshTransport( + machine='localhost', + proxy_jump='localhost,localhost,localhost', + timeout=30, + load_system_host_keys=True, + key_policy='AutoAddPolicy' + ): + pass + + def test_proxy_jump_invalid(self): + """Test proper error reporting when invalid host as a proxy""" + + # import is also that when Python is running with debug warnings `-Wd` + # no unclosed files are reported. + with self.assertRaises(paramiko.SSHException): + with SshTransport( + machine='localhost', + proxy_jump='localhost,nohost', + timeout=30, + load_system_host_keys=True, + key_policy='AutoAddPolicy' + ): + pass + + @staticmethod + def test_proxy_command(): + """Test the connection with a proxy command""" + with SshTransport( + machine='localhost', + proxy_command='ssh -W localhost:22 localhost', + timeout=30, + load_system_host_keys=True, + key_policy='AutoAddPolicy' + ): + pass + def test_no_host_key(self): """Test if there is no host key.""" # Disable logging to avoid output during test @@ -74,3 +123,22 @@ def test_gotocomputer(): """echo ' ** /remote_dir/' ; echo ' ** seems to have been deleted, I logout...' ; fi" """ ) assert cmd_str == expected_str + + +def test_gotocomputer_proxyjump(): + """Test gotocomputer""" + with SshTransport( + machine='localhost', + timeout=30, + use_login_shell=False, + key_policy='AutoAddPolicy', + proxy_jump='localhost', + ) as transport: + cmd_str = transport.gotocomputer_command('/remote_dir/') + + expected_str = ( + """ssh -t localhost -o ProxyJump='localhost' "if [ -d '/remote_dir/' ] ;""" + """ then cd '/remote_dir/' ; bash ; else echo ' ** The directory' ; """ + """echo ' ** /remote_dir/' ; echo ' ** seems to have been deleted, I logout...' ; fi" """ + ) + assert cmd_str == expected_str