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-24, 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_config`` 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_config:
+          - prompt: "Enter 2FA Token:"
+            totp: "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/configuration/utilities.py

@@ -1,17 +1,19 @@
+from cobald.daemon.plugins import YAMLTagSettings
 import yaml
 
 
 def enable_yaml_load(tag):
     def yaml_load_decorator(cls):
         def class_factory(loader, node):
+            settings = YAMLTagSettings.fetch(cls)
             new_cls = cls
             if isinstance(node, yaml.nodes.MappingNode):
-                parameters = loader.construct_mapping(node)
+                parameters = loader.construct_mapping(node, deep=settings.eager)
                 new_cls = cls(**parameters)
             elif isinstance(node, yaml.nodes.ScalarNode):
                 new_cls = cls()
             elif isinstance(node, yaml.nodes.SequenceNode):
-                parameters = loader.construct_sequence(node)
+                parameters = loader.construct_sequence(node, deep=settings.eager)
                 new_cls = cls(*parameters)
             return new_cls
 

tardis/utilities/executors/sshexecutor.py

@@ -1,16 +1,29 @@
 from typing import Optional
 from ...configuration.utilities import enable_yaml_load
+from ...exceptions.tardisexceptions import TardisAuthError
 from ...exceptions.executorexceptions import CommandExecutionFailure
 from ...interfaces.executor import Executor
 from ..attributedict import AttributeDict
+from cobald.daemon.plugins import yaml_tag
 
 import asyncio
 import asyncssh
+import logging
+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
+
+
+logger = logging.getLogger("cobald.runtime.tardis.utilities.executors.sshexecutor")
+
 
 async def probe_max_session(connection: asyncssh.SSHClientConnection):
     """
@@ -31,10 +44,58 @@ async def probe_max_session(connection: asyncssh.SSHClientConnection):
     return sessions
 
 
+class MFASSHClient(SSHClient):
+    def __init__(self, *args, mfa_config, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._mfa_responses = {}
+        for entry in mfa_config:
+            self._mfa_responses[entry["prompt"].strip()] = pyotp.TOTP(entry["totp"])
+
+    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]]
+        try:
+            return [self._mfa_responses[prompt[0].strip()].now() for prompt in prompts]
+        except KeyError as ke:
+            msg = f"Keyboard interactive authentication failed: Unexpected Prompt {ke}"
+            logger.error(msg)
+            raise TardisAuthError(msg) from ke
+
+
 @enable_yaml_load("!SSHExecutor")
+@yaml_tag(eager=True)
 class SSHExecutor(Executor):
     def __init__(self, **parameters):
         self._parameters = parameters
+        # enable Multi-factor Authentication if required
+        if mfa_config := self._parameters.pop("mfa_config", None):
+            self._parameters["client_factory"] = partial(
+                MFASSHClient, mfa_config=mfa_config
+            )
         # the current SSH connection or None if it must be (re-)established
         self._ssh_connection: Optional[asyncssh.SSHClientConnection] = None
         # the bound on MaxSession running concurrently

tests/utilities_t/executors_t/test_sshexecutor.py

@@ -1,7 +1,12 @@
 from tests.utilities.utilities import async_return, run_async
 from tardis.utilities.attributedict import AttributeDict
-from tardis.utilities.executors.sshexecutor import SSHExecutor, probe_max_session
+from tardis.utilities.executors.sshexecutor import (
+    SSHExecutor,
+    probe_max_session,
+    MFASSHClient,
+)
 from tardis.exceptions.executorexceptions import CommandExecutionFailure
+from tardis.exceptions.tardisexceptions import TardisAuthError
 
 from asyncssh import ChannelOpenError, ConnectionLost, DisconnectError, ProcessError
 
@@ -11,6 +16,7 @@
 import asyncio
 import yaml
 import contextlib
+import logging
 from asyncstdlib import contextmanager as asynccontextmanager
 
 
@@ -67,6 +73,63 @@ def test_max_sessions(self):
                 )
 
 
+class TestMFASSHClient(TestCase):
+    def setUp(self):
+        mfa_config = [
+            {
+                "prompt": "Enter MFA token:",
+                "totp": "EJL2DAWFOH7QPJ3D6I2DK2ARTBEJDBIB",
+            },
+            {
+                "prompt": "Yet another token:",
+                "totp": "D22246GDKKEDK7AAM77ZH5VRDRL7Z6W7",
+            },
+        ]
+        self.mfa_ssh_client = MFASSHClient(mfa_config=mfa_config)
+
+    def test_kbdint_auth_requested(self):
+        self.assertEqual(run_async(self.mfa_ssh_client.kbdint_auth_requested), "")
+
+    def test_kbdint_challenge_received(self):
+        def test_responses(prompts, num_of_expected_responses):
+            responses = run_async(
+                self.mfa_ssh_client.kbdint_challenge_received,
+                name="test",
+                instructions="no",
+                lang="en",
+                prompts=prompts,
+            )
+
+            self.assertEqual(len(responses), num_of_expected_responses)
+            for response in responses:
+                self.assertTrue(response.isdigit())
+
+        for prompts, num_of_expected_responses in (
+            ([("Enter MFA token:", False)], 1),
+            ([("Enter MFA token:", False), ("Yet another token: ", False)], 2),
+            ([], 0),
+        ):
+            test_responses(
+                prompts=prompts, num_of_expected_responses=num_of_expected_responses
+            )
+
+        prompts_to_fail = [("Enter MFA token:", False), ("Unknown token: ", False)]
+
+        with self.assertRaises(TardisAuthError) as tae:
+            with self.assertLogs(level=logging.ERROR):
+                run_async(
+                    self.mfa_ssh_client.kbdint_challenge_received,
+                    name="test",
+                    instructions="no",
+                    lang="en",
+                    prompts=prompts_to_fail,
+                )
+        self.assertIn(
+            "Keyboard interactive authentication failed: Unexpected Prompt",
+            str(tae.exception),
+        )
+
+
 class TestSSHExecutor(TestCase):
     mock_asyncssh = None
 
@@ -208,6 +271,17 @@ def test_run_command(self):
             run_async(raising_executor.run_command, command="Test", stdin_input="Test")
 
     def test_construction_by_yaml(self):
+        def test_yaml_construction(test_executor, *args, **kwargs):
+            self.assertEqual(
+                run_async(
+                    test_executor.run_command, command="Test", stdin_input="Test"
+                ).stdout,
+                "Test",
+            )
+            self.mock_asyncssh.connect.assert_called_with(*args, **kwargs)
+
+            self.mock_asyncssh.reset_mock()
+
         executor = yaml.safe_load(
             """
                    !SSHExecutor
@@ -218,10 +292,30 @@ def test_construction_by_yaml(self):
                    """
         )
 
-        self.assertEqual(
-            run_async(executor.run_command, command="Test", stdin_input="Test").stdout,
-            "Test",
+        test_yaml_construction(
+            executor,
+            host="test_host",
+            username="test",
+            client_keys=["TestKey"],
         )
-        self.mock_asyncssh.connect.assert_called_with(
-            host="test_host", username="test", client_keys=["TestKey"]
+
+        mfa_executor = yaml.safe_load(
+            """
+                   !SSHExecutor
+                   host: test_host
+                   username: test
+                   client_keys:
+                     - TestKey
+                   mfa_config:
+                     - prompt: 'Token: '
+                       totp: 123TopSecret
+                   """
+        )
+
+        test_yaml_construction(
+            mfa_executor,
+            host="test_host",
+            username="test",
+            client_keys=["TestKey"],
+            client_factory=mfa_executor._parameters["client_factory"],
         )