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)