Repository: google/turbinia
Pull request #415 (2019-07-31 00:42:25)

Add basic evidence validation
This validation happens on the worker instead of on the server so that the errors will still bubble up to the user.  By default this will just check that the attributes listed in the REQUIRED_ATTRIBUTES for the given evidence type are set, but each evidence object can potentially override this method to implement their own more stringent checks as required.


Patch:
diff --git a/turbinia/evidence.py b/turbinia/evidence.py
index 278036f..61a5bfc 100644
--- a/turbinia/evidence.py
+++ b/turbinia/evidence.py
@@ -100,6 +100,9 @@ class Evidence(object):
         relevant to us.
   """
 
+  # The list of attributes a given piece of Evidence requires to be set
+  REQUIRED_ATTRIBUTES = []
+
   def __init__(
       self, name=None, description=None, source=None, local_path=None,
       tags=None, request_id=None):
@@ -193,6 +196,26 @@ class Evidence(object):
     if self.parent_evidence:
       self.parent_evidence.postprocess()
 
+  def validate(self):
+    """Runs validation to verify evidence meets minimum requirements.
+
+    This default implementation will just check that the attributes listed in
+    REQUIRED_ATTRIBUTES are set, but other evidence types can override this
+    method to implement their own more stringent checks as needed.  This is
+    called by the worker, prior to the pre/post-processors running.
+
+    Raises:
+      TurbiniaException: If validation fails
+    """
+    for attribute in self.REQUIRED_ATTRIBUTES:
+      attribute_value = getattr(self, attribute, None)
+      if not attribute_value:
+        message = (
+            'Evidence validation failed: Required attribute {0:s} for class '
+            '{1:s} is not set. Please check original request.'.format(
+                attribute, self.name))
+        raise TurbiniaException(message)
+
 
 class Directory(Evidence):
   """Filesystem directory evidence."""
@@ -208,6 +231,8 @@ class ChromiumProfile(Evidence):
     format: Output format (default is sqlite, other options are xlsx and jsonl)
   """
 
+  REQUIRED_ATTRIBUTES = ['browser_type', 'output_format']
+
   def __init__(self, browser_type=None, output_format=None, *args, **kwargs):
     """Initialization for chromium profile evidence object."""
     super(ChromiumProfile, self).__init__(*args, **kwargs)
@@ -274,6 +299,8 @@ class BitlockerDisk(EncryptedDisk):
     unencrypted_path: A string to the unencrypted local path
   """
 
+  REQUIRED_ATTRIBUTES = ['recovery_key', 'password']
+
   def __init__(self, recovery_key=None, password=None, *args, **kwargs):
     """Initialization for Bitlocker disk evidence object"""
     self.recovery_key = recovery_key
@@ -292,6 +319,8 @@ class APFSEncryptedDisk(EncryptedDisk):
     unencrypted_path: A string to the unencrypted local path
   """
 
+  REQUIRED_ATTRIBUTES = ['recovery_key', 'password']
+
   def __init__(self, recovery_key=None, password=None, *args, **kwargs):
     """Initialization for Bitlocker disk evidence object"""
     self.recovery_key = recovery_key
@@ -309,6 +338,8 @@ class GoogleCloudDisk(RawDisk):
     disk_name: The cloud disk name.
   """
 
+  REQUIRED_ATTRIBUTES = ['disk_name', 'project', 'zone']
+
   def __init__(self, project=None, zone=None, disk_name=None, *args, **kwargs):
     """Initialization for Google Cloud Disk."""
     self.project = project
@@ -337,6 +368,8 @@ class GoogleCloudDiskRawEmbedded(GoogleCloudDisk):
     embedded_path: The path of the raw disk image inside the Persistent Disk
   """
 
+  REQUIRED_ATTRIBUTES = ['disk_name', 'project', 'zone', 'embedded_path']
+
   def __init__(self, embedded_path=None, *args, **kwargs):
     """Initialization for Google Cloud Disk."""
     self.embedded_path = embedded_path
@@ -404,6 +437,8 @@ class FilteredTextFile(TextFile):
 class ExportedFileArtifact(Evidence):
   """Exported file artifact."""
 
