Add MFA support for SSHExecutor

docs/source/changelog.rst

@@ -1,4 +1,4 @@
-.. Created by changelog.py at 2024-05-15, command
+.. Created by changelog.py at 2024-05-23, command
    '/Users/giffler/.cache/pre-commit/repoecmh3ah8/py_env-python3.12/bin/changelog docs/source/changes compile --categories Added Changed Fixed Security Deprecated --output=docs/source/changelog.rst'
    based on the format of 'https://keepachangelog.com/'
 

docs/source/conf.py

@@ -75,7 +75,7 @@
 #
 # This is also used if you do content translation via gettext catalogs.
 # Usually you set "language" from the command line for these cases.
-language = None
+language = "en"
 
 # List of patterns, relative to source directory, that match files and
 # directories to ignore when looking for source files.

docs/source/executors/executors.rst

@@ -46,6 +46,14 @@ SSH Executor
     directly passed as keyword arguments to `asyncssh` `connect` call. You can find all available parameters in the
     `asyncssh documentation`_
 
+    Additionally the ``SSHExecutor`` supports Multi-factor Authentication (MFA). In order to activate it, you need to
+    add ``mfa_secrets`` as parameter to the ``SSHExecutor`` containing a list of command line prompt to TOTP secrets
+    mappings.
+
+    .. note::
+        The prompt can be obtained by connecting to the server via ssh in a terminal. The prompt is the text the
+        terminal is showing in order to obtain the second factor for the ssh connection. (e.g. "Enter 2FA Token:")
+
     .. _asyncssh documentation: https://asyncssh.readthedocs.io/en/latest/api.html#connect
 
 .. content-tabs:: right-col
@@ -60,6 +68,20 @@ SSH Executor
         client_keys:
           - /opt/tardis/ssh/tardis
 
+    .. rubric:: Example configuration (Using Multi-factor Authentication)
+
+    .. code-block:: yaml
+
+      !TardisSSHExecutor
+        host: login.dorie.somewherein.de
+        username: clown
+        client_keys:
+          - /opt/tardis/ssh/tardis
+        mfa_secrets:
+          - prompt: "Enter 2FA Token:"
+            secret: "IMIZDDO2I45ZSTR6XDGFSPFDUY"
+
+
     .. rubric:: Example configuration (`COBalD` legacy object initialisation)
 
     .. code-block:: yaml

setup.py

@@ -95,6 +95,7 @@ def get_cryptography_version():
         "typing_extensions",
         "python-auditor==0.5.0",
         "tzlocal",
+        "pyotp",
         *REST_REQUIRES,
     ],
     extras_require={

tardis/utilities/executors/sshexecutor.py

@@ -6,11 +6,18 @@
 
 import asyncio
 import asyncssh
+import pyotp
+from asyncssh.auth import KbdIntPrompts, KbdIntResponse
+from asyncssh.client import SSHClient
+from asyncssh.misc import MaybeAwait
+
 from asyncstdlib import (
     ExitStack as AsyncExitStack,
     contextmanager as asynccontextmanager,
 )
 
+from functools import partial
+
 
 async def probe_max_session(connection: asyncssh.SSHClientConnection):
     """
@@ -31,9 +38,49 @@ async def probe_max_session(connection: asyncssh.SSHClientConnection):
     return sessions
 
 
+class MFASSHClient(SSHClient):
+    def __init__(self, *args, mfa_secrets, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._mfa_responses = {}
+        for mfa_secret in mfa_secrets:
+            self._mfa_responses[mfa_secret["prompt"].strip()] = pyotp.TOTP(
+                mfa_secret["secret"]
+            )
+
+    async def kbdint_auth_requested(self) -> MaybeAwait[Optional[str]]:
+        """
+        Keyboard-interactive authentication has been requested
+
+        This method should return a string containing a comma-separated
+        list of submethods that the server should use for
+        keyboard-interactive authentication. An empty string can be
+        returned to let the server pick the type of keyboard-interactive
+        authentication to perform.
+        """
+        return ""
+
+    async def kbdint_challenge_received(
+        self, name: str, instructions: str, lang: str, prompts: KbdIntPrompts
+    ) -> MaybeAwait[Optional[KbdIntResponse]]:
+        """
+        A keyboard-interactive auth challenge has been received
+
+        This method is called when the server sends a keyboard-interactive
+        authentication challenge.
+
+        The return value should be a list of strings of the same length
+        as the number of prompts provided if the challenge can be
+        answered, or `None` to indicate that some other form of
+        authentication should be attempted.
+        """
+        # prompts is of type Sequence[Tuple[str, bool]]
+        return [self._mfa_responses[prompt[0].strip()].now() for prompt in prompts]
+
+
 @enable_yaml_load("!SSHExecutor")
 class SSHExecutor(Executor):
     def __init__(self, **parameters):
+        self._mfa_secrets = parameters.pop("mfa_secrets", None)
         self._parameters = parameters
         # the current SSH connection or None if it must be (re-)established
         self._ssh_connection: Optional[asyncssh.SSHClientConnection] = None
@@ -42,17 +89,22 @@ def __init__(self, **parameters):
         self._lock = None
 
     async def _establish_connection(self):
+        client_factory = None
+        if self._mfa_secrets:
+            client_factory = partial(MFASSHClient, mfa_secrets=self._mfa_secrets)
         for retry in range(1, 10):
             try:
-                return await asyncssh.connect(**self._parameters)
+                return await asyncssh.connect(
+                    client_factory=client_factory, **self._parameters
+                )
             except (
                 ConnectionResetError,
                 asyncssh.DisconnectError,
                 asyncssh.ConnectionLost,
                 BrokenPipeError,
             ):
                 await asyncio.sleep(retry * 10)
-        return await asyncssh.connect(**self._parameters)
+        return await asyncssh.connect(client_factory=client_factory, **self._parameters)
 
     @property
     @asynccontextmanager

tests/utilities_t/executors_t/test_sshexecutor.py

@@ -102,7 +102,9 @@ def test_establish_connection(self):
             run_async(self.executor._establish_connection), MockConnection
         )
 
-        self.mock_asyncssh.connect.assert_called_with(**self.test_asyncssh_params)
+        self.mock_asyncssh.connect.assert_called_with(
+            client_factory=None, **self.test_asyncssh_params
+        )
 
         test_exceptions = [
             ConnectionResetError(),
@@ -163,7 +165,10 @@ async def is_queued(n: int):
     def test_run_command(self):
         self.assertIsNone(run_async(self.executor.run_command, command="Test").stdout)
         self.mock_asyncssh.connect.assert_called_with(
-            host="test_host", username="test", client_keys=["TestKey"]
+            client_factory=None,
+            host="test_host",
+            username="test",
+            client_keys=["TestKey"],
         )
         self.mock_asyncssh.reset_mock()
 
@@ -223,5 +228,8 @@ def test_construction_by_yaml(self):
             "Test",
         )
         self.mock_asyncssh.connect.assert_called_with(
-            host="test_host", username="test", client_keys=["TestKey"]
+            client_factory=None,
+            host="test_host",
+            username="test",
+            client_keys=["TestKey"],
         )