Repository: google/turbinia
Pull request #461 (2019-11-21 15:14:38)

Task Status Update
Added a feature for Turbiniactl status <request ID> to show live updates for queued/running tasks.

alimez: @aarontp PTAL.

Patch:
diff --git a/out.txt b/out.txt
new file mode 100644
index 0000000..e69de29
diff --git a/turbinia/client.py b/turbinia/client.py
index cde7c3a..cc298dc 100644
--- a/turbinia/client.py
+++ b/turbinia/client.py
@@ -95,6 +95,24 @@ log = logging.getLogger('turbinia')
 logger.setup()
 
 
+def get_turbinia_client(run_local=False):
+  """Return Turbinia client based on config.
+
+  Returns:
+    Initialized BaseTurbiniaClient or TurbiniaCeleryClient object.
+  """
+  config.LoadConfig()
+  # pylint: disable=no-else-return
+  if config.TASK_MANAGER.lower() == 'psq':
+    return BaseTurbiniaClient(run_local=run_local)
+  elif config.TASK_MANAGER.lower() == 'celery':
+    return TurbiniaCeleryClient(run_local=run_local)
+  else:
+    msg = 'Task Manager type "{0:s}" not implemented'.format(
+        config.TASK_MANAGER)
+    raise TurbiniaException(msg)
+
+
 def check_dependencies(dependencies):
   """Checks system dependencies.
 
@@ -233,7 +251,7 @@ class TurbiniaStats(object):
         self.description, self.count, self.min, self.mean, self.max)
 
 
-class TurbiniaClient(object):
+class BaseTurbiniaClient(object):
   """Client class for Turbinia.
 
   Attributes:
@@ -748,7 +766,7 @@ class TurbiniaClient(object):
     return 'Closed Task IDs: %s' % response.get('result')
 
 
-class TurbiniaCeleryClient(TurbiniaClient):
+class TurbiniaCeleryClient(BaseTurbiniaClient):
   """Client class for Turbinia (Celery).
 
   Overriding some things specific to Celery operation.
@@ -757,8 +775,8 @@ class TurbiniaCeleryClient(TurbiniaClient):
     redis (RedisStateManager): Redis datastore object
   """
 
-  def __init__(self, *_, **__):
-    super(TurbiniaCeleryClient, self).__init__()
+  def __init__(self, *args, **kwargs):
+    super(TurbiniaCeleryClient, self).__init__(*args, **kwargs)
     self.redis = RedisStateManager()
 
   def send_request(self, request):
@@ -818,7 +836,7 @@ class TurbiniaServer(object):
     self.task_manager.add_evidence(evidence_)
 
 
-class TurbiniaCeleryWorker(TurbiniaClient):
+class TurbiniaCeleryWorker(BaseTurbiniaClient):
   """Turbinia Celery Worker class.
 
   Attributes:
diff --git a/turbinia/state_manager.py b/turbinia/state_manager.py
index c85dd98..8906673 100644
--- a/turbinia/state_manager.py
+++ b/turbinia/state_manager.py
@@ -31,8 +31,6 @@ import six
 from turbinia import config
 from turbinia.config import DATETIME_FORMAT
 from turbinia import TurbiniaException
-from turbinia.workers import TurbiniaTask
-from turbinia.workers import TurbiniaTaskResult
 
 config.LoadConfig()
 if config.STATE_MANAGER.lower() == 'datastore':
@@ -112,6 +110,9 @@ class BaseStateManager(object):
     if task_dict.get('run_time'):
       task_dict['run_time'] = task_dict['run_time'].total_seconds()
 
+    #Importing these here to avoid circular dependencies.
+    from turbinia.workers import TurbiniaTask
+    from turbinia.workers import TurbiniaTaskResult
     # Set all non-existent keys to None
     all_attrs = set(
         TurbiniaTask.STORED_ATTRIBUTES + TurbiniaTaskResult.STORED_ATTRIBUTES)
diff --git a/turbinia/turbiniactl.py b/turbinia/turbiniactl.py
index 09e334b..8ba90d1 100644
--- a/turbinia/turbiniactl.py
+++ b/turbinia/turbiniactl.py
@@ -404,7 +404,7 @@ def main():
   # the config until after we parse the args so that we can use those arguments
   # to point to config paths.
   from turbinia import notify
-  from turbinia.client import TurbiniaClient
+  from turbinia import client as TurbiniaClientProvider
   from turbinia.client import TurbiniaCeleryClient
   from turbinia.client import TurbiniaServer
   from turbinia.client import TurbiniaCeleryWorker
@@ -454,15 +454,9 @@ def main():
           'Cannot open file {0:s} [{1!s}]'.format(args.filter_patterns_file, e))
 
   # Create Client object
+  client = None
   if args.command not in ('psqworker', 'server'):
-    if config.TASK_MANAGER.lower() == 'celery':
-      client = TurbiniaCeleryClient()
-    elif args.run_local:
-      client = TurbiniaClient(run_local=True)
-    else:
-      client = TurbiniaClient()
-  else:
-    client = None
+    client = TurbiniaClientProvider.get_turbinia_client(args.run_local)
 
   # Make sure run_local flags aren't conflicting with other server/client flags
   server_flags_set = args.server or args.command == 'server'
diff --git a/turbinia/workers/__init__.py b/turbinia/workers/__init__.py
index 62e4d0d..e87e5fc 100644
--- a/turbinia/workers/__init__.py
+++ b/turbinia/workers/__init__.py
@@ -37,6 +37,7 @@ from turbinia import config
 from turbinia.config import DATETIME_FORMAT
 from turbinia.evidence import evidence_decode
 from turbinia import output_manager
+from turbinia import state_manager
 from turbinia import TurbiniaException
 
 log = logging.getLogger('turbinia')
@@ -78,6 +79,8 @@ class TurbiniaTaskResult(object):
       task_id: Task ID of the parent task.
       task_name: Name of parent task.
       requester: The user who requested the task.
+      state_manager: (DatastoreStateManager|RedisStateManager): State manager
+        object to handle syncing with storage.
       worker_name: Name of worker task executed on.
       _log: A list of log messages
   """
