From 9e115d7d1b4cdfa1d1f4da2e04d35df68dac7d80 Mon Sep 17 00:00:00 2001 From: Chad Smith Date: Wed, 22 Apr 2020 12:42:31 -0600 Subject: [PATCH] repo: allow for pre_enable messaging interactions (#1036) * repo: allow for pre_enable/pre_disable messaging interactions Consolidate messging hook processing under a single function handle_message_operations This is groundwork for FIPS pre-enable and pre-disable custom messaging and prompts for #1031. * repo: make messaging an instance property so it can use instance vars --- uaclient/entitlements/cc.py | 33 ++++--- uaclient/entitlements/fips.py | 30 ++++--- uaclient/entitlements/repo.py | 59 ++++++++++-- uaclient/entitlements/tests/test_repo.py | 109 ++++++++++++++++++++++- 4 files changed, 199 insertions(+), 32 deletions(-) diff --git a/uaclient/entitlements/cc.py b/uaclient/entitlements/cc.py index 961f3a2718..de7a30a59c 100644 --- a/uaclient/entitlements/cc.py +++ b/uaclient/entitlements/cc.py @@ -1,5 +1,11 @@ from uaclient.entitlements import repo +try: + from typing import Callable, Dict, List, Tuple, Union # noqa +except ImportError: + # typing isn't available on trusty, so ignore its absence + pass + CC_README = "/usr/share/doc/ubuntu-commoncriteria/README" @@ -11,14 +17,19 @@ class CommonCriteriaEntitlement(repo.RepoEntitlement): description = "Common Criteria EAL2 Provisioning Packages" repo_key_file = "ubuntu-cc-keyring.gpg" packages = ["ubuntu-commoncriteria"] - messaging = { - "pre_install": [ - "(This will download more than 500MB of packages, so may take some" - " time.)" - ], - "post_enable": [ - "Please follow instructions in {} to configure EAL2".format( - CC_README - ) - ], - } + + @property + def messaging( + self + ) -> "Dict[str, List[Union[str, Tuple[Callable, Dict]]]]": + return { + "pre_install": [ + "(This will download more than 500MB of packages, so may take" + " some time.)" + ], + "post_enable": [ + "Please follow instructions in {} to configure EAL2".format( + CC_README + ) + ], + } diff --git a/uaclient/entitlements/fips.py b/uaclient/entitlements/fips.py index 5e51f03295..f3ed3bab48 100644 --- a/uaclient/entitlements/fips.py +++ b/uaclient/entitlements/fips.py @@ -2,7 +2,7 @@ from uaclient import apt, status, util try: - from typing import Dict, List, Set, Tuple # noqa + from typing import Callable, Dict, List, Set, Tuple, Union # noqa except ImportError: # typing isn't available on trusty, so ignore its absence pass @@ -60,25 +60,24 @@ class FIPSEntitlement(FIPSCommonEntitlement): name = "fips" title = "FIPS" description = "NIST-certified FIPS modules" - messaging = { - "post_enable": ["A reboot is required to complete the install"] - } origin = "UbuntuFIPS" static_affordances = ( ("Cannot install FIPS on a container", util.is_container, False), ) + @property + def messaging( + self + ) -> "Dict[str, List[Union[str, Tuple[Callable, Dict]]]]": + return { + "post_enable": ["A reboot is required to complete the install"] + } + class FIPSUpdatesEntitlement(FIPSCommonEntitlement): name = "fips-updates" title = "FIPS Updates" - messaging = { - "post_enable": [ - "FIPS Updates configured and pending, please reboot to make" - " active." - ] - } origin = "UbuntuFIPSUpdates" description = "Uncertified security updates to FIPS modules" static_affordances = ( @@ -88,3 +87,14 @@ class FIPSUpdatesEntitlement(FIPSCommonEntitlement): False, ), ) + + @property + def messaging( + self + ) -> "Dict[str, List[Union[str, Tuple[Callable, Dict]]]]": + return { + "post_enable": [ + "FIPS Updates configured and pending, please reboot to make" + " active." + ] + } diff --git a/uaclient/entitlements/repo.py b/uaclient/entitlements/repo.py index 265cbfce62..6cdca1a420 100644 --- a/uaclient/entitlements/repo.py +++ b/uaclient/entitlements/repo.py @@ -4,7 +4,16 @@ import re try: - from typing import Any, Dict, List, Optional, Tuple, Union # noqa: F401 + from typing import ( # noqa: F401 + Any, + Callable, + Dict, + List, + Optional, + Sequence, + Tuple, + Union, + ) except ImportError: # typing isn't available on trusty, so ignore its absence pass @@ -38,9 +47,13 @@ def repo_pin_priority(self) -> "Union[int, str, None]": def disable_apt_auth_only(self) -> bool: return False # Set True on ESM to only remove apt auth - # Any custom messages to emit pre or post enable or disable operations; - # currently post_enable is used in CommonCriteria - messaging = {} # type: Dict[str, List[str]] + # Any custom messages to emit to the console or callables which are + # handled at pre_enable, pre_disable, pre_install or post_enable stages + @property + def messaging( + self + ) -> "Dict[str, List[Union[str, Tuple[Callable, Dict]]]]": + return {} @property def packages(self) -> "List[str]": @@ -64,12 +77,16 @@ def enable(self, *, silent_if_inapplicable: bool = False) -> bool: """ if not self.can_enable(silent=silent_if_inapplicable): return False + msg_ops = self.messaging.get("pre_enable", []) + if not handle_message_operations(msg_ops): + return False self.setup_apt_config() if self.packages: try: print("Installing {title} packages".format(title=self.title)) - for msg in self.messaging.get("pre_install", []): - print(msg) + msg_ops = self.messaging.get("pre_install", []) + if not handle_message_operations(msg_ops): + return False apt.run_apt_command( ["apt-get", "install", "--assume-yes"] + self.packages, status.MESSAGE_ENABLED_FAILED_TMPL.format( @@ -80,13 +97,17 @@ def enable(self, *, silent_if_inapplicable: bool = False) -> bool: self._cleanup() raise print(status.MESSAGE_ENABLED_TMPL.format(title=self.title)) - for msg in self.messaging.get("post_enable", []): - print(msg) + msg_ops = self.messaging.get("post_enable", []) + if not handle_message_operations(msg_ops): + return False return True def disable(self, silent=False): if not self.can_disable(silent): return False + msg_ops = self.messaging.get("pre_disable", []) + if not handle_message_operations(msg_ops): + return False self._cleanup() return True @@ -301,3 +322,25 @@ def remove_apt_config(self): apt.run_apt_command( ["apt-get", "update"], status.MESSAGE_APT_UPDATE_FAILED ) + + +def handle_message_operations( + msg_ops: "List[Union[str, Tuple[Callable, Dict]]]" +) -> bool: + """Emit messages to the console for user interaction + + :param msg_op: A list of strings or tuples. Any string items are printed. + Any tuples will contain a callable and a dict of args to pass to the + callable. Callables are expected to return True on success and + False upon failure. + + :return: True upon success, False on failure. + """ + for msg_op in msg_ops: + if isinstance(msg_op, str): + print(msg_op) + else: # Then we are a callable and dict of args + functor, args = msg_op + if not functor(**args): + return False + return True diff --git a/uaclient/entitlements/tests/test_repo.py b/uaclient/entitlements/tests/test_repo.py index 3ce494af41..a84ef3cb60 100644 --- a/uaclient/entitlements/tests/test_repo.py +++ b/uaclient/entitlements/tests/test_repo.py @@ -7,7 +7,10 @@ from uaclient import apt from uaclient import config -from uaclient.entitlements.repo import RepoEntitlement +from uaclient.entitlements.repo import ( + RepoEntitlement, + handle_message_operations, +) from uaclient.entitlements.tests.conftest import machine_token from uaclient import exceptions from uaclient import status @@ -318,6 +321,73 @@ def test_enable_passes_silent_if_inapplicable_through( expected_call = mock.call(silent=bool(silent_if_inapplicable)) assert [expected_call] == m_can_enable.call_args_list + @pytest.mark.parametrize( + "pre_enable_msg, output, setup_apt_call_count", + ( + (["msg1", (lambda: False, {}), "msg2"], "msg1\n", 0), + ( + ["msg1", (lambda: True, {}), "msg2"], + "msg1\nmsg2\nRepo Test Class enabled\n", + 1, + ), + ), + ) + @mock.patch.object(RepoTestEntitlement, "setup_apt_config") + @mock.patch.object(RepoTestEntitlement, "can_enable", return_value=True) + def test_enable_can_exit_on_pre_enable_messaging_hooks( + self, + _can_enable, + setup_apt_config, + pre_enable_msg, + output, + setup_apt_call_count, + entitlement, + capsys, + ): + with mock.patch( + M_PATH + "RepoEntitlement.messaging", + new_callable=mock.PropertyMock, + ) as m_messaging: + m_messaging.return_value = {"pre_enable": pre_enable_msg} + with mock.patch.object(type(entitlement), "packages", []): + entitlement.enable() + stdout, _ = capsys.readouterr() + assert output == stdout + assert setup_apt_call_count == setup_apt_config.call_count + + @pytest.mark.parametrize( + "pre_disable_msg, output, remove_apt_call_count", + ( + (["msg1", (lambda: False, {}), "msg2"], "msg1\n", 0), + (["msg1", (lambda: True, {}), "msg2"], "msg1\nmsg2\n", 1), + ), + ) + @mock.patch(M_PATH + "util.subp", return_value=("", "")) + @mock.patch.object(RepoTestEntitlement, "remove_apt_config") + @mock.patch.object(RepoTestEntitlement, "can_disable", return_value=True) + def test_enable_can_exit_on_pre_disable_messaging_hooks( + self, + _can_disable, + remove_apt_config, + m_subp, + pre_disable_msg, + output, + remove_apt_call_count, + entitlement, + capsys, + ): + messaging = {"pre_disable": pre_disable_msg} + with mock.patch.object(type(entitlement), "messaging", messaging): + with mock.patch.object(type(entitlement), "packages", []): + entitlement.disable() + stdout, _ = capsys.readouterr() + assert output == stdout + assert remove_apt_call_count == remove_apt_config.call_count + if remove_apt_call_count > 0: + assert [ + mock.call(["apt-get", "remove", "--assume-yes"]) + ] == m_subp.call_args_list + @pytest.mark.parametrize("with_pre_install_msg", (False, True)) @pytest.mark.parametrize("packages", (["a"], [], None)) @mock.patch(M_PATH + "util.subp", return_value=("", "")) @@ -344,8 +414,10 @@ def test_enable_calls_adds_apt_repo_and_calls_apt_update( pre_install_msgs = ["Some pre-install information", "Some more info"] if with_pre_install_msg: - messaging_patch = mock.patch.object( - entitlement, "messaging", {"pre_install": pre_install_msgs} + messaging_patch = mock.patch( + M_PATH + "RepoEntitlement.messaging", + new_callable=mock.PropertyMock, + return_value={"pre_install": pre_install_msgs}, ) else: messaging_patch = mock.MagicMock() @@ -792,3 +864,34 @@ def test_enabled_status_by_apt_policy( expected_explanation = "Repo Test Class is not configured" assert expected_status == application_status assert expected_explanation == explanation + + +def success_call(): + print("success") + return True + + +def fail_call(a=None): + print("fail %s" % a) + return False + + +class TestHandleMessageOperations: + @pytest.mark.parametrize( + "msg_ops, retval, output", + ( + ([], True, ""), + (["msg1", "msg2"], True, "msg1\nmsg2\n"), + ( + [(success_call, {}), "msg1", (fail_call, {"a": 1}), "msg2"], + False, + "success\nmsg1\nfail 1\n", + ), + ), + ) + def test_handle_message_operations_for_strings_and_callables( + self, msg_ops, retval, output, capsys + ): + assert retval is handle_message_operations(msg_ops) + out, _err = capsys.readouterr() + assert output == out