+  REQUIRED_ATTRIBUTES = ['artifact_name']
+
   def __init__(self, artifact_name=None):
     """Initializes an exported file artifact."""
     super(ExportedFileArtifact, self).__init__()
@@ -424,6 +459,8 @@ class RawMemory(Evidence):
     module_list (list): Module used for the analysis
     """
 
+  REQUIRED_ATTRIBUTES = ['module_list', 'profile']
+
   def __init__(self, module_list=None, profile=None, *args, **kwargs):
     """Initialization for raw memory evidence object."""
     super(RawMemory, self).__init__(*args, **kwargs)
diff --git a/turbinia/workers/__init__.py b/turbinia/workers/__init__.py
index ce65023..3548cad 100644
--- a/turbinia/workers/__init__.py
+++ b/turbinia/workers/__init__.py
@@ -592,6 +592,7 @@ class TurbiniaTask(object):
       try:
         self.result = self.setup(evidence)
         original_result_id = self.result.id
+        evidence.validate()
 
         if self.turbinia_version != turbinia.__version__:
           message = (


Test patch:
diff --git a/turbinia/evidence_test.py b/turbinia/evidence_test.py
index 15a9c35..d1c02f4 100644
--- a/turbinia/evidence_test.py
+++ b/turbinia/evidence_test.py
@@ -46,3 +46,24 @@ class TestTurbiniaEvidence(unittest.TestCase):
     """Test that evidence_decode throws error on dict with no type attribute."""
     test = {1: 2, 3: 4}
     self.assertRaises(TurbiniaException, evidence.evidence_decode, test)
+
+  def testEvidenceValidation(self):
+    """Test successful evidence validation."""
+    rawdisk = evidence.RawDisk(
+        name='My Evidence', local_path='/tmp/foo', mount_path='/mnt/foo')
+    rawdisk.REQUIRED_ATTRIBUTES = ['name', 'local_path', 'mount_path']
+    rawdisk.validate()
+
+  def testEvidenceValidationEmptyAttribute(self):
+    """Test failed evidence validation with an empty attribute."""
+    rawdisk = evidence.RawDisk(
+        name='My Evidence', local_path=None, mount_path='/mnt/foo')
+    rawdisk.REQUIRED_ATTRIBUTES = ['name', 'local_path', 'mount_path']
+    self.assertRaises(TurbiniaException, rawdisk.validate)
+
+  def testEvidenceValidationNoAttribute(self):
+    """Test failed evidence validation with no attribute."""
+    rawdisk = evidence.RawDisk(
+        name='My Evidence', local_path='/tmp/foo', mount_path='/mnt/foo')
+    rawdisk.REQUIRED_ATTRIBUTES = ['doesnotexist']
+    self.assertRaises(TurbiniaException, rawdisk.validate)
diff --git a/turbinia/workers/workers_test.py b/turbinia/workers/workers_test.py
index 65a1c8a..1b45874 100644
--- a/turbinia/workers/workers_test.py
+++ b/turbinia/workers/workers_test.py
@@ -191,6 +191,18 @@ class TestTurbiniaTask(TestTurbiniaTaskBase):
     self.assertEqual(type(new_result), TurbiniaTaskResult)
     self.assertNotEqual(new_result.error, {})
 
+  @mock.patch('turbinia.workers.evidence_decode')
+  def testTurbiniaTaskEvidenceValidationFailure(self, evidence_decode_mock):
+    """Tests Task fails when evidence validation fails."""
+    self.setResults()
+    test_evidence = evidence.RawDisk()
+    test_evidence.REQUIRED_ATTRIBUTES = ['doesnotexist']
+    evidence_decode_mock.return_value = test_evidence
+    test_result = self.task.run_wrapper(test_evidence.__dict__)
+    test_result = TurbiniaTaskResult.deserialize(test_result)
+    self.assertFalse(test_result.successful)
+    self.assertIn('validation failed', test_result.status)
+
   @mock.patch('turbinia.workers.subprocess.Popen')
   def testTurbiniaTaskExecute(self, popen_mock):
     """Test execution with success case."""