From 1de0bee61364d1366fa62d57d81b0c262cdb1826 Mon Sep 17 00:00:00 2001 From: Manuel Giffels Date: Thu, 23 May 2024 16:52:24 +0200 Subject: [PATCH 01/10] Add MFA support for SSHExecutor --- docs/source/changelog.rst | 2 +- setup.py | 1 + tardis/utilities/executors/sshexecutor.py | 56 ++++++++++++++++++++++- 3 files changed, 56 insertions(+), 3 deletions(-) diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst index 90535446..567415ec 100644 --- a/docs/source/changelog.rst +++ b/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/' diff --git a/setup.py b/setup.py index c949c69f..f26d6fb4 100644 --- a/setup.py +++ b/setup.py @@ -95,6 +95,7 @@ def get_cryptography_version(): "typing_extensions", "python-auditor==0.5.0", "tzlocal", + "pyotp", *REST_REQUIRES, ], extras_require={ diff --git a/tardis/utilities/executors/sshexecutor.py b/tardis/utilities/executors/sshexecutor.py index d0d4e2a1..515c23f3 100644 --- a/tardis/utilities/executors/sshexecutor.py +++ b/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,9 +89,14 @@ 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, @@ -52,7 +104,7 @@ async def _establish_connection(self): BrokenPipeError, ): await asyncio.sleep(retry * 10) - return await asyncssh.connect(**self._parameters) + return await asyncssh.connect(client_factory=client_factory, **self._parameters) @property @asynccontextmanager From 6f108fe261666cb0ecb70e27005e2327b7e6be3a Mon Sep 17 00:00:00 2001 From: Manuel Giffels Date: Thu, 23 May 2024 17:31:46 +0200 Subject: [PATCH 02/10] Adjust existing unittests --- tests/utilities_t/executors_t/test_sshexecutor.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/tests/utilities_t/executors_t/test_sshexecutor.py b/tests/utilities_t/executors_t/test_sshexecutor.py index c789ece6..a73effd1 100644 --- a/tests/utilities_t/executors_t/test_sshexecutor.py +++ b/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"], ) From 50a571ee9bc759710f0c71bd05d36a2395d92c31 Mon Sep 17 00:00:00 2001 From: Manuel Giffels Date: Thu, 23 May 2024 17:51:04 +0200 Subject: [PATCH 03/10] Add documentation for MFA support --- docs/source/conf.py | 2 +- docs/source/executors/executors.rst | 22 ++++++++++++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/docs/source/conf.py b/docs/source/conf.py index 7665a165..38162b7a 100644 --- a/docs/source/conf.py +++ b/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. diff --git a/docs/source/executors/executors.rst b/docs/source/executors/executors.rst index 6169f8d9..63a18849 100644 --- a/docs/source/executors/executors.rst +++ b/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 From 61ba55aa74b9aab0ed8e0ca6f63aabc88758ee1f Mon Sep 17 00:00:00 2001 From: Manuel Giffels Date: Fri, 24 May 2024 10:18:37 +0200 Subject: [PATCH 04/10] Revert "Adjust existing unittests" This reverts commit 6f108fe261666cb0ecb70e27005e2327b7e6be3a. --- tests/utilities_t/executors_t/test_sshexecutor.py | 14 +++----------- 1 file changed, 3 insertions(+), 11 deletions(-) diff --git a/tests/utilities_t/executors_t/test_sshexecutor.py b/tests/utilities_t/executors_t/test_sshexecutor.py index a73effd1..c789ece6 100644 --- a/tests/utilities_t/executors_t/test_sshexecutor.py +++ b/tests/utilities_t/executors_t/test_sshexecutor.py @@ -102,9 +102,7 @@ def test_establish_connection(self): run_async(self.executor._establish_connection), MockConnection ) - self.mock_asyncssh.connect.assert_called_with( - client_factory=None, **self.test_asyncssh_params - ) + self.mock_asyncssh.connect.assert_called_with(**self.test_asyncssh_params) test_exceptions = [ ConnectionResetError(), @@ -165,10 +163,7 @@ 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( - client_factory=None, - host="test_host", - username="test", - client_keys=["TestKey"], + host="test_host", username="test", client_keys=["TestKey"] ) self.mock_asyncssh.reset_mock() @@ -228,8 +223,5 @@ def test_construction_by_yaml(self): "Test", ) self.mock_asyncssh.connect.assert_called_with( - client_factory=None, - host="test_host", - username="test", - client_keys=["TestKey"], + host="test_host", username="test", client_keys=["TestKey"] ) From 85796823e682c971a8a944fb4cd5fb49b73a0d6e Mon Sep 17 00:00:00 2001 From: Manuel Giffels Date: Fri, 24 May 2024 10:19:25 +0200 Subject: [PATCH 05/10] Refactor implementation --- docs/source/changelog.rst | 2 +- tardis/utilities/executors/sshexecutor.py | 17 ++++++++--------- 2 files changed, 9 insertions(+), 10 deletions(-) diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst index 567415ec..56801890 100644 --- a/docs/source/changelog.rst +++ b/docs/source/changelog.rst @@ -1,4 +1,4 @@ -.. Created by changelog.py at 2024-05-23, 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/' diff --git a/tardis/utilities/executors/sshexecutor.py b/tardis/utilities/executors/sshexecutor.py index 515c23f3..d38bfa5b 100644 --- a/tardis/utilities/executors/sshexecutor.py +++ b/tardis/utilities/executors/sshexecutor.py @@ -80,8 +80,12 @@ async def kbdint_challenge_received( @enable_yaml_load("!SSHExecutor") class SSHExecutor(Executor): def __init__(self, **parameters): - self._mfa_secrets = parameters.pop("mfa_secrets", None) - self._parameters = parameters + self._parameters = dict(parameters) + # enable Multi-factor Authentication if required + if mfa_secrets := self._parameters.pop("mfa_secrets", None): + self._parameters["client_factory"] = partial( + MFASSHClient, mfa_secrets=mfa_secrets + ) # 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 @@ -89,14 +93,9 @@ 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( - client_factory=client_factory, **self._parameters - ) + return await asyncssh.connect(**self._parameters) except ( ConnectionResetError, asyncssh.DisconnectError, @@ -104,7 +103,7 @@ async def _establish_connection(self): BrokenPipeError, ): await asyncio.sleep(retry * 10) - return await asyncssh.connect(client_factory=client_factory, **self._parameters) + return await asyncssh.connect(**self._parameters) @property @asynccontextmanager From b4c2312169dae4a2484a1f25ce4c302ddb1bfb27 Mon Sep 17 00:00:00 2001 From: Manuel Giffels Date: Fri, 24 May 2024 13:07:25 +0200 Subject: [PATCH 06/10] Add a bit error handling --- tardis/utilities/executors/sshexecutor.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/tardis/utilities/executors/sshexecutor.py b/tardis/utilities/executors/sshexecutor.py index d38bfa5b..7383f8d6 100644 --- a/tardis/utilities/executors/sshexecutor.py +++ b/tardis/utilities/executors/sshexecutor.py @@ -1,11 +1,13 @@ 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 import asyncio import asyncssh +import logging import pyotp from asyncssh.auth import KbdIntPrompts, KbdIntResponse from asyncssh.client import SSHClient @@ -19,6 +21,9 @@ from functools import partial +logger = logging.getLogger("cobald.runtime.tardis.utilities.executors.sshexecutor") + + async def probe_max_session(connection: asyncssh.SSHClientConnection): """ Probe the sshd `MaxSessions`, i.e. the multiplexing limit per connection @@ -74,7 +79,12 @@ async def kbdint_challenge_received( authentication should be attempted. """ # prompts is of type Sequence[Tuple[str, bool]] - return [self._mfa_responses[prompt[0].strip()].now() for prompt in prompts] + 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") From d9e781ac5bbf88b3fc853e00c72760652bfa9a85 Mon Sep 17 00:00:00 2001 From: Manuel Giffels Date: Fri, 24 May 2024 13:08:15 +0200 Subject: [PATCH 07/10] Enable eager evaluation of yaml constructors --- tardis/configuration/utilities.py | 6 ++++-- tardis/utilities/executors/sshexecutor.py | 2 ++ 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/tardis/configuration/utilities.py b/tardis/configuration/utilities.py index aaf84e64..5e96962a 100644 --- a/tardis/configuration/utilities.py +++ b/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 diff --git a/tardis/utilities/executors/sshexecutor.py b/tardis/utilities/executors/sshexecutor.py index 7383f8d6..d50333ec 100644 --- a/tardis/utilities/executors/sshexecutor.py +++ b/tardis/utilities/executors/sshexecutor.py @@ -4,6 +4,7 @@ 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 @@ -88,6 +89,7 @@ async def kbdint_challenge_received( @enable_yaml_load("!SSHExecutor") +@yaml_tag(eager=True) class SSHExecutor(Executor): def __init__(self, **parameters): self._parameters = dict(parameters) From ce68e203d5e5acc12fa6f7362de7cad69c0d4898 Mon Sep 17 00:00:00 2001 From: Manuel Giffels Date: Fri, 24 May 2024 13:09:50 +0200 Subject: [PATCH 08/10] Add tests for MFA support --- .../executors_t/test_sshexecutor.py | 106 +++++++++++++++++- 1 file changed, 100 insertions(+), 6 deletions(-) diff --git a/tests/utilities_t/executors_t/test_sshexecutor.py b/tests/utilities_t/executors_t/test_sshexecutor.py index c789ece6..4f4ae296 100644 --- a/tests/utilities_t/executors_t/test_sshexecutor.py +++ b/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_secrets = [ + { + "prompt": "Enter MFA token:", + "secret": "EJL2DAWFOH7QPJ3D6I2DK2ARTBEJDBIB", + }, + { + "prompt": "Yet another token:", + "secret": "D22246GDKKEDK7AAM77ZH5VRDRL7Z6W7", + }, + ] + self.mfa_ssh_client = MFASSHClient(mfa_secrets=mfa_secrets) + + 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_secrets: + - prompt: 'Token: ' + secret: 123TopSecret + """ + ) + + test_yaml_construction( + mfa_executor, + host="test_host", + username="test", + client_keys=["TestKey"], + client_factory=mfa_executor._parameters["client_factory"], ) From 2858c8111e4e298843f43504cd554c00395e14f5 Mon Sep 17 00:00:00 2001 From: Manuel Giffels Date: Fri, 24 May 2024 15:24:28 +0200 Subject: [PATCH 09/10] Remove explicit copy of parameters dict MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Max Kühn --- tardis/utilities/executors/sshexecutor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tardis/utilities/executors/sshexecutor.py b/tardis/utilities/executors/sshexecutor.py index d50333ec..c8df5394 100644 --- a/tardis/utilities/executors/sshexecutor.py +++ b/tardis/utilities/executors/sshexecutor.py @@ -92,7 +92,7 @@ async def kbdint_challenge_received( @yaml_tag(eager=True) class SSHExecutor(Executor): def __init__(self, **parameters): - self._parameters = dict(parameters) + self._parameters = parameters # enable Multi-factor Authentication if required if mfa_secrets := self._parameters.pop("mfa_secrets", None): self._parameters["client_factory"] = partial( From 23aba4a0ebf40fe4d1377e7909bf4a82ef3ea543 Mon Sep 17 00:00:00 2001 From: Manuel Giffels Date: Fri, 24 May 2024 15:40:34 +0200 Subject: [PATCH 10/10] Rename the mfa paramter to mfa_config --- docs/source/executors/executors.rst | 6 +++--- tardis/utilities/executors/sshexecutor.py | 12 +++++------- tests/utilities_t/executors_t/test_sshexecutor.py | 12 ++++++------ 3 files changed, 14 insertions(+), 16 deletions(-) diff --git a/docs/source/executors/executors.rst b/docs/source/executors/executors.rst index 63a18849..741b18be 100644 --- a/docs/source/executors/executors.rst +++ b/docs/source/executors/executors.rst @@ -47,7 +47,7 @@ SSH Executor `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 + add ``mfa_config`` as parameter to the ``SSHExecutor`` containing a list of command line prompt to TOTP secrets mappings. .. note:: @@ -77,9 +77,9 @@ SSH Executor username: clown client_keys: - /opt/tardis/ssh/tardis - mfa_secrets: + mfa_config: - prompt: "Enter 2FA Token:" - secret: "IMIZDDO2I45ZSTR6XDGFSPFDUY" + totp: "IMIZDDO2I45ZSTR6XDGFSPFDUY" .. rubric:: Example configuration (`COBalD` legacy object initialisation) diff --git a/tardis/utilities/executors/sshexecutor.py b/tardis/utilities/executors/sshexecutor.py index c8df5394..9bd600a4 100644 --- a/tardis/utilities/executors/sshexecutor.py +++ b/tardis/utilities/executors/sshexecutor.py @@ -45,13 +45,11 @@ async def probe_max_session(connection: asyncssh.SSHClientConnection): class MFASSHClient(SSHClient): - def __init__(self, *args, mfa_secrets, **kwargs): + def __init__(self, *args, mfa_config, **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"] - ) + for entry in mfa_config: + self._mfa_responses[entry["prompt"].strip()] = pyotp.TOTP(entry["totp"]) async def kbdint_auth_requested(self) -> MaybeAwait[Optional[str]]: """ @@ -94,9 +92,9 @@ class SSHExecutor(Executor): def __init__(self, **parameters): self._parameters = parameters # enable Multi-factor Authentication if required - if mfa_secrets := self._parameters.pop("mfa_secrets", None): + if mfa_config := self._parameters.pop("mfa_config", None): self._parameters["client_factory"] = partial( - MFASSHClient, mfa_secrets=mfa_secrets + MFASSHClient, mfa_config=mfa_config ) # the current SSH connection or None if it must be (re-)established self._ssh_connection: Optional[asyncssh.SSHClientConnection] = None diff --git a/tests/utilities_t/executors_t/test_sshexecutor.py b/tests/utilities_t/executors_t/test_sshexecutor.py index 4f4ae296..5a490015 100644 --- a/tests/utilities_t/executors_t/test_sshexecutor.py +++ b/tests/utilities_t/executors_t/test_sshexecutor.py @@ -75,17 +75,17 @@ def test_max_sessions(self): class TestMFASSHClient(TestCase): def setUp(self): - mfa_secrets = [ + mfa_config = [ { "prompt": "Enter MFA token:", - "secret": "EJL2DAWFOH7QPJ3D6I2DK2ARTBEJDBIB", + "totp": "EJL2DAWFOH7QPJ3D6I2DK2ARTBEJDBIB", }, { "prompt": "Yet another token:", - "secret": "D22246GDKKEDK7AAM77ZH5VRDRL7Z6W7", + "totp": "D22246GDKKEDK7AAM77ZH5VRDRL7Z6W7", }, ] - self.mfa_ssh_client = MFASSHClient(mfa_secrets=mfa_secrets) + 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), "") @@ -306,9 +306,9 @@ def test_yaml_construction(test_executor, *args, **kwargs): username: test client_keys: - TestKey - mfa_secrets: + mfa_config: - prompt: 'Token: ' - secret: 123TopSecret + totp: 123TopSecret """ )