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 testTaskRunner(self, mock_task_deserialize):
"""Basic test for task_runner."""
task = PlasoTask()
task.run_wrapper = lambda x: x
mock_task_deserialize.return_value = task
task_dict = task.serialize()
ret = task_utils.task_runner(task_dict, 'testValue')
self.assertEqual(ret, 'testValue')
def testTaskDeserialize(self):
"""Basic test for task_deserialize."""
task = PlasoTask(request_id='testRequestID', requester='testRequester')
task_dict = task.serialize()
test_task = task_utils.task_deserialize(task_dict)
self.assertEqual(test_task.request_id, 'testRequestID')
self.assertEqual(test_task.requester, 'testRequester')
self.assertIsInstance(test_task, PlasoTask)
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.job_name = 'PlasoJob'
self.task.output_manager = mock.MagicMock()
self.task.output_manager.get_local_output_dirs.return_value = (None, None)
self.task.get_metrics = mock.MagicMock()
# Set up RawDisk Evidence
test_disk_path = tempfile.mkstemp(dir=self.base_output_dir)[1]
self.remove_files.append(test_disk_path)
self.test_stdout_path = tempfile.mkstemp(dir=self.base_output_dir)[1]
self.remove_files.append(self.test_stdout_path)
self.evidence = evidence.RawDisk(source_path=test_disk_path)
self.evidence.config['abort'] = False
self.evidence.config['globals'] = {}
self.evidence.preprocess = mock.MagicMock()
# Set up TurbiniaTaskResult
self.result = TurbiniaTaskResult(base_output_dir=self.base_output_dir)
self.result.output_dir = self.base_output_dir
def create_tasks(self, evidence):
"""Create task for Plaso.
Args:
evidence: List of evidence objects to process
Returns:
A list of PlasoTasks.
"""
return [PlasoTask() for _ in evidence]
def create_task(src_path, out_path, job_id=None, workers=1):
"""Create task for bulk_extractor.
Args:
src_path: Path to the data to process.
out_path: Path to where to put the result.
job_id: Unique identifier for the job (optional).
workers: Number of workers to run the Job on.
Returns:
A Celery task (instance of celery.Task).
"""
if not job_id:
job_id = uuid.uuid4().hex
task = PlasoTask().delay(src_path, out_path, job_id, workers=workers)
return task, job_id
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)
