repo: allow for pre_enable messaging interactions

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
+                )
+            ],
+        }

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."
+            ]
+        }

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

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