def testTurbiniaTaskRunWrapperBadResult(self, _):
"""Test that the run wrapper recovers from run returning bad result."""
self.unregisterMetrics()
bad_result = 'Not a TurbiniaTaskResult'
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.__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)
class TestPSQStateManager(unittest.TestCase):
"""Test PSQStateManager class."""
def _get_state_manager(self):
"""Gets a Datastore State Manager object for test."""
config.STATE_MANAGER = 'Datastore'
return state_manager.get_state_manager()
@mock.patch('turbinia.state_manager.datastore.Client')
def setUp(self, _):
self.remove_files = []
self.remove_dirs = []
self.state_manager = None
config.LoadConfig()
self.state_manager_save = config.STATE_MANAGER
self.test_data = {
'name': 'TestTask',
'request_id': 'TestRequestId',
'status': 'TestStatus',
'saved_paths': ['testpath1', 'testpath2']
}
# Set up TurbiniaTask
self.base_output_dir = tempfile.mkdtemp()
self.task = TurbiniaTask(base_output_dir=self.base_output_dir,
name=self.test_data['name'],
request_id=self.test_data['request_id'])
self.task.output_manager = mock.MagicMock()
self.task.output_manager.get_local_output_dirs.return_value = (
'/fake/tmp/dir', self.base_output_dir)
# Set up TurbiniaTaskResult
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
def tearDown(self):
config.STATE_MANAGER = self.state_manager_save
[os.remove(f) for f in self.remove_files if os.path.exists(f)]
[os.rmdir(d) for d in self.remove_dirs if os.path.exists(d)]
os.rmdir(self.base_output_dir)
@mock.patch('turbinia.state_manager.datastore.Client')
def testStateManagerGetTaskDict(self, _):
"""Test State Manager get_task_dict()."""
self.state_manager = self._get_state_manager()
task_dict = self.state_manager.get_task_dict(self.task)
# Make the returned task_dict contains all of our test data
self.assertEqual(task_dict['name'], self.test_data['name'])
self.assertEqual(task_dict['request_id'], self.test_data['request_id'])
self.assertEqual(task_dict['status'], self.test_data['status'])
self.assertEqual(len(task_dict['saved_paths']), 2)
self.assertTrue('instance' in task_dict)
self.assertIn(self.test_data['saved_paths'][0],
task_dict['saved_paths'])
@mock.patch('turbinia.state_manager.datastore.Client')
def testStateManagerValidateDataValidDict(self, _):
"""Test State Manager _validate_data() base case."""
self.state_manager = self._get_state_manager()
# pylint: disable=protected-access
test_data = self.state_manager._validate_data(self.test_data)
self.assertDictEqual(test_data, self.test_data)
@mock.patch('turbinia.state_manager.datastore.Client')
def testStateManagerValidateDataInvalidDict(self, _):
"""Test State Manager _validate_data() base case."""
self.state_manager = self._get_state_manager()
invalid_dict = copy.deepcopy(self.test_data)
invalid_dict[
'status'] = 'A' * state_manager.MAX_DATASTORE_STRLEN + 'BORKEN'
# pylint: disable=protected-access
test_data = self.state_manager._validate_data(invalid_dict)
self.assertListEqual(list(test_data.keys()),
list(self.test_data.keys()))
self.assertNotEqual(test_data['status'], self.test_data['status'])
self.assertLessEqual(len(test_data['status']),
state_manager.MAX_DATASTORE_STRLEN)
class TestTurbiniaTaskBase(unittest.TestCase):
"""Test TurbiniaTask class.
Attributes:
class_task(TurbiniaTask): The class the test should instantiated
remove_file(list(str)): Files that will be removed after the test run
remove_dirs(list(str)): Dirs that will be removed after the test run
base_output_dir(str): The base output directory used by the Task
task(TurbiniaTask): The instantiated Task under test
evidence(Evidence): The test evidence object used by the Task
result(TurbiniaResult): The result object used by the Task
"""
def setUp(self, task_class=TurbiniaTask, evidence_class=evidence.RawDisk):
self.task_class = task_class
self.evidence_class = evidence_class
self.remove_files = []
self.remove_dirs = []
# Set up Tasks under test
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 = self.task_class(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)
# Set up Evidence
test_artifact_path = tempfile.mkstemp(dir=self.base_output_dir)[1]
self.remove_files.append(test_artifact_path)
self.evidence = self.evidence_class(local_path=test_artifact_path)
# Set up TurbiniaTaskResult
self.result = TurbiniaTaskResult(base_output_dir=self.base_output_dir)
self.result.setup(self.task)
self.result.output_dir = self.base_output_dir
def tearDown(self):
for remove_file in self.remove_files:
if os.path.exists(remove_file):
os.remove(remove_file)
for directory in self.remove_dirs:
if os.path.exists(directory):
os.rmdir(directory)
os.rmdir(self.base_output_dir)
def setResults(self,
setup=None,
run=None,
validate_result=None,
mock_run=True):
"""Set up mock returns in TurbiniaTaskResult object.
Args:
setup: What value to return from setup()
run: What value to return from run()
validate_result: What value to return from validate_result()
mock_run(bool): Whether to mock out the run method
"""
if setup is None:
setup = self.result
if run is None:
run = self.result
if validate_result is None:
validate_result = self.result
self.result.status = 'TestStatus'
self.result.close = mock.MagicMock()
self.task.setup = mock.MagicMock(return_value=setup)
if mock_run:
self.task.run = mock.MagicMock(return_value=run)
self.task.validate_result = mock.MagicMock(
return_value=validate_result)
class TestTurbiniaTask(unittest.TestCase):
"""Test TurbiniaTask class."""
def setUp(self):
self.remove_files = []
self.remove_dirs = []
# 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)
# Set up RawDisk Evidence
test_disk_path = tempfile.mkstemp(dir=self.base_output_dir)[1]
self.remove_files.append(test_disk_path)
self.evidence = evidence.RawDisk(local_path=test_disk_path)
# Set up TurbiniaTaskResult
self.result = TurbiniaTaskResult(base_output_dir=self.base_output_dir)
self.result.setup(self.task)
self.result.output_dir = self.base_output_dir
def tearDown(self):
for remove_file in self.remove_files:
if os.path.exists(remove_file):
os.remove(remove_file)
for directory in self.remove_dirs:
if os.path.exists(directory):
os.rmdir(directory)
os.rmdir(self.base_output_dir)
def setResults(self, setup=None, run=None, validate_result=None):
"""Set up mock returns.
Args:
setup: What value to return from setup()
run: What value to return from run()
validate_result: What value to return from validate_result()
"""
if setup is None:
setup = self.result
if run is None:
run = self.result
if validate_result is None:
validate_result = self.result
self.result.status = 'TestStatus'
self.result.close = mock.MagicMock()
self.task.setup = mock.MagicMock(return_value=setup)
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.__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.__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(
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.__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)
def testTurbiniaTaskRunWrapperExceptionThrown(self):
"""Test that the run wrapper recovers from run throwing an exception."""
self.setResults()
self.task.run = mock.MagicMock(side_effect=TurbiniaException)
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)
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)
self.remove_files.append(
os.path.join(self.task.base_output_dir, 'worker-log.txt'))
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)
def testTurbiniaTaskValidateResultGoodResult(self):
"""Tests validate_result with good result."""
self.result.status = 'GoodStatus'
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, _):
"""Tests validate_result with bad result."""
# Passing in an unpickleable object (json module) and getting back a
# TurbiniaTaskResult
new_result = self.task.validate_result(json)
self.assertEqual(type(new_result), TurbiniaTaskResult)
self.assertNotEqual(new_result.error, {})
@mock.patch('turbinia.workers.subprocess.Popen')
def testTurbiniaTaskExecute(self, popen_mock):
"""Test execution with success case."""
cmd = 'test cmd'
output = ('test stdout', 'test stderr')
self.result.close = mock.MagicMock()
proc_mock = mock.MagicMock()
proc_mock.communicate.return_value = output
proc_mock.returncode = 0
popen_mock.return_value = proc_mock
self.task.execute(cmd, self.result, close=True)
# Command was executed, has the correct output saved and
# TurbiniaTaskResult.close() was called with successful status.
popen_mock.assert_called_with(cmd)
self.assertEqual(self.result.error['stdout'], output[0])
self.assertEqual(self.result.error['stderr'], output[1])
self.result.close.assert_called_with(self.task, success=True)
@mock.patch('turbinia.workers.subprocess.Popen')
def testTurbiniaTaskExecuteFailure(self, popen_mock):
"""Test execution with failure case."""
cmd = 'test cmd'
output = ('test stdout', 'test stderr')
self.result.close = mock.MagicMock()
proc_mock = mock.MagicMock()
proc_mock.communicate.return_value = output
proc_mock.returncode = 1
popen_mock.return_value = proc_mock
self.task.execute(cmd, self.result, close=True)
# Command was executed and TurbiniaTaskResult.close() was called with
# unsuccessful status.
popen_mock.assert_called_with(cmd)
self.result.close.assert_called_with(self.task,
success=False,
status=mock.ANY)
@mock.patch('turbinia.workers.subprocess.Popen')
def testTurbiniaTaskExecuteEvidenceExists(self, popen_mock):
"""Test execution with new evidence that has valid a local_path."""
cmd = 'test cmd'
output = ('test stdout', 'test stderr')
self.result.close = mock.MagicMock()
proc_mock = mock.MagicMock()
proc_mock.communicate.return_value = output
proc_mock.returncode = 0
popen_mock.return_value = proc_mock
# Create our evidence local path file
with open(self.evidence.local_path, 'w') as evidence_path:
evidence_path.write('test')
self.task.execute(cmd,
self.result,
new_evidence=[self.evidence],
close=True)
self.assertIn(self.evidence, self.result.evidence)
@mock.patch('turbinia.workers.subprocess.Popen')
def testTurbiniaTaskExecuteEvidenceDoesNotExist(self, popen_mock):
"""Test execution with new evidence that does not have a local_path."""
cmd = 'test cmd'
output = ('test stdout', 'test stderr')
self.result.close = mock.MagicMock()
proc_mock = mock.MagicMock()
proc_mock.communicate.return_value = output
proc_mock.returncode = 0
popen_mock.return_value = proc_mock
os.remove(self.evidence.local_path)
self.task.execute(cmd,
self.result,
new_evidence=[self.evidence],
close=True)
self.assertNotIn(self.evidence, self.result.evidence)
@mock.patch('turbinia.workers.subprocess.Popen')
def testTurbiniaTaskExecuteEvidenceExistsButEmpty(self, popen_mock):
"""Test execution with new evidence local_path that exists but is empty."""
cmd = 'test cmd'
output = ('test stdout', 'test stderr')
self.result.close = mock.MagicMock()
proc_mock = mock.MagicMock()
proc_mock.communicate.return_value = output
proc_mock.returncode = 0
popen_mock.return_value = proc_mock
# Exists and is empty
self.assertTrue(os.path.exists(self.evidence.local_path))
self.assertEqual(os.path.getsize(self.evidence.local_path), 0)
self.task.execute(cmd,
self.result,
new_evidence=[self.evidence],
close=True)
self.assertNotIn(self.evidence, self.result.evidence)
