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 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 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.status = 'CheckedResult' self.setResults(run=bad_result, validate_result=checked_result) new_result = self.task.run_wrapper(self.evidence) self.task.validate_result.assert_any_call(bad_result) self.assertEqual(type(new_result), TurbiniaTaskResult) self.assertIn('CheckedResult', new_result.status)
def run(results): """Celery task that reduces the results into one SQLite database. Args: results: List of returned values from Celery tasks. Returns: Task result object (instance of TurbiniaTaskResult) as JSON. """ job_id = results[0] cmd_output = subprocess.check_output( ['/usr/local/bin/be_reducer.sh', job_id]) result = TurbiniaTaskResult() result.add_result(result_type='PATH', result=cmd_output) return result.to_json()
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 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 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 testTurbiniaTaskJobUnavailable(self): """Test that the run wrapper can fail if the job doesn't exist.""" self.setResults() self.task.job_name = 'non_exist' canary_status = ('Task will not run due to the job: ' 'non_exist being disabled on the worker.') new_result = self.task.run_wrapper(self.evidence.__dict__) new_result = TurbiniaTaskResult.deserialize(new_result) self.assertEqual(new_result.status, canary_status)
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)
def testTurbiniaTaskRunWrapper(self): """Test that the run wrapper executes task run.""" self.unregisterMetrics() 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 run(src_path, out_path, job_id, **kwargs): """Task that process data with Plaso. Args: src_path: Path to data to process. out_path: Path to temporary storage of results. job_id: Unique ID for this task. Returns: Task result object (instance of TurbiniaTaskResult) as JSON. """ out_path = "{0:s}/{1:s}".format(out_path, job_id) if not os.path.exists(out_path): os.makedirs(out_path) cmd_output = subprocess.check_output(["/usr/local/bin/plaso_wrapper.sh", src_path, out_path, job_id]) res, version, metadata = cmd_output.split(" ", 2) result = TurbiniaTaskResult(version=version, metadata=json.loads(metadata)) result.add_result(result_type="PATH", result=res) return result.to_json()
def run(src_path, out_path, job_id, **kwargs): """Task that process data with Plaso. Args: src_path: Path to data to process. out_path: Path to temporary storage of results. job_id: Unique ID for this task. Returns: Task result object (instance of TurbiniaTaskResult) as JSON. """ out_path = '{0:s}/{1:s}'.format(out_path, job_id) if not os.path.exists(out_path): os.makedirs(out_path) cmd_output = subprocess.check_output( ['/usr/local/bin/plaso_wrapper.sh', src_path, out_path, job_id]) res, version, metadata = cmd_output.split(' ', 2) result = TurbiniaTaskResult(version=version, metadata=json.loads(metadata)) result.add_result(result_type='PATH', result=res) return result.to_json()
def testTurbiniaTaskRunWrapperSetupFail(self): """Test that the run wrapper recovers from setup failing.""" self.task.result = None canary_status = 'exception_message' self.task.setup = mock.MagicMock( side_effect=TurbiniaException('exception_message')) 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 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)
def setUp(self): self.remove_files = [] self.remove_dirs = [] # Set up TurbiniaTask self.base_output_dir = tempfile.mkdtemp() self.task = TurbiniaTask(base_output_dir=self.base_output_dir) self.task.output_manager = 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.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.output_dir = self.base_output_dir
def testTurbiniaTaskRunWrapperSetupFail(self, _, mock_create_result): """Test that the run wrapper recovers from setup failing.""" self.task.result = None canary_status = 'exception_message' self.task.setup = mock.MagicMock( side_effect=TurbiniaException('exception_message')) self.result.no_output_manager = True mock_create_result.return_value = self.result 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) # Checking specifically for `False` value and not whether this evaluates to # `False` because we don't want the `None` case to pass. self.assertEqual(new_result.successful, False) create_results_args = mock_create_result.call_args.kwargs self.assertIn(canary_status, create_results_args['message'])
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)
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 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)