def testFlowReportsErrorWhenCollectingSingleFileAboveTotalLimit(self): with temp.AutoTempFilePath() as temp_file_path: target_bytes = 2**20 # 1 MiB less_than_necessary_bytes = target_bytes // 2 with io.open(temp_file_path, "wb") as fd: fd.write(b"1" * target_bytes) table = f""" [ {{ "collect_column": "{temp_file_path}" }} ] """ with mock.patch.object(osquery_flow, "FILE_COLLECTION_MAX_TOTAL_BYTES", less_than_necessary_bytes): with osquery_test_lib.FakeOsqueryiOutput(stdout=table, stderr=""): flow_id = self._InitializeFlow( file_collection_columns=["collect_column"], check_flow_errors=False) progress = flow_test_lib.GetFlowProgress(self.client_id, flow_id) self.assertEqual( "Files for collection exceed the total size limit of " f"{less_than_necessary_bytes} bytes.", progress.error_message)
def testFailsAfterRetryOnFailedFetchOnWindows(self): temp_bar_file = self.create_tempfile() temp_bar_file.write_bytes(b"bar") file_bar_path = temp_bar_file.full_path with mock.patch.object(flow_base.FlowBase, "client_os", "Windows"): with mock.patch.object(vfs, "VFSOpen", side_effect=IOError("mock err")): flow_id = flow_test_lib.TestFlowHelper( file.CollectFilesByKnownPath.__name__, self.client_mock, client_id=self.client_id, paths=[file_bar_path], creator=self.test_username) progress = flow_test_lib.GetFlowProgress(self.client_id, flow_id) self.assertEqual( rdf_file_finder.CollectFilesByKnownPathProgress( num_in_progress=0, num_raw_fs_access_retries=1, num_collected=0, num_failed=1, ), progress) results = flow_test_lib.GetFlowResults(self.client_id, flow_id) self.assertLen(results, 1) self.assertEqual(results[0].stat.pathspec.path, file_bar_path) self.assertEqual( results[0].stat.pathspec.pathtype, config.CONFIG["Server.raw_filesystem_access_pathtype"]) self.assertIsNone(results[0].hash.sha1) self.assertEqual( results[0].status, rdf_file_finder.CollectFilesByKnownPathResult.Status.FAILED)
def testFileNotFound(self): temp_dir = self.create_tempdir() non_existent_file_path = os.path.join(temp_dir.full_path, "/non_existent") flow_id = flow_test_lib.TestFlowHelper( file.CollectFilesByKnownPath.__name__, self.client_mock, client_id=self.client_id, paths=[non_existent_file_path], creator=self.test_username) progress = flow_test_lib.GetFlowProgress(self.client_id, flow_id) self.assertEqual( rdf_file_finder.CollectFilesByKnownPathProgress( num_in_progress=0, num_raw_fs_access_retries=0, num_collected=0, num_failed=1, ), progress) results = flow_test_lib.GetFlowResults(self.client_id, flow_id) self.assertLen(results, 1) self.assertEqual(results[0].stat.pathspec.path, non_existent_file_path) self.assertEqual(results[0].stat.pathspec.pathtype, rdf_paths.PathSpec.PathType.OS) self.assertEqual( results[0].status, rdf_file_finder.CollectFilesByKnownPathResult.Status.NOT_FOUND)
def testCorrectlyReportsProgressInFlight(self): flow_id = flow_test_lib.StartFlow(file.CollectSingleFile, client_id=self.client_id, path="/some/path") progress = flow_test_lib.GetFlowProgress(self.client_id, flow_id) self.assertEqual( progress.status, rdf_file_finder.CollectSingleFileProgress.Status.IN_PROGRESS)
def testFetchIsRetriedWithRawOnWindows(self): temp_bar_file = self.create_tempfile() temp_bar_file.write_bytes(b"bar") file_bar_path = temp_bar_file.full_path file_bar_hash = hashlib.sha1(b"bar").hexdigest() with mock.patch.object(flow_base.FlowBase, "client_os", "Windows"): with _PatchVfs(): flow_id = flow_test_lib.TestFlowHelper( file.CollectFilesByKnownPath.__name__, self.client_mock, client_id=self.client_id, paths=[file_bar_path], creator=self.test_username) progress = flow_test_lib.GetFlowProgress(self.client_id, flow_id) self.assertEqual( rdf_file_finder.CollectFilesByKnownPathProgress( num_in_progress=0, num_raw_fs_access_retries=1, num_collected=1, num_failed=0, ), progress) results = flow_test_lib.GetFlowResults(self.client_id, flow_id) self.assertLen(results, 4) for result in results: self.assertEqual(result.stat.pathspec.path, file_bar_path) expected_pathtypes = [ # First attempt (failure) rdf_paths.PathSpec.PathType.OS, # Second attempt (success) config.CONFIG["Server.raw_filesystem_access_pathtype"], config.CONFIG["Server.raw_filesystem_access_pathtype"], config.CONFIG["Server.raw_filesystem_access_pathtype"] ] self.assertCountEqual(expected_pathtypes, self._getResultsPathType(results)) expected_statuses = [ # First attempt (failure) rdf_file_finder.CollectFilesByKnownPathResult.Status.IN_PROGRESS, # Second attempt (success) rdf_file_finder.CollectFilesByKnownPathResult.Status.IN_PROGRESS, rdf_file_finder.CollectFilesByKnownPathResult.Status.IN_PROGRESS, rdf_file_finder.CollectFilesByKnownPathResult.Status.COLLECTED ] self.assertCountEqual(expected_statuses, self._getResultsStatus(results)) expected_hashes = ["None", "None", file_bar_hash, file_bar_hash] self.assertCountEqual(expected_hashes, self._getResultsHashes(results))
def testProgressCorrectlyIndicatesNotFoundStatus(self): flow_id = flow_test_lib.StartAndRunFlow(file.CollectSingleFile, self.client_mock, client_id=self.client_id, check_flow_errors=False, path="/nonexistent") progress = flow_test_lib.GetFlowProgress(self.client_id, flow_id) self.assertEqual( progress.status, rdf_file_finder.CollectSingleFileProgress.Status.NOT_FOUND)
def _RunCollectBrowserHistory(self, **kwargs): flow_args = webhistory.CollectBrowserHistoryArgs(**kwargs) flow_id = flow_test_lib.StartAndRunFlow( webhistory.CollectBrowserHistory, creator=self.token.username, client_mock=action_mocks.ActionMock(), client_id=self.client_id, flow_args=flow_args) results = flow_test_lib.GetFlowResults(self.client_id, flow_id) progress = flow_test_lib.GetFlowProgress(self.client_id, flow_id) return flow_id, results, progress
def testFlowReportsErrorWhenCollectingColumnsAboveLimit(self): with temp.AutoTempFilePath() as temp_file_path: with io.open(temp_file_path, mode="w", encoding="utf-8") as fd: fd.write("Just sample text to put in the file.") with mock.patch.object(osquery_flow, "FILE_COLLECTION_MAX_COLUMNS", 1): flow_id = self._InitializeFlow( file_collection_columns=["collect1", "collect2"], check_flow_errors=False) progress = flow_test_lib.GetFlowProgress(self.client_id, flow_id) self.assertEqual( "Requested file collection for 2 columns, but the limit is 1 columns.", progress.error_message)
def testProgressContainsResultOnSuccess(self): flow_id = flow_test_lib.StartAndRunFlow( file.CollectSingleFile, self.client_mock, client_id=self.client_id, path=self.files["bar"].path) progress = flow_test_lib.GetFlowProgress(self.client_id, flow_id) self.assertEqual(progress.status, rdf_file_finder.CollectSingleFileProgress.Status.COLLECTED) self.assertEqual(progress.result.stat.pathspec.path, self.files["bar"].path) self.assertEqual(progress.result.stat.pathspec.pathtype, rdf_paths.PathSpec.PathType.OS) self.assertEqual(str(progress.result.hash.sha1), self.files["bar"].sha1)
def testSmallerTruncationLimit(self): two_row_table = """ [ { "col1": "cell-1-1", "col2": "cell-1-2", "col3": "cell-1-3" }, { "col1": "cell-2-1", "col2": "cell-2-2", "col3": "cell-2-3" } ] """ max_rows = 1 with osquery_test_lib.FakeOsqueryiOutput(stdout=two_row_table, stderr=""): with mock.patch.object(osquery_flow, "TRUNCATED_ROW_COUNT", max_rows): flow_id = self._InitializeFlow() progress = flow_test_lib.GetFlowProgress(self.client_id, flow_id) self.assertEqual(len(progress.partial_table.rows), max_rows)
def testProgressCorrectlyIndicatesErrorStatus(self): with mock.patch.object(flow_base.FlowBase, "client_os", "Windows"): with mock.patch.object(vfs, "VFSOpen", side_effect=IOError("mock err")): flow_id = flow_test_lib.StartAndRunFlow( file.CollectSingleFile, self.client_mock, client_id=self.client_id, check_flow_errors=False, path="/nonexistent") progress = flow_test_lib.GetFlowProgress(self.client_id, flow_id) self.assertEqual(progress.status, rdf_file_finder.CollectSingleFileProgress.Status.FAILED) self.assertEqual( progress.error_description, f"mock err when fetching /nonexistent with {config.CONFIG['Server.raw_filesystem_access_pathtype']}" )
def testTotalRowCountIncludesAllChunks(self): row_count = 100 split_pieces = 10 cell_value = "fixed" table = [{"column1": cell_value}] * row_count table_json = json.dumps(table) table_bytes = row_count * len(cell_value.encode("utf-8")) chunk_bytes = table_bytes // split_pieces with test_lib.ConfigOverrider({"Osquery.max_chunk_size": chunk_bytes}): with osquery_test_lib.FakeOsqueryiOutput(stdout=table_json, stderr=""): flow_id = self._InitializeFlow() progress = flow_test_lib.GetFlowProgress(self.client_id, flow_id) self.assertEqual(progress.total_row_count, row_count)
def testChunksSmallerThanTruncation(self): row_count = 100 max_rows = 70 split_pieces = 10 cell_value = "fixed" table = [{"column1": cell_value}] * row_count table_json = json.dumps(table) table_bytes = row_count * len(cell_value.encode("utf-8")) chunk_bytes = table_bytes // split_pieces with test_lib.ConfigOverrider({"Osquery.max_chunk_size": chunk_bytes}): with osquery_test_lib.FakeOsqueryiOutput(stdout=table_json, stderr=""): with mock.patch.object(osquery_flow, "TRUNCATED_ROW_COUNT", max_rows): flow_id = self._InitializeFlow() progress = flow_test_lib.GetFlowProgress(self.client_id, flow_id) self.assertEqual(len(progress.partial_table.rows), max_rows)
def testCorrectlyReportProgressForSuccessfullyCollectedFiles(self): path = os.path.join(self.fixture_path, "b*") flow_id = flow_test_lib.TestFlowHelper( file.CollectMultipleFiles.__name__, self.client_mock, client_id=self.client_id, path_expressions=[path], creator=self.test_username) progress = flow_test_lib.GetFlowProgress(self.client_id, flow_id) self.assertEqual( rdf_file_finder.CollectMultipleFilesProgress( num_collected=2, num_failed=0, num_found=2, num_in_progress=0, num_raw_fs_access_retries=0, ), progress)
def testReturnsMultipleFilesStat(self): temp_bar_file = self.create_tempfile() temp_bar_file.write_bytes(b"bar") file_bar_path = temp_bar_file.full_path temp_foo_file = self.create_tempfile() temp_foo_file.write_bytes(b"foo") file_foo_path = temp_foo_file.full_path flow_id = flow_test_lib.TestFlowHelper( file.CollectFilesByKnownPath.__name__, self.client_mock, client_id=self.client_id, paths=[file_bar_path, file_foo_path], collection_level=rdf_file_finder.CollectFilesByKnownPathArgs. CollectionLevel.STAT, creator=self.test_username) progress = flow_test_lib.GetFlowProgress(self.client_id, flow_id) self.assertEqual( rdf_file_finder.CollectFilesByKnownPathProgress( num_in_progress=0, num_raw_fs_access_retries=0, num_collected=2, num_failed=0, ), progress) results = flow_test_lib.GetFlowResults(self.client_id, flow_id) self.assertLen(results, 2) self.assertEqual(results[0].stat.pathspec.path, file_bar_path) self.assertEqual(results[0].stat.pathspec.pathtype, rdf_paths.PathSpec.PathType.OS) self.assertFalse(results[0].hash) self.assertEqual( results[0].status, rdf_file_finder.CollectFilesByKnownPathResult.Status.COLLECTED) self.assertEqual(results[1].stat.pathspec.path, file_foo_path) self.assertEqual(results[1].stat.pathspec.pathtype, rdf_paths.PathSpec.PathType.OS) self.assertFalse(results[1].hash) self.assertEqual( results[1].status, rdf_file_finder.CollectFilesByKnownPathResult.Status.COLLECTED)
def testReturnsSingleFile(self): temp_bar_file = self.create_tempfile() temp_bar_file.write_bytes(b"bar") file_bar_path = temp_bar_file.full_path file_bar_hash = hashlib.sha1(b"bar").hexdigest() flow_id = flow_test_lib.TestFlowHelper( file.CollectFilesByKnownPath.__name__, self.client_mock, client_id=self.client_id, paths=[file_bar_path], creator=self.test_username) progress = flow_test_lib.GetFlowProgress(self.client_id, flow_id) self.assertEqual( rdf_file_finder.CollectFilesByKnownPathProgress( num_in_progress=0, num_raw_fs_access_retries=0, num_collected=1, num_failed=0, ), progress) results = flow_test_lib.GetFlowResults(self.client_id, flow_id) self.assertLen(results, 3) for result in results: self.assertEqual(result.stat.pathspec.pathtype, rdf_paths.PathSpec.PathType.OS) expected_paths = [file_bar_path, file_bar_path, file_bar_path] self.assertCountEqual(expected_paths, self._getResultsPaths(results)) expected_hashes = ["None", file_bar_hash, file_bar_hash] self.assertCountEqual(expected_hashes, self._getResultsHashes(results)) expected_statuses = [ rdf_file_finder.CollectFilesByKnownPathResult.Status.IN_PROGRESS, rdf_file_finder.CollectFilesByKnownPathResult.Status.IN_PROGRESS, rdf_file_finder.CollectFilesByKnownPathResult.Status.COLLECTED ] self.assertCountEqual(expected_statuses, self._getResultsStatus(results))
def testFlowReportsErrorWhenCollectingUnexistingColumn(self): with temp.AutoTempFilePath() as temp_file_path: with io.open(temp_file_path, mode="w", encoding="utf-8") as fd: fd.write("Just sample text to put in the file.") table = f""" [ {{ "collect_column": "{temp_file_path}" }} ] """ with osquery_test_lib.FakeOsqueryiOutput(stdout=table, stderr=""): flow_id = self._InitializeFlow( file_collection_columns=["non_existent_column"], check_flow_errors=False) progress = flow_test_lib.GetFlowProgress(self.client_id, flow_id) self.assertEqual( "No such column 'non_existent_column' to collect files from.", progress.error_message)
def testFlowReportsErrorWhenCollectingRowsAboveLimit(self): with temp.AutoTempFilePath() as temp_file_path: with io.open(temp_file_path, mode="w", encoding="utf-8") as fd: fd.write("Just sample text to put in the file.") table = f""" [ {{ "collect_column": "{temp_file_path}"}}, {{ "collect_column": "{temp_file_path}"}} ] """ with mock.patch.object(osquery_flow, "FILE_COLLECTION_MAX_ROWS", 1): with osquery_test_lib.FakeOsqueryiOutput(stdout=table, stderr=""): flow_id = self._InitializeFlow( file_collection_columns=["collect_column"], check_flow_errors=False) progress = flow_test_lib.GetFlowProgress(self.client_id, flow_id) self.assertEqual( "Requested file collection on a table with 2 rows, " "but the limit is 1 rows.", progress.error_message)