@@ -115,6 +118,7 @@ class TurbiniaTaskResult(object):
     self.status = None
     self.error = {}
     self.worker_name = platform.node()
+    self.state_manager = None
     # TODO(aarontp): Create mechanism to grab actual python logging data.
     self._log = []
 
@@ -134,6 +138,7 @@ class TurbiniaTaskResult(object):
     self.task_id = task.id
     self.task_name = task.name
     self.requester = task.requester
+    self.state_manager = state_manager.get_state_manager()
     if task.output_manager.is_setup:
       _, self.output_dir = task.output_manager.get_local_output_dirs()
     else:
@@ -192,7 +197,7 @@ class TurbiniaTaskResult(object):
     # also fail.
     # pylint: disable=broad-except
     except Exception as exception:
-      message = 'Evidence post-processing for {0:s} failed: {1!s}'.format(
+      message = 'Evidence post-processing for {0!s} failed: {1!s}'.format(
           self.input_evidence.name, exception)
       self.log(message, level=logging.ERROR)
 
@@ -237,6 +242,18 @@ class TurbiniaTaskResult(object):
     if traceback_:
       self.result.set_error(message, traceback_)
 
+  def update_task_status(self, task, status):
+    """Updates the task status and pushes it directly to datastor.
+
+    Args:
+      task: The calling Task object
+      status: One line descriptive task status.
+    """
+    task.result.status = 'Task {0!s} is {1!s} on {2!s}'.format(
+        self.task_name, status, self.worker_name)
+
+    self.state_manager.update_task(task)
+
   def add_evidence(self, evidence, evidence_config):
     """Populate the results list.
 
@@ -273,11 +290,13 @@ class TurbiniaTaskResult(object):
     Returns:
       dict: Object dictionary that is JSON serializable.
     """
+    self.state_manager = None
     self.run_time = self.run_time.total_seconds() if self.run_time else None
     self.start_time = self.start_time.strftime(DATETIME_FORMAT)
     if self.input_evidence:
       self.input_evidence = self.input_evidence.serialize()
     self.evidence = [x.serialize() for x in self.evidence]
+
     return self.__dict__
 
   @classmethod
@@ -292,6 +311,8 @@ class TurbiniaTaskResult(object):
     """
     result = TurbiniaTaskResult()
     result.__dict__.update(input_dict)
+    if result.state_manager:
+      result.state_manager = None
     if result.run_time:
       result.run_time = timedelta(seconds=result.run_time)
     result.start_time = datetime.strptime(result.start_time, DATETIME_FORMAT)
@@ -522,7 +543,6 @@ class TurbiniaTask(object):
       raise TurbiniaException(
           'Evidence source path {0:s} does not exist'.format(
               evidence.source_path))
-    evidence.preprocess(self.tmp_dir)
     return self.result
 
   def touch(self):
@@ -609,14 +629,40 @@ class TurbiniaTask(object):
 
     log.debug('Task {0:s} {1:s} awaiting execution'.format(self.name, self.id))
     evidence = evidence_decode(evidence)
+    try:
+      self.result = self.setup(evidence)
+      self.result.update_task_status(self, 'queued')
+    except TurbiniaException as exception:
+      message = (
+          '{0:s} Task setup failed with exception: [{1!s}]'.format(
+              self.name, exception))
+      # Logging explicitly here because the result is in an unknown state
+      trace = traceback.format_exc()
+      log.error(message)
+      log.error(trace)
+      if self.result:
+        if hasattr(exception, 'message'):
+          self.result.set_error(exception.message, traceback.format_exc())
+        else:
+          self.result.set_error(exception.__class__, traceback.format_exc())
+        self.result.status = message
+      else:
+        self.result = TurbiniaTaskResult(
+            base_output_dir=self.base_output_dir, request_id=self.request_id,
+            job_id=self.job_id)
+        self.result.setup(self)
+        self.result.status = message
+        self.result.set_error(message, traceback.format_exc())
+      return self.result.serialize()
     with filelock.FileLock(config.LOCK_FILE):
       log.info('Starting Task {0:s} {1:s}'.format(self.name, self.id))
       original_result_id = None
       try:
-        self.result = self.setup(evidence)
         original_result_id = self.result.id
         evidence.validate()
 
+        # Preprocessor must be called after evidence validation.
+        evidence.preprocess(self.tmp_dir)
         # TODO(wyassine): refactor it so the result task does not
         # have to go through the preprocess stage. At the moment
         # self.results.setup is required to be called to set its status.
@@ -637,7 +683,7 @@ class TurbiniaTask(object):
           self.result.log(message, level=logging.ERROR)
           self.result.status = message
           return self.result.serialize()
-
+        self.result.update_task_status(self, 'running')
         self._evidence_config = evidence.config
         self.result = self.run(evidence, self.result)
       # pylint: disable=broad-except
diff --git a/turbinia/workers/finalize_request.py b/turbinia/workers/finalize_request.py
index bb90463..2750edc 100644
--- a/turbinia/workers/finalize_request.py
+++ b/turbinia/workers/finalize_request.py
@@ -40,8 +40,8 @@ class FinalizeRequestTask(TurbiniaTask):
         TurbiniaTaskResult: Task execution results.
     """
     # Doing a delayed import to avoid circular dependencies.
-    from turbinia.client import TurbiniaClient
-    client = TurbiniaClient()
+    from turbinia import client as TurbiniaClientProvider
+    client = TurbiniaClientProvider.get_turbinia_client()
 
     report_file = os.path.join(
         self.tmp_dir, 'final_turbinia_report_{0:s}.md'.format(self.id))


Test patch:
diff --git a/turbinia/client_test.py b/turbinia/client_test.py
index b27626f..ebd2b2d 100644
--- a/turbinia/client_test.py
+++ b/turbinia/client_test.py
@@ -28,7 +28,7 @@ import textwrap
 import mock
 
 from turbinia import config
-from turbinia.client import TurbiniaClient
+from turbinia import client as TurbiniaClientProvider
 from turbinia.client import TurbiniaServer
 from turbinia.client import TurbiniaStats
 from turbinia.client import TurbiniaPsqWorker
@@ -205,7 +205,7 @@ class TestTurbiniaClient(unittest.TestCase):
   def testTurbiniaClientInit(self, _, __):
     """Basic test for client."""
     config.LoadConfig()
-    client = TurbiniaClient()
+    client = TurbiniaClientProvider.get_turbinia_client()
     self.assertTrue(hasattr(client, 'task_manager'))
 
   @mock.patch('turbinia.client.GoogleCloudFunction.ExecuteFunction')
@@ -221,7 +221,7 @@ class TestTurbiniaClient(unittest.TestCase):
     gcf_result = json.dumps(gcf_result)
     function_return = {'result': gcf_result}
     mock_cloud_function.return_value = function_return
-    client = TurbiniaClient()
+    client = TurbiniaClientProvider.get_turbinia_client()
     task_data = client.get_task_data('inst', 'proj', 'reg')
     # get_task_data() converts this back into a timedelta(). We returned it
     # seconds from the GCF function call because that is what it is stored in
@@ -235,7 +235,7 @@ class TestTurbiniaClient(unittest.TestCase):
   def testTurbiniaClientGetTaskDataNoResults(self, _, __, mock_cloud_function):
     """Test for exception after empty results from cloud functions."""
     mock_cloud_function.return_value = {}
-    client = TurbiniaClient()
+    client = TurbiniaClientProvider.get_turbinia_client()
     self.assertRaises(
         TurbiniaException, client.get_task_data, "inst", "proj", "reg")
 
@@ -246,7 +246,7 @@ class TestTurbiniaClient(unittest.TestCase):
       self, _, __, mock_cloud_function):
     """Test for exception after bad json results from cloud functions."""
     mock_cloud_function.return_value = {'result': None}
