Repository: google/turbinia
Pull request #308 (2018-11-20 06:08:23)

[Celery] Use JSON serialization instead of Pickle
* No more RCE worry; pickle is gone & JSON is in!
* Removes the awkward `fexec` method, we now just call `task_runner` like normal
* `task_runner` takes a serializable dict as input, which is then decoded into the proper task object
* Adds a `task_decode` method in the `client` module, this behaves similarly to `evidence_decode` (placed it here because the client module already imports every existing *Task module)
* `TurbiniaTask` objects now return dicts that are JSON serializable; those dicts are later decoded back into `TurbiniaTaskResult` objects once received by the server

I'm not sure how PSQ handles serialization (I think it converts to JSON automatically, right?), but everything here should work with either backend.

ericzinnikas: > Global variable 'task_runner' undefined at the module level

Yeah, not sure the idiomatic way to do this. Celery redefines methods (so you can call `delay` on them to send them to workers) and I'm just doing it manually myself, so we don't need to use the `@celery.app.task` decorator (to support Celery & PSQ).
ericzinnikas: Thoughts on making passing evidence objects nicer (instead of passing `evidence_.__dict__` everywhere)?
aarontp: Thanks for the PR.  I'm just catching up after being out over the holidays, but I'll take a look at this tomorrow.
ericzinnikas: Good to go for a re-review.
ericzinnikas: Hmm, I'm pretty sure I avoided any cyclic dependency issues by importing `turbinia.client` inside the `deserialize` method. Tests complete fine, etc. but CodeFactor still doesn't seem to like that.
aarontp: Have you run a full test of this with a client/server/worker?  I tried checking out your branch and running it, but the server never picked up the request.  I haven't tried to debug it at all yet though.
ericzinnikas: There's actually one final issue with encoding/bytes/str that I'm trying to resolve (was working on the plane so this comment didn't post).
ericzinnikas: This last change fixes some py3/py2 issues (encoding). Tested a full run on Python 2.7.11 & Python 3.6.3rc1+ and it completes fine, except for:

```
[ERROR] HadoopAnalysisTask Task failed with exception: [write() argument must be str, not bytes]
[ERROR] Traceback (most recent call last):
  File "/home/ericwz/virtualenvs/turbinia/lib/python3.6/site-packages/turbinia/workers/__init__.py", line 533, in run_wrapper
    self.result = self.run(evidence, self.result)
  File "/home/ericwz/virtualenvs/turbinia/lib/python3.6/site-packages/turbinia/workers/hadoop.py", line 111, in run
    fh.write(output_evidence.text_data.encode('utf8'))
TypeError: write() argument must be str, not bytes
```

Reading https://docs.python.org/3/howto/pyporting.html seems this is due to how py3/py2 handle  `open()` and encoding/bytes differently, so that should probably be solved in a separate PR. Looks like converting all `open()` to `io.open()` is the way to go (or just making sure to open files in `wb` vs. `b`).
aarontp: Oops, looks like merging the last PR caused a conflict here.  I'll let you resolve it so I don't cause an issue with your branch (it looks like you do rebasing instead of merging as well).

I'll debug a bit to see what's going on with the GCP version.

I see you've pushed a couple new changes to the branch as well.  LMK  when it's ready so I can merge after I've tested.
ericzinnikas: Ah yeah, I forgot that other PR needed to be based on this one (ya we do rebase, no merging).

Good to go now for review/test (though CodeFactor sees more cyclic after the rebase -- works fine when I test though), last 3 commits didn't change anything substantial, just clarified a few things.
ericzinnikas: Just rebased again.
ericzinnikas: CI build is failing now, but seems related to yapf changes and not anything in this branch. @aarontp get a chance to test out the GCP issue?
aarontp: I did try to run this again at the end of last week and saw similar behavior (request wasn't picked up), but I didn't have much time to debug fully.  I'll build a fresh environment to make sure it's not something there and will debug.
aarontp: Ok, I just created a new environment and things run just fine there, so it must have been something in my environment.   Tests seem to be failing, but I think you just need to re-run YAPF (and resolve the conflict), and then I can merge this in.  Thanks!
ericzinnikas: Hmm, so `$ yapf -r -d turbinia/` returns cleanly when I run it, but Travis CI is still complaining.
aarontp: @ericzinnikas Do you have the latest version of YAPF installed?  It looks like Travis is using 0.26.0;  Does that match yours?  Maybe just need to `pip install --upgrade yapf` ?
ericzinnikas: Ah good call, seems like that resolved it.
aarontp: Great!  Merging now.

Patch:
diff --git a/turbinia/celery.py b/turbinia/celery.py
index 6137025..0212bd4 100644
--- a/turbinia/celery.py
+++ b/turbinia/celery.py
@@ -46,36 +46,6 @@ class TurbiniaCelery(object):
   def __init__(self):
     """Celery configurations."""
     self.app = None
-    self.fexec = None
-
-  def fexec(f, *args, **kwargs):
-    # pylint: disable=no-self-argument,method-hidden
-    """Placeholder function, overwritten with the _fexec closure once setup()
-    is called. Once overwritten, will call the specified function with
-    whichever arguments you pass it.
-
-    Arguments:
-      f (function): an arbitrary function
-      args: any positional arguments to be passsed to the function
-      kwargs: any keyword arguments to be passed to the function
-    """
-    raise NotImplementedError
-
-  def _fexec(self):
-    """Closure used to pass functions to workers. We use this instead of having
-    to annotate each individual TurbiniaTask with the Celery @app.task
-    decorator.
-
-    Returns:
-      function: the fexec() function.
-    """
-
-    @self.app.task(name='fexec')
-    def fexec(f, *args, **kwargs):
-      """Lets us pass in an arbitrary function without Celery annotations"""
-      return f(*args, **kwargs)
-
-    return fexec
 
   def setup(self):
     """Set up Celery"""
@@ -84,13 +54,7 @@ class TurbiniaCelery(object):
         'turbinia', broker=config.CELERY_BROKER, backend=config.CELERY_BACKEND)
     self.app.conf.update(
         task_default_queue=config.INSTANCE_ID,
-        # TODO(ericzinnikas): pickle is not secure, we need to replace it with
-        # the default json serializer, but need to figure out how to register
-        # the TurbiniaTask objects with Celery
-        event_serializer='pickle',
-        result_serializer='pickle',
-        task_serializer='pickle',
-        accept_content=['pickle'],
+        accept_content=['json'],
         # TODO(ericzinnikas): Without task_acks_late Celery workers will start
         # on one task and prefetch another (i.e. can result in 1 worker getting
         # 2 plaso jobs while another worker is free). But enabling this causes
@@ -100,7 +64,6 @@ class TurbiniaCelery(object):
         worker_concurrency=1,
         worker_prefetch_multiplier=1,
     )
-    self.fexec = self._fexec()
 
 
 class TurbiniaKombu(TurbiniaMessageBase):
diff --git a/turbinia/client.py b/turbinia/client.py
index 8e5d647..c27cb81 100644
--- a/turbinia/client.py
+++ b/turbinia/client.py
@@ -439,7 +439,8 @@ class TurbiniaCeleryWorker(TurbiniaClient):
   def start(self):
     """Start Turbinia Celery Worker."""
     log.info('Running Turbinia Celery Worker.')
-    argv = ['celery', 'worker', '--loglevel=info']
+    self.worker.task(task_manager.task_runner, name='task_runner')
+    argv = ['celery', 'worker', '--loglevel=info', '--pool=solo']
     self.worker.start(argv)
 
 
diff --git a/turbinia/processors/mount_local.py b/turbinia/processors/mount_local.py
index 13fc474..bde60c2 100644
--- a/turbinia/processors/mount_local.py
+++ b/turbinia/processors/mount_local.py
@@ -45,7 +45,8 @@ def PreprocessLosetup(source_path):
   losetup_command = ['sudo', 'losetup', '--show', '--find', '-P', source_path]
   log.info('Running command {0:s}'.format(' '.join(losetup_command)))
   try:
-    losetup_device = subprocess.check_output(losetup_command).strip()
+    losetup_device = subprocess.check_output(
+        losetup_command, universal_newlines=True).strip()
   except subprocess.CalledProcessError as e:
     raise TurbiniaException('Could not set losetup devices {0!s}'.format(e))
 
diff --git a/turbinia/task_manager.py b/turbinia/task_manager.py
index c4e7587..838390b 100644
--- a/turbinia/task_manager.py
+++ b/turbinia/task_manager.py
@@ -20,6 +20,7 @@ import logging
 import time
 
 import turbinia
+from turbinia import workers
 from turbinia import evidence
 from turbinia import config
 from turbinia import state_manager
@@ -73,6 +74,7 @@ def task_runner(obj, *args, **kwargs):
   Returns:
     Output from TurbiniaTask (should be TurbiniaTaskResult).
   """
+  obj = workers.TurbiniaTask.deserialize(obj)
   return obj.run_wrapper(*args, **kwargs)
 
 
@@ -293,11 +295,13 @@ class CeleryTaskManager(BaseTaskManager):
   Attributes:
     celery (TurbiniaCelery): Celery task queue, handles worker tasks.
     kombu (TurbiniaKombu): Kombu queue, handles receiving evidence.
+    celery_runner: task_runner method, but wrapped for Celery usage.
   """
 
   def __init__(self):
     self.celery = None
     self.kombu = None
+    self.celery_runner = None
     config.LoadConfig()
     super(CeleryTaskManager, self).__init__()
 
@@ -306,6 +310,7 @@ class CeleryTaskManager(BaseTaskManager):
     self.celery.setup()
     self.kombu = turbinia_celery.TurbiniaKombu(config.KOMBU_CHANNEL)
     self.kombu.setup()
+    self.celery_runner = self.celery.app.task(task_runner, name="task_runner")
 
   def process_tasks(self):
     """Determine the current state of our tasks.
@@ -324,7 +329,7 @@ class CeleryTaskManager(BaseTaskManager):
         log.warning('Task {0:s} failed.'.format(celery_task.id))
         completed_tasks.append(task)
       elif celery_task.status == celery_states.SUCCESS:
-        task.result = celery_task.result
+        task.result = workers.TurbiniaTaskResult.deserialize(celery_task.result)
         completed_tasks.append(task)
       else:
         log.debug('Task {0:s} status unknown'.format(celery_task.id))
@@ -356,7 +361,8 @@ class CeleryTaskManager(BaseTaskManager):
     log.info(
         'Adding Celery task {0:s} with evidence {1:s} to queue'.format(
             task.name, evidence_.name))
-    task.stub = self.celery.fexec.delay(task_runner, task, evidence_)
+    task.stub = self.celery_runner.delay(
+        task.serialize(), evidence_.serialize())
 
 
 class PSQTaskManager(BaseTaskManager):
@@ -413,7 +419,8 @@ class PSQTaskManager(BaseTaskManager):
         log.warning('Task {0:s} failed.'.format(psq_task.id))
         completed_tasks.append(task)
       else:
-        task.result = task.stub.result(timeout=PSQ_TASK_TIMEOUT)
+        task.result = workers.TurbiniaTaskResult.deserialize(
+            task.stub.result(timeout=PSQ_TASK_TIMEOUT))
         completed_tasks.append(task)
 
     outstanding_task_count = len(self.tasks) - len(completed_tasks)
@@ -438,4 +445,5 @@ class PSQTaskManager(BaseTaskManager):
     log.info(
         'Adding PSQ task {0:s} with evidence {1:s} to queue'.format(
             task.name, evidence_.name))
-    task.stub = self.psq.enqueue(task_runner, task, evidence_)
+    task.stub = self.psq.enqueue(
+        task_runner, task.serialize(), evidence_.serialize())
diff --git a/turbinia/workers/__init__.py b/turbinia/workers/__init__.py
index d2af59d..8a5dff5 100644
--- a/turbinia/workers/__init__.py
+++ b/turbinia/workers/__init__.py
@@ -16,7 +16,8 @@
 
 from __future__ import unicode_literals
 
-from datetime import datetime
+from copy import deepcopy
+from datetime import datetime, timedelta
 import getpass
 import logging
 import os
@@ -24,6 +25,7 @@ import pickle
 import platform
 import pprint
 import subprocess
+import sys
 import traceback
 import uuid
 import turbinia
@@ -31,6 +33,7 @@ import turbinia
 import filelock
 
 from turbinia import config
+from turbinia.evidence import evidence_decode
 from turbinia import output_manager
 from turbinia import TurbiniaException
 
@@ -65,26 +68,21 @@ class TurbiniaTaskResult(object):
   STORED_ATTRIBUTES = ['worker_name', 'status', 'saved_paths', 'successful']
 
   def __init__(
-      self, task, evidence=None, input_evidence=None, base_output_dir=None,
+      self, evidence=None, input_evidence=None, base_output_dir=None,
       request_id=None):
-    """Initialize the TurbiniaTaskResult object.
-
-    Args:
-      task (TurbiniaTask): The calling Task object
-
-    Raises:
-      TurbiniaException: If the Output Manager is not setup.
-    """
+    """Initialize the TurbiniaTaskResult object."""
 
     self.closed = False
     self.evidence = evidence if evidence else []
     self.input_evidence = input_evidence
     self.id = uuid.uuid4().hex
-    self.task_id = task.id
-    self.task_name = task.name
     self.base_output_dir = base_output_dir
     self.request_id = request_id
-    self.user = task.user
+
+    self.task_id = None
+    self.task_name = None
+    self.user = None
+    self.output_dir = None
 
     self.start_time = datetime.now()
     self.run_time = None
@@ -95,14 +93,28 @@ class TurbiniaTaskResult(object):
     self.worker_name = platform.node()
     # TODO(aarontp): Create mechanism to grab actual python logging data.
     self._log = []
+
+  def __str__(self):
+    return pprint.pformat(vars(self), depth=3)
+
+  def setup(self, task):
+    """Handles initializing task based attributes, after object creation.
+
+    Args:
+      task (TurbiniaTask): The calling Task object
+
+    Raises:
+      TurbiniaException: If the Output Manager is not setup.
+    """
+
+    self.task_id = task.id
+    self.task_name = task.name
+    self.user = task.user
     if task.output_manager.is_setup:
       _, self.output_dir = task.output_manager.get_local_output_dirs()
     else:
       raise TurbiniaException('Output Manager is not setup yet.')
 
-  def __str__(self):
-    return pprint.pformat(vars(self), depth=3)
-
   def close(self, task, success, status=None):
     """Handles closing of this result and writing logs.
 
@@ -217,6 +229,41 @@ class TurbiniaTaskResult(object):
     self.error['error'] = error
     self.error['traceback'] = traceback_
 
+  def serialize(self):
+    """Prepares result object for serialization.
+
+    Returns:
+      dict: Object dictionary that is JSON serializable.
+    """
+    self.run_time = self.run_time.total_seconds() if self.run_time else None
+    self.start_time = str(self.start_time)
+    if self.input_evidence:
+      self.input_evidence = self.input_evidence.serialize()
+    self.evidence = [x.serialize() for x in self.evidence]
+    return self.__dict__
+
+  @classmethod
+  def deserialize(cls, input_dict):
+    """Converts an input dictionary back into a TurbiniaTaskResult object.
+
+    Args:
+      input_dict (dict): TurbiniaTaskResult object dictionary.
+
+    Returns:
+      TurbiniaTaskResult: Deserialized object.
+    """
+    result = TurbiniaTaskResult()
+    result.__dict__.update(input_dict)
+    if result.run_time:
+      result.run_time = timedelta(seconds=result.run_time)
+    result.start_time = datetime.strptime(
+        result.start_time, '%Y-%m-%d %H:%M:%S.%f')
+    if result.input_evidence:
+      result.input_evidence = evidence_decode(result.input_evidence)
+    result.evidence = [evidence_decode(x) for x in result.evidence]
+
+    return result
+
 
 class TurbiniaTask(object):
   """Base class for Turbinia tasks.
@@ -268,6 +315,45 @@ class TurbiniaTask(object):
     self.user = user if user else getpass.getuser()
     self._evidence_config = {}
 
+  def serialize(self):
+    """Converts the TurbiniaTask object into a serializable dict.
+
+    Returns:
+      Dict: Dictionary representing this object, ready to be serialized.
+    """
+    task_copy = deepcopy(self.__dict__)
+    task_copy['output_manager'] = self.output_manager.__dict__
+    task_copy['last_update'] = str(self.last_update)
+    return task_copy
+
+  @classmethod
+  def deserialize(cls, input_dict):
+    """Converts an input dictionary back into a TurbiniaTask object.
+
+    Args:
+      input_dict (dict): TurbiniaTask object dictionary.
+
+    Returns:
+      TurbiniaTask: Deserialized object.
+    """
+    from turbinia import client  # Avoid circular imports
+
+    type_ = input_dict['name']
+    try:
+      task = getattr(sys.modules['turbinia.client'], type_)()
+    except AttributeError:
+      msg = (
+          "Could not import {0:s} object! Make sure it is imported where "
+          "this method is defined.".format(type_))
+      log.error(msg)
+      raise TurbiniaException(msg)
+    task.__dict__.update(input_dict)
+    task.output_manager = output_manager.OutputManager()
+    task.output_manager.__dict__.update(input_dict['output_manager'])
+    task.last_update = datetime.strptime(
+        input_dict['last_update'], '%Y-%m-%d %H:%M:%S.%f')
+    return task
+
   def execute(
       self, cmd, result, save_files=None, log_files=None, new_evidence=None,
       close=False, shell=False, success_codes=None):
@@ -364,8 +450,9 @@ class TurbiniaTask(object):
     self.tmp_dir, self.output_dir = self.output_manager.get_local_output_dirs()
     if not self.result:
       self.result = TurbiniaTaskResult(
-          task=self, input_evidence=evidence,
-          base_output_dir=self.base_output_dir, request_id=self.request_id)
+          input_evidence=evidence, base_output_dir=self.base_output_dir,
+          request_id=self.request_id)
+      self.result.setup(self)
 
     if not self.run_local:
       if evidence.copyable and not config.SHARED_FILESYSTEM:
@@ -423,8 +510,8 @@ class TurbiniaTask(object):
         old_status = 'No previous status'
 
       result = TurbiniaTaskResult(
-          task=self, base_output_dir=self.base_output_dir,
-          request_id=self.request_id)
+          base_output_dir=self.base_output_dir, request_id=self.request_id)
+      result.setup(self)
       result.status = '{0:s}. Previous status: [{1:s}]'.format(
           bad_message, old_status)
       result.set_error(bad_message, traceback.format_exc())
@@ -451,12 +538,13 @@ class TurbiniaTask(object):
       - Locking to make sure only one task is active at a time
 
     Args:
-      evidence: Evidence object
+      evidence (dict): To be decoded into Evidence object
 
     Returns:
       A TurbiniaTaskResult object
     """
     log.debug('Task {0:s} {1:s} awaiting execution'.format(self.name, self.id))
+    evidence = evidence_decode(evidence)
     with filelock.FileLock(config.LOCK_FILE):
       log.info('Starting Task {0:s} {1:s}'.format(self.name, self.id))
       original_result_id = None
@@ -532,9 +620,7 @@ class TurbiniaTask(object):
       log.debug(
           'Returning original result object {0:s} after task execution'.format(
               self.result.id))
-    # TODO(aarontp): Find a better way to ensure this gets unset.
-    self.output_manager = None
-    return self.result
+    return self.result.serialize()
 
   def run(self, evidence, result):
     """Entry point to execute the task.
diff --git a/turbinia/workers/artifact.py b/turbinia/workers/artifact.py
index 1adc5e3..5a30c83 100644
--- a/turbinia/workers/artifact.py
+++ b/turbinia/workers/artifact.py
@@ -26,7 +26,7 @@ from turbinia.workers import TurbiniaTask
 class FileArtifactExtractionTask(TurbiniaTask):
   """Task to run image_export (log2timeline)."""
 
-  def __init__(self, artifact_name):
+  def __init__(self, artifact_name='FileArtifact'):
     super(FileArtifactExtractionTask, self).__init__()
     self.artifact_name = artifact_name
 


Test patch:
diff --git a/turbinia/state_manager_test.py b/turbinia/state_manager_test.py
index e9e66f2..ca1c72c 100644
--- a/turbinia/state_manager_test.py
+++ b/turbinia/state_manager_test.py
@@ -62,8 +62,8 @@ class TestPSQStateManager(unittest.TestCase):
         '/fake/tmp/dir', self.base_output_dir)
 
     # Set up TurbiniaTaskResult
-    self.result = TurbiniaTaskResult(
-        task=self.task, base_output_dir=self.base_output_dir)
+    self.result = TurbiniaTaskResult(base_output_dir=self.base_output_dir)
+    self.result.setup(self.task)
     self.result.status = self.test_data['status']
     self.result.saved_paths = self.test_data['saved_paths']
     self.task.result = self.result
diff --git a/turbinia/workers/workers_test.py b/turbinia/workers/workers_test.py
index c068494..f90bb8d 100644
--- a/turbinia/workers/workers_test.py
+++ b/turbinia/workers/workers_test.py
@@ -26,6 +26,7 @@ from turbinia import evidence
 from turbinia import TurbiniaException
 from turbinia.workers import TurbiniaTask
 from turbinia.workers import TurbiniaTaskResult
+from turbinia.workers.plaso import PlasoTask
 
 
 class TestTurbiniaTask(unittest.TestCase):
@@ -37,6 +38,10 @@ class TestTurbiniaTask(unittest.TestCase):
 
     # Set up TurbiniaTask
     self.base_output_dir = tempfile.mkdtemp()
+    self.plaso_task = PlasoTask(base_output_dir=self.base_output_dir)
+    self.plaso_task.output_manager = mock.MagicMock()
+    self.plaso_task.output_manager.get_local_output_dirs.return_value = (
+        None, None)
     self.task = TurbiniaTask(base_output_dir=self.base_output_dir)
     self.task.output_manager = mock.MagicMock()
     self.task.output_manager.get_local_output_dirs.return_value = (None, None)
@@ -47,8 +52,8 @@ class TestTurbiniaTask(unittest.TestCase):
     self.evidence = evidence.RawDisk(local_path=test_disk_path)
 
     # Set up TurbiniaTaskResult
-    self.result = TurbiniaTaskResult(
-        task=self.task, base_output_dir=self.base_output_dir)
+    self.result = TurbiniaTaskResult(base_output_dir=self.base_output_dir)
+    self.result.setup(self.task)
 
     self.result.output_dir = self.base_output_dir
 
@@ -84,30 +89,43 @@ class TestTurbiniaTask(unittest.TestCase):
     self.task.run = mock.MagicMock(return_value=run)
     self.task.validate_result = mock.MagicMock(return_value=validate_result)
 
+  def testTurbiniaTaskSerialize(self):
+    """Test that we can properly serialize/deserialize tasks."""
+    out_dict = self.plaso_task.serialize()
+    out_obj = TurbiniaTask.deserialize(out_dict)
+    self.assertIsInstance(out_obj, PlasoTask)
+    # Nuke output_manager so we don't deal with class equality
+    self.plaso_task.output_manager = None
+    out_obj.output_manager = None
+    self.assertEqual(out_obj.__dict__, self.plaso_task.__dict__)
+
   def testTurbiniaTaskRunWrapper(self):
     """Test that the run wrapper executes task run."""
     self.setResults()
     self.result.closed = True
-    new_result = self.task.run_wrapper(self.evidence)
+    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()
 
   def testTurbiniaTaskRunWrapperAutoClose(self):
     """Test that the run wrapper closes the task."""
     self.setResults()
-    new_result = self.task.run_wrapper(self.evidence)
+    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_called()
 
   def testTurbiniaTaskRunWrapperBadResult(self):
     """Test that the run wrapper recovers from run returning bad result."""
     bad_result = 'Not a TurbiniaTaskResult'
-    checked_result = TurbiniaTaskResult(
-        task=self.task, base_output_dir=self.base_output_dir)
+    checked_result = TurbiniaTaskResult(base_output_dir=self.base_output_dir)
+    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)
+    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)
@@ -118,7 +136,8 @@ class TestTurbiniaTask(unittest.TestCase):
     self.setResults()
     self.task.run = mock.MagicMock(side_effect=TurbiniaException)
 
-    new_result = self.task.run_wrapper(self.evidence)
+    new_result = self.task.run_wrapper(self.evidence.__dict__)
+    new_result = TurbiniaTaskResult.deserialize(new_result)
     self.assertEqual(type(new_result), TurbiniaTaskResult)
     self.assertIn('failed', new_result.status)
 
@@ -132,7 +151,8 @@ class TestTurbiniaTask(unittest.TestCase):
     self.remove_files.append(
         os.path.join(self.task.base_output_dir, 'worker-log.txt'))
 
-    new_result = self.task.run_wrapper(self.evidence)
+    new_result = self.task.run_wrapper(self.evidence.__dict__)
+    new_result = TurbiniaTaskResult.deserialize(new_result)
     self.assertEqual(type(new_result), TurbiniaTaskResult)
     self.assertIn(canary_status, new_result.status)