From 85b5bc0cb2cff93ee129f64e2a3a9ed1f279d16d Mon Sep 17 00:00:00 2001
From: Brad Warren <bradmw@umich.edu>
Date: Mon, 29 Jun 2015 17:31:48 -0700
Subject: [PATCH] Reimplemented exception handling

---
 letsencrypt/cli.py            | 63 ++++++++++++++++++-----------------
 letsencrypt/tests/cli_test.py | 22 ++++++++++--
 2 files changed, 52 insertions(+), 33 deletions(-)

diff --git a/letsencrypt/cli.py b/letsencrypt/cli.py
index 9bb4438a5..8f833522d 100644
--- a/letsencrypt/cli.py
+++ b/letsencrypt/cli.py
@@ -2,11 +2,13 @@
 # TODO: Sanity check all input.  Be sure to avoid shell code etc...
 import argparse
 import atexit
+import functools
 import logging
 import logging.handlers
 import os
 import sys
 import time
+import traceback
 
 import configargparse
 import zope.component
@@ -642,15 +644,29 @@ def _setup_logging(args):
     logger.info("Saving debug log to %s", log_file_name)
 
 
-def main2(cli_args, args, config, plugins):
-    """Continued main script execution."""
+def _handle_exception(exc_type, exc_value, trace, args):
+    logger.debug(
+        "Exiting abnormally:\n%s",
+        "".join(traceback.format_exception(exc_type, exc_value, trace)))
 
-    # Displayer
-    if args.text_mode:
-        displayer = display_util.FileDisplay(sys.stdout)
+    if issubclass(exc_type, errors.Error) and (not args or not args.debug):
+        sys.exit(exc_value)
+    elif issubclass(exc_type, Exception) and args and not args.debug:
+        sys.exit(
+            "An unexpected error occurred. Please see the logfiles in {0} for "
+            "more details.".format(args.logs_dir))
     else:
-        displayer = display_util.NcursesDisplay()
-    zope.component.provideUtility(displayer)
+        traceback.print_exception(exc_type, exc_value, trace, file=sys.stderr)
+
+
+def main(cli_args=sys.argv[1:]):
+    """Command line argument parsing and main script execution."""
+    sys.excepthook = functools.partial(_handle_exception, args=None)
+
+    # note: arg parser internally handles --help (and exits afterwards)
+    plugins = plugins_disco.PluginsRegistry.find_all()
+    args = create_parser(plugins, cli_args).parse_args(cli_args)
+    config = configuration.NamespaceConfig(args)
 
     # Setup logging ASAP, otherwise "No handlers could be found for
     # logger ..." TODO: this should be done before plugins discovery
@@ -666,6 +682,15 @@ def main2(cli_args, args, config, plugins):
     logger.debug("Arguments: %r", cli_args)
     logger.debug("Discovered plugins: %r", plugins)
 
+    sys.excepthook = functools.partial(_handle_exception, args=args)
+
+    # Displayer
+    if args.text_mode:
+        displayer = display_util.FileDisplay(sys.stdout)
+    else:
+        displayer = display_util.NcursesDisplay()
+    zope.component.provideUtility(displayer)
+
     # Reporter
     report = reporter.Reporter()
     zope.component.provideUtility(report)
@@ -685,29 +710,5 @@ def main2(cli_args, args, config, plugins):
     return args.func(args, config, plugins)
 
 
-def main(cli_args=sys.argv[1:]):
-    """Command line argument parsing and main script execution."""
-    # note: arg parser internally handles --help (and exits afterwards)
-    plugins = plugins_disco.PluginsRegistry.find_all()
-    args = create_parser(plugins, cli_args).parse_args(cli_args)
-    config = configuration.NamespaceConfig(args)
-
-    def handle_exception_common():
-        """Logs the exception and reraises it if in debug mode."""
-        logger.debug("Exiting abnormally", exc_info=True)
-        if args.debug:
-            raise
-
-    try:
-        return main2(cli_args, args, config, plugins)
-    except errors.Error as error:
-        handle_exception_common()
-        return error
-    except Exception: # pylint: disable=broad-except
-        handle_exception_common()
-        return ("An unexpected error occured. Please see the logfiles in {0} "
-                "for more details.".format(args.logs_dir))
-
-
 if __name__ == "__main__":
     sys.exit(main())  # pragma: no cover
diff --git a/letsencrypt/tests/cli_test.py b/letsencrypt/tests/cli_test.py
index 798ec227a..8da009864 100644
--- a/letsencrypt/tests/cli_test.py
+++ b/letsencrypt/tests/cli_test.py
@@ -66,12 +66,30 @@ class CLITest(unittest.TestCase):
         attrs = {'view_config_changes.side_effect' : error}
         self.assertRaises(
             errors.Error, self._call, ['--debug'] + cmd_arg, attrs)
-        self._call(cmd_arg, attrs)
 
         attrs['view_config_changes.side_effect'] = [ValueError]
         self.assertRaises(
             ValueError, self._call, ['--debug'] + cmd_arg, attrs)
-        self._call(cmd_arg, attrs)
+
+    @mock.patch("letsencrypt.cli.sys")
+    def test_handle_exception(self, mock_sys):
+        # pylint: disable=protected-access
+        import StringIO
+        from letsencrypt import cli
+        from letsencrypt import errors
+
+        cli._handle_exception(errors.Error, "detail", None, None)
+        mock_sys.exit.assert_called_once_with("detail")
+
+        args = mock.MagicMock(debug=False)
+        cli._handle_exception(ValueError, "detail", None, args)
+        self.assertTrue("logfile" in mock_sys.exit.call_args_list[1][0][0])
+
+        mock_sys.stderr = StringIO.StringIO()
+        exc_value = "A very specific string"
+        cli._handle_exception(KeyboardInterrupt, exc_value, None, None)
+        self.assertTrue(exc_value in mock_sys.stderr.getvalue())
+
 
 if __name__ == '__main__':
     unittest.main()  # pragma: no cover