Better critical error handling (log2timeline#397)

* Raise CriticalError instead of exit-ing

* Soothe linter

* Catch exception around SetupModule and RunModules

* Restore message in CriticalError

dftimewolf/cli/dftimewolf_recipes.py

@@ -332,25 +332,33 @@ def Main():
 
  try:
  tool.ReadRecipes()
- except (KeyError, errors.RecipeParseError) as exception:
+ except (KeyError, errors.RecipeParseError, errors.CriticalError) as exception:
  logger.critical(str(exception))
  return False
 
  try:
  tool.ParseArguments(sys.argv[1:])
- except (errors.CommandLineParseError, errors.RecipeParseError) as exception:
+ except (errors.CommandLineParseError,
+ errors.RecipeParseError,
+ errors.CriticalError) as exception:
  logger.critical(str(exception))
  return False
 
  tool.state.LogExecutionPlan()
 
  tool.RunPreflights()
 
- # TODO: log errors if this fails.
- tool.SetupModules()
+ try:
+ tool.SetupModules()
+ except errors.CriticalError as exception:
+ logger.critical(str(exception))
+ return False
 
- # TODO: log errors if this fails.
- tool.RunModules()
+ try:
+ tool.RunModules()
+ except errors.CriticalError as exception:
+ logger.critical(str(exception))
+ return False
 
  tool.CleanUpPreflights()
 

dftimewolf/lib/errors.py

@@ -33,3 +33,6 @@ class RecipeParseError(DFTimewolfError):
 
 class CommandLineParseError(DFTimewolfError):
  """Error when parsing the command-line arguments."""
+
+class CriticalError(DFTimewolfError):
+ """Critical error that should abort the whole workflow."""

dftimewolf/lib/state.py

@@ -5,7 +5,6 @@
 """
 
 import logging
-import sys
 import threading
 import traceback
 import importlib
@@ -421,5 +420,4 @@ def CheckErrors(self, is_global=False):
  'Please consider opening an issue: {0:s}'.format(NEW_ISSUE_URL))
 
  if critical_errors:
- logger.critical('Critical error found. Aborting.')
- sys.exit(1)
+ raise errors.CriticalError('Critical error found. Aborting.')

tests/lib/state.py

@@ -12,6 +12,7 @@
 from dftimewolf.lib.containers import containers
 from dftimewolf.lib.modules import manager as modules_manager
 from dftimewolf.lib.recipes import manager as recipes_manager
+from dftimewolf.lib import errors
 
 from tests.test_modules import modules
 from tests.test_modules import test_recipe
@@ -183,20 +184,19 @@ def testProcessNamedModules(self, mock_process1, mock_process2):
 
  @mock.patch('tests.test_modules.modules.DummyModule2.Process')
  @mock.patch('tests.test_modules.modules.DummyModule1.Process')
- @mock.patch('sys.exit')
- def testProcessErrors(self, mock_exit, mock_process1, mock_process2):
+ def testProcessErrors(self, mock_process1, mock_process2):
  """Tests that module's errors are correctly caught."""
  test_state = state.DFTimewolfState(config.Config)
  test_state.command_line_options = {}
  test_state.LoadRecipe(test_recipe.contents, TEST_MODULES)
  mock_process1.side_effect = Exception('asd')
  test_state.SetupModules()
- test_state.RunModules()
+ with self.assertRaises(errors.CriticalError):
+ test_state.RunModules()
  mock_process1.assert_called_with()
- # Procesds() in module 2 is never called since the failure in Module1
+ # Process() in module 2 is never called since the failure in Module1
  # will abort execution
  mock_process2.assert_not_called()
- mock_exit.assert_called_with(1)
  self.assertEqual(len(test_state.global_errors), 1)
  error = test_state.global_errors[0]
  self.assertIn('An unknown error occurred in module DummyModule1: asd',