-    client = TurbiniaClient()
+    client = TurbiniaClientProvider.get_turbinia_client()
     self.assertRaises(
         TurbiniaException, client.get_task_data, "inst", "proj", "reg")
 
@@ -255,7 +255,7 @@ class TestTurbiniaClient(unittest.TestCase):
   @mock.patch('turbinia.state_manager.get_state_manager')
   def testClientFormatTaskStatistics(self, _, __, ___):
     """Tests format_task_statistics() report output."""
-    client = TurbiniaClient()
+    client = TurbiniaClientProvider.get_turbinia_client()
     client.get_task_data = mock.MagicMock()
     client.get_task_data.return_value = self.task_data
     stats_report = client.format_task_statistics('inst', 'proj', 'reg')
@@ -267,7 +267,7 @@ class TestTurbiniaClient(unittest.TestCase):
   @mock.patch('turbinia.state_manager.get_state_manager')
   def testClientFormatTaskStatisticsCsv(self, _, __, ___):
     """Tests format_task_statistics() CSV report output."""
-    client = TurbiniaClient()
+    client = TurbiniaClientProvider.get_turbinia_client()
     client.get_task_data = mock.MagicMock()
     client.get_task_data.return_value = self.task_data
     stats_report = client.format_task_statistics(
@@ -280,7 +280,7 @@ class TestTurbiniaClient(unittest.TestCase):
   @mock.patch('turbinia.state_manager.get_state_manager')
   def testClientGetTaskStatistics(self, _, __, ___):
     """Tests get_task_statistics() basic functionality."""
-    client = TurbiniaClient()
+    client = TurbiniaClientProvider.get_turbinia_client()
     client.get_task_data = mock.MagicMock()
     client.get_task_data.return_value = self.task_data
     task_stats = client.get_task_statistics('inst', 'proj', 'reg')
@@ -313,7 +313,7 @@ class TestTurbiniaClient(unittest.TestCase):
   @mock.patch('turbinia.state_manager.get_state_manager')
   def testClientFormatTaskStatus(self, _, __, ___):
     """Tests format_task_status() with empty report_priority."""
-    client = TurbiniaClient()
+    client = TurbiniaClientProvider.get_turbinia_client()
     client.get_task_data = mock.MagicMock()
     self.task_data[0]['report_priority'] = None
     self.task_data[1]['report_priority'] = ''
@@ -327,7 +327,7 @@ class TestTurbiniaClient(unittest.TestCase):
   @mock.patch('turbinia.state_manager.get_state_manager')
   def testClientFormatTaskStatusShortReport(self, _, __, ___):
     """Tests format_task_status() has valid output with short report."""
-    client = TurbiniaClient()
+    client = TurbiniaClientProvider.get_turbinia_client()
     client.get_task_data = mock.MagicMock()
     client.get_task_data.return_value = self.task_data
     result = client.format_task_status('inst', 'proj', 'reg')
@@ -338,7 +338,7 @@ class TestTurbiniaClient(unittest.TestCase):
   @mock.patch('turbinia.state_manager.get_state_manager')
   def testClientFormatTaskStatusFullReport(self, _, __, ___):
     """Tests format_task_status() has valid output with full report."""
-    client = TurbiniaClient()
+    client = TurbiniaClientProvider.get_turbinia_client()
     client.get_task_data = mock.MagicMock()
     client.get_task_data.return_value = self.task_data
     result = client.format_task_status('inst', 'proj', 'reg', full_report=True)
@@ -349,7 +349,7 @@ class TestTurbiniaClient(unittest.TestCase):
   @mock.patch('turbinia.state_manager.get_state_manager')
   def testClientFormatTaskStatusFiles(self, _, __, ___):
     """Tests format_task_status() has valid output with report and files."""
-    client = TurbiniaClient()
+    client = TurbiniaClientProvider.get_turbinia_client()
     client.get_task_data = mock.MagicMock()
     client.get_task_data.return_value = self.task_data
     result = client.format_task_status(
diff --git a/turbinia/task_manager_test.py b/turbinia/task_manager_test.py
index 456c55f..4807de4 100644
--- a/turbinia/task_manager_test.py
+++ b/turbinia/task_manager_test.py
@@ -181,7 +181,8 @@ class TestTaskManager(TestTurbiniaTaskBase):
     self.assertIsNone(self.manager.process_result(self.result))
     self.assertIsInstance(self.result.evidence, list)
 
-  def testFinalizeJobGenerateJobFinalizeTasks(self):
+  @mock.patch('turbinia.task_manager.state_manager.get_state_manager')
+  def testFinalizeJobGenerateJobFinalizeTasks(self, _):
     """Tests process_job method generates Job finalize Task."""
     request_id = 'testRequestID'
     self.task.id = 'createdFinalizeTask'
diff --git a/turbinia/workers/workers_test.py b/turbinia/workers/workers_test.py
index 53551d6..0a4a9ad 100644
--- a/turbinia/workers/workers_test.py
+++ b/turbinia/workers/workers_test.py
@@ -27,6 +27,7 @@ from turbinia import TurbiniaException
 from turbinia.workers import TurbiniaTask
 from turbinia.workers import TurbiniaTaskResult
 from turbinia.workers.plaso import PlasoTask
+from turbinia import state_manager
 
 
 class TestTurbiniaTaskBase(unittest.TestCase):
@@ -63,7 +64,7 @@ class TestTurbiniaTaskBase(unittest.TestCase):
     test_disk_path = tempfile.mkstemp(dir=self.base_output_dir)[1]
     self.remove_files.append(test_disk_path)
     self.evidence = evidence.RawDisk(source_path=test_disk_path)
-
+    self.evidence.preprocess = mock.MagicMock()
     # Set up TurbiniaTaskResult
     self.result = TurbiniaTaskResult(base_output_dir=self.base_output_dir)
     self.result.setup(self.task)
@@ -99,8 +100,11 @@ class TestTurbiniaTaskBase(unittest.TestCase):
       validate_result = self.result
 
     self.result.status = 'TestStatus'
+    self.result.update_task_status = mock.MagicMock()
     self.result.close = mock.MagicMock()
     self.task.setup = mock.MagicMock(return_value=setup)
+    self.result.worker_name = 'worker1'
+    self.result.state_manager = None
     if mock_run:
       self.task.run = mock.MagicMock(return_value=run)
     self.task.validate_result = mock.MagicMock(return_value=validate_result)
@@ -124,6 +128,7 @@ class TestTurbiniaTask(TestTurbiniaTaskBase):
     self.setResults()
     self.result.closed = True
     new_result = self.task.run_wrapper(self.evidence.__dict__)
+
     new_result = TurbiniaTaskResult.deserialize(new_result)
     self.assertEqual(new_result.status, 'TestStatus')
     self.result.close.assert_not_called()
@@ -143,10 +148,8 @@ class TestTurbiniaTask(TestTurbiniaTaskBase):
     checked_result.setup(self.task)
     checked_result.status = 'CheckedResult'
     self.setResults(run=bad_result, validate_result=checked_result)
-
     new_result = self.task.run_wrapper(self.evidence.__dict__)
     new_result = TurbiniaTaskResult.deserialize(new_result)
-
     self.task.validate_result.assert_any_call(bad_result)
     self.assertEqual(type(new_result), TurbiniaTaskResult)
     self.assertIn('CheckedResult', new_result.status)
@@ -175,10 +178,9 @@ class TestTurbiniaTask(TestTurbiniaTaskBase):
   def testTurbiniaTaskRunWrapperSetupFail(self):
     """Test that the run wrapper recovers from setup failing."""
     self.task.result = None
-    canary_status = 'ReturnedFromValidateResult'
-    self.result.status = canary_status
-    self.task.validate_result = mock.MagicMock(return_value=self.result)
-    self.task.setup = mock.MagicMock(side_effect=TurbiniaException)
+    canary_status = 'exception_message'
+    self.task.setup = mock.MagicMock(
+        side_effect=TurbiniaException('exception_message'))
     self.remove_files.append(
         os.path.join(self.task.base_output_dir, 'worker-log.txt'))
 
@@ -190,12 +192,14 @@ class TestTurbiniaTask(TestTurbiniaTaskBase):
   def testTurbiniaTaskValidateResultGoodResult(self):
     """Tests validate_result with good result."""
     self.result.status = 'GoodStatus'
+    self.result.state_manager = None
     new_result = self.task.validate_result(self.result)
     self.assertEqual(new_result.status, 'GoodStatus')
     self.assertDictEqual(new_result.error, {})
 
   @mock.patch('turbinia.workers.TurbiniaTaskResult.close')
-  def testTurbiniaTaskValidateResultBadResult(self, _):
+  @mock.patch('turbinia.state_manager.get_state_manager')
+  def testTurbiniaTaskValidateResultBadResult(self, _, __):
     """Tests validate_result with bad result."""
     # Passing in an unpickleable object (json module) and getting back a
     # TurbiniaTaskResult