Enhance verb (#5596)

* Add the cli parameters

* Tests and error messages

* Requested fixes

* Only handle SSL vhosts

* Interactive cert-name selection if not defined on CLI

* Address PR comments

* Address review comments

* Added tests and addressed review comments

* Move cert manager tests to correct file

* Add display ops tests

* Use display util constants instead of hardcoded values in tests

certbot/cert_manager.py

@@ -46,7 +46,7 @@ def rename_lineage(config):
     """
     disp = zope.component.getUtility(interfaces.IDisplay)
 
-    certname = _get_certnames(config, "rename")[0]
+    certname = get_certnames(config, "rename")[0]
 
     new_certname = config.new_certname
     if not new_certname:
@@ -88,7 +88,7 @@ def certificates(config):
 
 def delete(config):
     """Delete Certbot files associated with a certificate lineage."""
-    certnames = _get_certnames(config, "delete", allow_multiple=True)
+    certnames = get_certnames(config, "delete", allow_multiple=True)
     for certname in certnames:
         storage.delete_files(config, certname)
         disp = zope.component.getUtility(interfaces.IDisplay)
@@ -288,11 +288,7 @@ def human_readable_cert_info(config, cert, skip_filter_checks=False):
                          cert.privkey))
     return "".join(certinfo)
 
-###################
-# Private Helpers
-###################
-
-def _get_certnames(config, verb, allow_multiple=False):
+def get_certnames(config, verb, allow_multiple=False, custom_prompt=None):
     """Get certname from flag, interactively, or error out.
     """
     certname = config.certname
@@ -305,22 +301,32 @@ def _get_certnames(config, verb, allow_multiple=False):
         if not choices:
             raise errors.Error("No existing certificates found.")
         if allow_multiple:
+            if not custom_prompt:
+                prompt = "Which certificate(s) would you like to {0}?".format(verb)
+            else:
+                prompt = custom_prompt
             code, certnames = disp.checklist(
-                                    "Which certificate(s) would you like to {0}?".format(verb),
-                                    choices, cli_flag="--cert-name",
-                                    force_interactive=True)
+                prompt, choices, cli_flag="--cert-name", force_interactive=True)
             if code != display_util.OK:
                 raise errors.Error("User ended interaction.")
         else:
-            code, index = disp.menu("Which certificate would you like to {0}?".format(verb),
-                                    choices, cli_flag="--cert-name",
-                                    force_interactive=True)
+            if not custom_prompt:
+                prompt = "Which certificate would you like to {0}?".format(verb)
+            else:
+                prompt = custom_prompt
+
+            code, index = disp.menu(
+                prompt, choices, cli_flag="--cert-name", force_interactive=True)
 
             if code != display_util.OK or index not in range(0, len(choices)):
                 raise errors.Error("User ended interaction.")
             certnames = [choices[index]]
     return certnames
 
+###################
+# Private Helpers
+###################
+
 def _report_lines(msgs):
     """Format a results report for a category of single-line renewal outcomes"""
     return "  " + "\n  ".join(str(msg) for msg in msgs)

certbot/cli.py

@@ -76,6 +76,7 @@ obtain, install, and renew certificates:
     (default) run   Obtain & install a certificate in your current webserver
     certonly        Obtain or renew a certificate, but do not install it
     renew           Renew all previously obtained certificates that are near expiry
+    enhance         Add security enhancements to your existing configuration
    -d DOMAINS       Comma-separated list of domains to obtain a certificate for
 
   %s
@@ -415,6 +416,12 @@ VERB_HELP = [
                   os.path.join(flag_default("config_dir"), "live"))),
         "usage": "\n\n  certbot update_symlinks [options]\n\n"
     }),
+    ("enhance", {
+        "short": "Add security enhancements to your existing configuration",
+        "opts": ("Helps to harden the TLS configration by adding security enhancements "
+                 "to already existing configuration."),
+        "usage": "\n\n  certbot enhance [options]\n\n"
+    }),
 
 ]
 # VERB_HELP is a list in order to preserve order, but a dict is sometimes useful
@@ -449,6 +456,7 @@ class HelpfulArgumentParser(object):
             "update_symlinks": main.update_symlinks,
             "certificates": main.certificates,
             "delete": main.delete,
+            "enhance": main.enhance,
         }
 
         # Get notification function for printing
@@ -883,21 +891,22 @@ def prepare_and_parse_args(plugins, args, detect_defaults=False):  # pylint: dis
              "flag to 0 disables log rotation entirely, causing "
              "Certbot to always append to the same log file.")
     helpful.add(
-        [None, "automation", "run", "certonly"], "-n", "--non-interactive", "--noninteractive",
+        [None, "automation", "run", "certonly", "enhance"],
+        "-n", "--non-interactive", "--noninteractive",
         dest="noninteractive_mode", action="store_true",
         default=flag_default("noninteractive_mode"),
         help="Run without ever asking for user input. This may require "
               "additional command line flags; the client will try to explain "
               "which ones are required if it finds one missing")
     helpful.add(
-        [None, "register", "run", "certonly"],
+        [None, "register", "run", "certonly", "enhance"],
         constants.FORCE_INTERACTIVE_FLAG, action="store_true",
         default=flag_default("force_interactive"),
         help="Force Certbot to be interactive even if it detects it's not "
              "being run in a terminal. This flag cannot be used with the "
              "renew subcommand.")
     helpful.add(
-        [None, "run", "certonly", "certificates"],
+        [None, "run", "certonly", "certificates", "enhance"],
         "-d", "--domains", "--domain", dest="domains",
         metavar="DOMAIN", action=_DomainsAction,
         default=flag_default("domains"),
@@ -913,8 +922,8 @@ def prepare_and_parse_args(plugins, args, detect_defaults=False):  # pylint: dis
              "name. In the case of a name collision it will append a number "
              "like 0001 to the file path name. (default: Ask)")
     helpful.add(
-        [None, "run", "certonly", "manage", "delete", "certificates", "renew"],
-        "--cert-name", dest="certname",
+        [None, "run", "certonly", "manage", "delete", "certificates",
+         "renew", "enhance"], "--cert-name", dest="certname",
         metavar="CERTNAME", default=flag_default("certname"),
         help="Certificate name to apply. This name is used by Certbot for housekeeping "
              "and in file paths; it doesn't affect the content of the certificate itself. "
@@ -1085,7 +1094,8 @@ def prepare_and_parse_args(plugins, args, detect_defaults=False):  # pylint: dis
         dest="must_staple", default=flag_default("must_staple"),
         help=config_help("must_staple"))
     helpful.add(
-        "security", "--redirect", action="store_true", dest="redirect",
+        ["security", "enhance"],
+        "--redirect", action="store_true", dest="redirect",
         default=flag_default("redirect"),
         help="Automatically redirect all HTTP traffic to HTTPS for the newly "
              "authenticated vhost. (default: Ask)")
@@ -1095,7 +1105,8 @@ def prepare_and_parse_args(plugins, args, detect_defaults=False):  # pylint: dis
         help="Do not automatically redirect all HTTP traffic to HTTPS for the newly "
              "authenticated vhost. (default: Ask)")
     helpful.add(
-        "security", "--hsts", action="store_true", dest="hsts", default=flag_default("hsts"),
+        ["security", "enhance"],
+        "--hsts", action="store_true", dest="hsts", default=flag_default("hsts"),
         help="Add the Strict-Transport-Security header to every HTTP response."
              " Forcing browser to always use SSL for the domain."
              " Defends against SSL Stripping.")
@@ -1103,7 +1114,8 @@ def prepare_and_parse_args(plugins, args, detect_defaults=False):  # pylint: dis
         "security", "--no-hsts", action="store_false", dest="hsts",
         default=flag_default("hsts"), help=argparse.SUPPRESS)
     helpful.add(
-        "security", "--uir", action="store_true", dest="uir", default=flag_default("uir"),
+        ["security", "enhance"],
+        "--uir", action="store_true", dest="uir", default=flag_default("uir"),
         help='Add the "Content-Security-Policy: upgrade-insecure-requests"'
              ' header to every HTTP response. Forcing the browser to use'
              ' https:// for every http:// resource.')

certbot/client.py

@@ -466,7 +466,7 @@ class Client(object):
             # sites may have been enabled / final cleanup
             self.installer.restart()
 
-    def enhance_config(self, domains, chain_path):
+    def enhance_config(self, domains, chain_path, ask_redirect=True):
         """Enhance the configuration.
 
         :param list domains: list of domains to configure
@@ -493,8 +493,9 @@ class Client(object):
         for config_name, enhancement_name, option in enhancement_info:
             config_value = getattr(self.config, config_name)
             if enhancement_name in supported:
-                if config_name == "redirect" and config_value is None:
-                    config_value = enhancements.ask(enhancement_name)
+                if ask_redirect:
+                    if config_name == "redirect" and config_value is None:
+                        config_value = enhancements.ask(enhancement_name)
                 if config_value:
                     self.apply_enhancement(domains, enhancement_name, option)
                     enhanced = True
@@ -530,8 +531,12 @@ class Client(object):
                 try:
                     self.installer.enhance(dom, enhancement, options)
                 except errors.PluginEnhancementAlreadyPresent:
-                    logger.warning("Enhancement %s was already set.",
-                            enhancement)
+                    if enhancement == "ensure-http-header":
+                        logger.warning("Enhancement %s was already set.",
+                                options)
+                    else:
+                        logger.warning("Enhancement %s was already set.",
+                                enhancement)
                 except errors.PluginError:
                     logger.warning("Unable to set enhancement %s for %s",
                             enhancement, dom)

certbot/display/ops.py

@@ -86,13 +86,31 @@ def choose_account(accounts):
     else:
         return None
 
+def choose_values(values, question=None):
+    """Display screen to let user pick one or multiple values from the provided
+    list.
 
-def choose_names(installer):
+    :param list values: Values to select from
+
+    :returns: List of selected values
+    :rtype: list
+    """
+    code, items = z_util(interfaces.IDisplay).checklist(
+        question, tags=values, force_interactive=True)
+    if code == display_util.OK and items:
+        return items
+    else:
+        return []
+
+def choose_names(installer, question=None):
     """Display screen to select domains to validate.
 
     :param installer: An installer object
     :type installer: :class:`certbot.interfaces.IInstaller`
 
+    :param `str` question: Overriding dialog question to ask the user if asked
+        to choose from domain names.
+
     :returns: List of selected names
     :rtype: `list` of `str`
 
@@ -108,7 +126,7 @@ def choose_names(installer):
         return _choose_names_manually(
             "No names were found in your configuration files. ")
 
-    code, names = _filter_names(names)
+    code, names = _filter_names(names, question)
     if code == display_util.OK and names:
         return names
     else:
@@ -142,7 +160,7 @@ def _sort_names(FQDNs):
     return sorted(FQDNs, key=lambda fqdn: fqdn.split('.')[::-1][1:])
 
 
-def _filter_names(names):
+def _filter_names(names, override_question=None):
     """Determine which names the user would like to select from a list.
 
     :param list names: domain names
@@ -155,10 +173,12 @@ def _filter_names(names):
     """
     #Sort by domain first, and then by subdomain
     sorted_names = _sort_names(names)
-
+    if override_question:
+        question = override_question
+    else:
+        question = "Which names would you like to activate HTTPS for?"
     code, names = z_util(interfaces.IDisplay).checklist(
-        "Which names would you like to activate HTTPS for?",
-        tags=sorted_names, cli_flag="--domains", force_interactive=True)
+        question, tags=sorted_names, cli_flag="--domains", force_interactive=True)
     return code, [str(s) for s in names]
 
 

certbot/main.py

@@ -382,7 +382,7 @@ def _ask_user_to_confirm_new_names(config, new_domains, certname, old_domains):
     if not obj.yesno(msg, "Update cert", "Cancel", default=True):
         raise errors.ConfigurationError("Specified mismatched cert name and domains.")
 
-def _find_domains_or_certname(config, installer):
+def _find_domains_or_certname(config, installer, question=None):
     """Retrieve domains and certname from config or user input.
 
     :param config: Configuration object
@@ -391,6 +391,8 @@ def _find_domains_or_certname(config, installer):
     :param installer: Installer object
     :type installer: interfaces.IInstaller
 
+    :param `str` question: Overriding dialog question to ask the user if asked
+        to choose from domain names.
 
     :returns: Two-part tuple of domains and certname
     :rtype: `tuple` of list of `str` and `str`
@@ -411,7 +413,7 @@ def _find_domains_or_certname(config, installer):
     # that certname might not have existed, or there was a problem.
     # try to get domains from the user.
     if not domains:
-        domains = display_ops.choose_names(installer)
+        domains = display_ops.choose_names(installer, question)
 
     if not domains and not certname:
         raise errors.Error("Please specify --domains, or --installer that "
@@ -859,6 +861,53 @@ def plugins_cmd(config, plugins):
     logger.debug("Prepared plugins: %s", available)
     notify(str(available))
 
+def enhance(config, plugins):
+    """Add security enhancements to existing configuration
+
+    :param config: Configuration object
+    :type config: interfaces.IConfig
+
+    :param plugins: List of plugins
+    :type plugins: `list` of `str`
+
+    :returns: `None`
+    :rtype: None
+
+    """
+    supported_enhancements = ["hsts", "redirect", "uir", "staple"]
+    # Check that at least one enhancement was requested on command line
+    if not any([getattr(config, enh) for enh in supported_enhancements]):
+        msg = ("Please specify one or more enhancement types to configure. To list "
+               "the available enhancement types, run:\n\n%s --help enhance\n")
+        logger.warning(msg, sys.argv[0])
+        raise errors.MisconfigurationError("No enhancements requested, exiting.")
+
+    try:
+        installer, _ = plug_sel.choose_configurator_plugins(config, plugins, "enhance")
+    except errors.PluginSelectionError as e:
+        return str(e)
+
+    certname_question = ("Which certificate would you like to use to enhance "
+                         "your configuration?")
+    config.certname = cert_manager.get_certnames(
+        config, "enhance", allow_multiple=False,
+        custom_prompt=certname_question)[0]
+    cert_domains = cert_manager.domains_for_certname(config, config.certname)
+    if config.noninteractive_mode:
+        domains = cert_domains
+    else:
+        domain_question = ("Which domain names would you like to enable the "
+                           "selected enhancements for?")
+        domains = display_ops.choose_values(cert_domains, domain_question)
+        if not domains:
+            raise errors.Error("User cancelled the domain selection. No domains "
+                               "defined, exiting.")
+    if not config.chain_path:
+        lineage = cert_manager.lineage_for_certname(config, config.certname)
+        config.chain_path = lineage.chain_path
+    le_client = _init_le_client(config, authenticator=None, installer=installer)
+    le_client.enhance_config(domains, config.chain_path, ask_redirect=False)
+
 
 def rollback(config, plugins):
     """Rollback server configuration changes made during install.

certbot/plugins/selection.py

@@ -147,6 +147,7 @@ def record_chosen_plugins(config, plugins, auth, inst):
 
 
 def choose_configurator_plugins(config, plugins, verb):
+    # pylint: disable=too-many-branches
     """
     Figure out which configurator we're going to use, modifies
     config.authenticator and config.installer strings to reflect that choice if
@@ -159,6 +160,11 @@ def choose_configurator_plugins(config, plugins, verb):
     """
 
     req_auth, req_inst = cli_plugin_requests(config)
+    installer_question = None
+
+    if verb == "enhance":
+        installer_question = ("Which installer would you like to use to "
+                              "configure the selected enhancements?")
 
     # Which plugins do we need?
     if verb == "run":
@@ -176,11 +182,11 @@ def choose_configurator_plugins(config, plugins, verb):
         need_inst = need_auth = False
     if verb == "certonly":
         need_auth = True
-    if verb == "install":
+    if verb == "install" or verb == "enhance":
         need_inst = True
         if config.authenticator:
-            logger.warning("Specifying an authenticator doesn't make sense in install mode")
-
+            logger.warning("Specifying an authenticator doesn't make sense when "
+                           "running Certbot with verb \"%s\"", verb)
     # Try to meet the user's request and/or ask them to pick plugins
     authenticator = installer = None
     if verb == "run" and req_auth == req_inst:
@@ -189,7 +195,7 @@ def choose_configurator_plugins(config, plugins, verb):
         authenticator = installer = pick_configurator(config, req_inst, plugins)
     else:
         if need_inst or req_inst:
-            installer = pick_installer(config, req_inst, plugins)
+            installer = pick_installer(config, req_inst, plugins, installer_question)
         if need_auth:
             authenticator = pick_authenticator(config, req_auth, plugins)
     logger.debug("Selected authenticator %s and installer %s", authenticator, installer)

certbot/tests/cert_manager_test.py

@@ -569,5 +569,103 @@ class MatchAndCheckOverlaps(storage_test.BaseRenewableCertTest):
         self.assertRaises(errors.OverlappingMatchFound, self._call, self.config, None, None, None)
 
 
+class GetCertnameTest(unittest.TestCase):
+    """Tests for certbot.cert_manager."""
+
+    def setUp(self):
+        self.get_utility_patch = test_util.patch_get_utility()
+        self.mock_get_utility = self.get_utility_patch.start()
+        self.config = mock.MagicMock()
+        self.config.certname = None
+
+    def tearDown(self):
+        self.get_utility_patch.stop()
+
+    @mock.patch('certbot.storage.renewal_conf_files')
+    @mock.patch('certbot.storage.lineagename_for_filename')
+    def test_get_certnames(self, mock_name, mock_files):
+        mock_files.return_value = ['example.com.conf']
+        mock_name.return_value = 'example.com'
+        from certbot import cert_manager
+        prompt = "Which certificate would you"
+        self.mock_get_utility().menu.return_value = (display_util.OK, 0)
+        self.assertEquals(
+            cert_manager.get_certnames(
+                self.config, "verb", allow_multiple=False), ['example.com'])
+        self.assertTrue(
+            prompt in self.mock_get_utility().menu.call_args[0][0])
+
+    @mock.patch('certbot.storage.renewal_conf_files')
+    @mock.patch('certbot.storage.lineagename_for_filename')
+    def test_get_certnames_custom_prompt(self, mock_name, mock_files):
+        mock_files.return_value = ['example.com.conf']
+        mock_name.return_value = 'example.com'
+        from certbot import cert_manager
+        prompt = "custom prompt"
+        self.mock_get_utility().menu.return_value = (display_util.OK, 0)
+        self.assertEquals(
+            cert_manager.get_certnames(
+                self.config, "verb", allow_multiple=False, custom_prompt=prompt),
+            ['example.com'])
+        self.assertEquals(self.mock_get_utility().menu.call_args[0][0],
+                          prompt)
+
+    @mock.patch('certbot.storage.renewal_conf_files')
+    @mock.patch('certbot.storage.lineagename_for_filename')
+    def test_get_certnames_user_abort(self, mock_name, mock_files):
+        mock_files.return_value = ['example.com.conf']
+        mock_name.return_value = 'example.com'
+        from certbot import cert_manager
+        self.mock_get_utility().menu.return_value = (display_util.CANCEL, 0)
+        self.assertRaises(
+            errors.Error,
+            cert_manager.get_certnames,
+            self.config, "erroring_anyway", allow_multiple=False)
+
+    @mock.patch('certbot.storage.renewal_conf_files')
+    @mock.patch('certbot.storage.lineagename_for_filename')
+    def test_get_certnames_allow_multiple(self, mock_name, mock_files):
+        mock_files.return_value = ['example.com.conf']
+        mock_name.return_value = 'example.com'
+        from certbot import cert_manager
+        prompt = "Which certificate(s) would you"
+        self.mock_get_utility().checklist.return_value = (display_util.OK,
+                                                          ['example.com'])
+        self.assertEquals(
+            cert_manager.get_certnames(
+                self.config, "verb", allow_multiple=True), ['example.com'])
+        self.assertTrue(
+            prompt in self.mock_get_utility().checklist.call_args[0][0])
+
+    @mock.patch('certbot.storage.renewal_conf_files')
+    @mock.patch('certbot.storage.lineagename_for_filename')
+    def test_get_certnames_allow_multiple_custom_prompt(self, mock_name, mock_files):
+        mock_files.return_value = ['example.com.conf']
+        mock_name.return_value = 'example.com'
+        from certbot import cert_manager
+        prompt = "custom prompt"
+        self.mock_get_utility().checklist.return_value = (display_util.OK,
+                                                          ['example.com'])
+        self.assertEquals(
+            cert_manager.get_certnames(
+                self.config, "verb", allow_multiple=True, custom_prompt=prompt),
+            ['example.com'])
+        self.assertEquals(
+            self.mock_get_utility().checklist.call_args[0][0],
+            prompt)
+
+    @mock.patch('certbot.storage.renewal_conf_files')
+    @mock.patch('certbot.storage.lineagename_for_filename')
+    def test_get_certnames_allow_multiple_user_abort(self, mock_name, mock_files):
+        mock_files.return_value = ['example.com.conf']
+        mock_name.return_value = 'example.com'
+        from certbot import cert_manager
+        self.mock_get_utility().checklist.return_value = (display_util.CANCEL, [])
+        self.assertRaises(
+            errors.Error,
+            cert_manager.get_certnames,
+            self.config, "erroring_anyway", allow_multiple=True)
+
+
 if __name__ == "__main__":
     unittest.main()  # pragma: no cover

certbot/tests/client_test.py

@@ -433,6 +433,22 @@ class EnhanceConfigTest(ClientTestCommon):
         self.client.installer.enhance.assert_not_called()
         mock_enhancements.ask.assert_not_called()
 
+    @mock.patch("certbot.client.logger")
+    def test_already_exists_header(self, mock_log):
+        self.config.hsts = True
+        self._test_with_already_existing()
+        self.assertTrue(mock_log.warning.called)
+        self.assertEquals(mock_log.warning.call_args[0][1],
+                          'Strict-Transport-Security')
+
+    @mock.patch("certbot.client.logger")
+    def test_already_exists_redirect(self, mock_log):
+        self.config.redirect = True
+        self._test_with_already_existing()
+        self.assertTrue(mock_log.warning.called)
+        self.assertEquals(mock_log.warning.call_args[0][1],
+                          'redirect')
+
     def test_no_ask_hsts(self):
         self.config.hsts = True
         self._test_with_all_supported()
@@ -508,6 +524,13 @@ class EnhanceConfigTest(ClientTestCommon):
         self.assertEqual(self.client.installer.save.call_count, 1)
         self.assertEqual(self.client.installer.restart.call_count, 1)
 
+    def _test_with_already_existing(self):
+        self.client.installer = mock.MagicMock()
+        self.client.installer.supported_enhancements.return_value = [
+            "ensure-http-header", "redirect", "staple-ocsp"]
+        self.client.installer.enhance.side_effect = errors.PluginEnhancementAlreadyPresent()
+        self.client.enhance_config([self.domain], None)
+
 
 class RollbackTest(unittest.TestCase):
     """Tests for certbot.client.rollback."""

certbot/tests/display/ops_test.py

@@ -207,9 +207,9 @@ class ChooseNamesTest(unittest.TestCase):
         self.mock_install = mock.MagicMock()
 
     @classmethod
-    def _call(cls, installer):
+    def _call(cls, installer, question=None):
         from certbot.display.ops import choose_names
-        return choose_names(installer)
+        return choose_names(installer, question)
 
     @mock.patch("certbot.display.ops._choose_names_manually")
     def test_no_installer(self, mock_manual):
@@ -281,6 +281,15 @@ class ChooseNamesTest(unittest.TestCase):
         self.assertEqual(names, ["example.com"])
         self.assertEqual(mock_util().checklist.call_count, 1)
 
+    @test_util.patch_get_utility("certbot.display.ops.z_util")
+    def test_filter_namees_override_question(self, mock_util):
+        self.mock_install.get_all_names.return_value = set(["example.com"])
+        mock_util().checklist.return_value = (display_util.OK, ["example.com"])
+        names = self._call(self.mock_install, "Custom")
+        self.assertEqual(names, ["example.com"])
+        self.assertEqual(mock_util().checklist.call_count, 1)
+        self.assertEqual(mock_util().checklist.call_args[0][0], "Custom")
+
     @test_util.patch_get_utility("certbot.display.ops.z_util")
     def test_filter_names_nothing_selected(self, mock_util):
         self.mock_install.get_all_names.return_value = set(["example.com"])
@@ -481,5 +490,42 @@ class ValidatorTests(unittest.TestCase):
                           self.__validator, "msg", default="")
 
 
+class ChooseValuesTest(unittest.TestCase):
+    """Test choose_values."""
+    @classmethod
+    def _call(cls, values, question):
+        from certbot.display.ops import choose_values
+        return choose_values(values, question)
+
+    @test_util.patch_get_utility("certbot.display.ops.z_util")
+    def test_choose_names_success(self, mock_util):
+        items = ["first", "second", "third"]
+        mock_util().checklist.return_value = (display_util.OK, [items[2]])
+        result = self._call(items, None)
+        self.assertEquals(result, [items[2]])
+        self.assertTrue(mock_util().checklist.called)
+        self.assertEquals(mock_util().checklist.call_args[0][0], None)
+
+    @test_util.patch_get_utility("certbot.display.ops.z_util")
+    def test_choose_names_success_question(self, mock_util):
+        items = ["first", "second", "third"]
+        question = "Which one?"
+        mock_util().checklist.return_value = (display_util.OK, [items[1]])
+        result = self._call(items, question)
+        self.assertEquals(result, [items[1]])
+        self.assertTrue(mock_util().checklist.called)
+        self.assertEquals(mock_util().checklist.call_args[0][0], question)
+
+    @test_util.patch_get_utility("certbot.display.ops.z_util")
+    def test_choose_names_user_cancel(self, mock_util):
+        items = ["first", "second", "third"]
+        question = "Want to cancel?"
+        mock_util().checklist.return_value = (display_util.CANCEL, [])
+        result = self._call(items, question)
+        self.assertEquals(result, [])
+        self.assertTrue(mock_util().checklist.called)
+        self.assertEquals(mock_util().checklist.call_args[0][0], question)
+
+
 if __name__ == "__main__":
     unittest.main()  # pragma: no cover

certbot/tests/main_test.py

@@ -1533,5 +1533,114 @@ class MakeOrVerifyNeededDirs(test_util.ConfigTestCase):
                 strict=self.config.strict_permissions)
 
 
+class EnhanceTest(unittest.TestCase):
+    """Tests for certbot.main.enhance."""
+
+    def setUp(self):
+        self.get_utility_patch = test_util.patch_get_utility()
+        self.mock_get_utility = self.get_utility_patch.start()
+
+    def tearDown(self):
+        self.get_utility_patch.stop()
+
+    def _call(self, args):
+        plugins = disco.PluginsRegistry.find_all()
+        config = configuration.NamespaceConfig(
+            cli.prepare_and_parse_args(plugins, args))
+
+        with mock.patch('certbot.cert_manager.get_certnames') as mock_certs:
+            mock_certs.return_value = ['example.com']
+            with mock.patch('certbot.cert_manager.domains_for_certname') as mock_dom:
+                mock_dom.return_value = ['example.com']
+                with mock.patch('certbot.main._init_le_client') as mock_init:
+                    mock_client = mock.MagicMock()
+                    mock_client.config = config
+                    mock_init.return_value = mock_client
+                    main.enhance(config, plugins)
+                    return mock_client # returns the client
+
+    @mock.patch('certbot.main.plug_sel.record_chosen_plugins')
+    @mock.patch('certbot.cert_manager.lineage_for_certname')
+    @mock.patch('certbot.main.display_ops.choose_values')
+    @mock.patch('certbot.main._find_domains_or_certname')
+    def test_selection_question(self, mock_find, mock_choose, mock_lineage, _rec):
+        mock_lineage.return_value = mock.MagicMock(chain_path="/tmp/nonexistent")
+        mock_choose.return_value = ['example.com']
+        mock_find.return_value = (None, None)
+        with mock.patch('certbot.main.plug_sel.pick_installer') as mock_pick:
+            self._call(['enhance', '--redirect'])
+            self.assertTrue(mock_pick.called)
+            # Check that the message includes "enhancements"
+            self.assertTrue("enhancements" in mock_pick.call_args[0][3])
+
+    @mock.patch('certbot.main.plug_sel.record_chosen_plugins')
+    @mock.patch('certbot.cert_manager.lineage_for_certname')
+    @mock.patch('certbot.main.display_ops.choose_values')
+    @mock.patch('certbot.main._find_domains_or_certname')
+    def test_selection_auth_warning(self, mock_find, mock_choose, mock_lineage, _rec):
+        mock_lineage.return_value = mock.MagicMock(chain_path="/tmp/nonexistent")
+        mock_choose.return_value = ["example.com"]
+        mock_find.return_value = (None, None)
+        with mock.patch('certbot.main.plug_sel.pick_installer'):
+            with mock.patch('certbot.main.plug_sel.logger.warning') as mock_log:
+                mock_client = self._call(['enhance', '-a', 'webroot', '--redirect'])
+                self.assertTrue(mock_log.called)
+                self.assertTrue("make sense" in mock_log.call_args[0][0])
+                self.assertTrue(mock_client.enhance_config.called)
+
+    @mock.patch('certbot.cert_manager.lineage_for_certname')
+    @mock.patch('certbot.main.display_ops.choose_values')
+    @mock.patch('certbot.main.plug_sel.record_chosen_plugins')
+    def test_enhance_config_call(self, _rec, mock_choose, mock_lineage):
+        mock_lineage.return_value = mock.MagicMock(chain_path="/tmp/nonexistent")
+        mock_choose.return_value = ["example.com"]
+        with mock.patch('certbot.main.plug_sel.pick_installer'):
+            mock_client = self._call(['enhance', '--redirect', '--hsts'])
+            req_enh = ["redirect", "hsts"]
+            not_req_enh = ["uir"]
+            self.assertTrue(mock_client.enhance_config.called)
+            self.assertTrue(
+                all([getattr(mock_client.config, e) for e in req_enh]))
+            self.assertFalse(
+                any([getattr(mock_client.config, e) for e in not_req_enh]))
+            self.assertTrue(
+                "example.com" in mock_client.enhance_config.call_args[0][0])
+
+    @mock.patch('certbot.cert_manager.lineage_for_certname')
+    @mock.patch('certbot.main.display_ops.choose_values')
+    @mock.patch('certbot.main.plug_sel.record_chosen_plugins')
+    def test_enhance_noninteractive(self, _rec, mock_choose, mock_lineage):
+        mock_lineage.return_value = mock.MagicMock(
+            chain_path="/tmp/nonexistent")
+        mock_choose.return_value = ["example.com"]
+        with mock.patch('certbot.main.plug_sel.pick_installer'):
+            mock_client = self._call(['enhance', '--redirect',
+                                      '--hsts', '--non-interactive'])
+            self.assertTrue(mock_client.enhance_config.called)
+            self.assertFalse(mock_choose.called)
+
+    @mock.patch('certbot.main.display_ops.choose_values')
+    @mock.patch('certbot.main.plug_sel.record_chosen_plugins')
+    def test_user_abort_domains(self, _rec, mock_choose):
+        mock_choose.return_value = []
+        with mock.patch('certbot.main.plug_sel.pick_installer'):
+            self.assertRaises(errors.Error,
+                              self._call,
+                              ['enhance', '--redirect', '--hsts'])
+
+    def test_no_enhancements_defined(self):
+        self.assertRaises(errors.MisconfigurationError,
+                          self._call, ['enhance'])
+
+    @mock.patch('certbot.main.plug_sel.choose_configurator_plugins')
+    @mock.patch('certbot.main.display_ops.choose_values')
+    @mock.patch('certbot.main.plug_sel.record_chosen_plugins')
+    def test_plugin_selection_error(self, _rec, mock_choose, mock_pick):
+        mock_choose.return_value = ["example.com"]
+        mock_pick.return_value = (None, None)
+        mock_pick.side_effect = errors.PluginSelectionError()
+        mock_client = self._call(['enhance', '--hsts'])
+        self.assertFalse(mock_client.enhance_config.called)
+
 if __name__ == '__main__':
     unittest.main()  # pragma: no cover