transports/ssh: support proxy_jump

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<username>[^@]+)@)?(?P<host>[^@:]+)(?::(?P<port>\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:
@@ -1334,13 +1423,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)

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 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 <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 (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 [email protected]
 
 AiiDA configuration
 ^^^^^^^^^^^^^^^^^^^
 
-When :ref:`configuring the computer in AiiDA <how-to:run-codes:computer:configuration>`, AiiDA will automatically parse the required information from your ``~/.ssh/config`` file.
+When :ref:`configuring the computer in AiiDA <how-to:run-codes:computer:configuration>`, 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
+.. dropdown:: 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.