def testFlowReportsErrorWhenCollectingFileAboveSingleLimit(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_SINGLE_FILE_BYTES", less_than_necessary_bytes): with osquery_test_lib.FakeOsqueryiOutput(stdout=table, stderr=""): flow_id = self._InitializeFlow("Doesn't matter", ["collect_column"], check_flow_errors=False) progress = flow_test_lib.GetFlowProgress( self.client_id, flow_id) self.assertEqual( f"File with path {temp_file_path} is too big: " f"{target_bytes} bytes when the limit is " f"{less_than_necessary_bytes} bytes.", 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)
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 testFlowDoesntFailWhenCollectingFilesFromEmptyResult(self): empty_table = """ [ ] """ with osquery_test_lib.FakeOsqueryiOutput(stdout=empty_table, stderr=""): results = self._RunFlow(file_collection_columns=["collect_column"]) self.assertEmpty(results)
def testFlowDoesntFailWhenCollectingFilesFromEmptyResult(self): empty_table = """ [ ] """ with osquery_test_lib.FakeOsqueryiOutput(stdout=empty_table, stderr=""): results = self._RunFlow("Doesn't matter", ["collect_column"]) self.assertLen(results, 1) self.assertIsInstance(results[0], rdf_osquery.OsqueryResult)
def testFlowFailsWhenCollectingUnexistingColumn(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=""): with self.assertRaises(RuntimeError): self._RunFlow("Doesn't matter", ["non_existent_column"])
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 _RunOsqueryExportResults(self, stdout: str) -> utils.BinaryChunkIterator: client_id = self.SetupClient(0) with osquery_test_lib.FakeOsqueryiOutput(stdout=stdout, stderr=""): flow_id = flow_test_lib.TestFlowHelper( osquery_flow.OsqueryFlow.__name__, action_mocks.OsqueryClientMock(), client_id=client_id, token=self.token, query="doesn't matter") result_flow = self.api.Client(client_id=client_id).Flow(flow_id) result_flow.WaitUntilDone() format_csv = api_osquery_pb2.ApiGetOsqueryResultsArgs.Format.CSV return result_flow.GetOsqueryResults(format_csv)
def testFlowDoesntCollectWhenRowsAboveLimit(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=""): with self.assertRaises(RuntimeError): self._RunFlow(file_collection_columns=["collect_column"])
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 testOutput(self): stdout = """ [ { "foo": "bar", "quux": "norf" }, { "foo": "baz", "quux": "thud" } ] """ with osquery_test_lib.FakeOsqueryiOutput(stdout=stdout, stderr=""): results = _Query("SELECT foo, quux FROM blargh;") self.assertLen(results, 1) table = results[0].table self.assertLen(table.header.columns, 2) self.assertEqual(table.header.columns[0].name, "foo") self.assertEqual(table.header.columns[1].name, "quux") self.assertEqual(list(table.Column("foo")), ["bar", "baz"]) self.assertEqual(list(table.Column("quux")), ["norf", "thud"])
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 testFlowDoesntCollectSingleFileAboveTotalLimit(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=""): with self.assertRaises(RuntimeError): self._RunFlow(file_collection_columns=["collect_column"])
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 testOsquery(self): data_store.REL_DB.WriteClientMetadata( client_id=ClientTest.FAKE_CLIENT_ID, fleetspeak_enabled=False) client = grr_colab.Client.with_id(ClientTest.FAKE_CLIENT_ID) stdout = """ [ { "foo": "test1", "bar": "test2" }, { "foo": "test3", "bar": "test4" } ] """ with osquery_test_lib.FakeOsqueryiOutput(stdout=stdout, stderr=''): table = client.osquery('SELECT foo, bar FROM table;') self.assertLen(table.header.columns, 2) self.assertEqual(table.header.columns[0].name, 'foo') self.assertEqual(table.header.columns[1].name, 'bar') self.assertEqual(list(list(table.rows)[0].values), ['test1', 'test2']) self.assertEqual(list(list(table.rows)[1].values), ['test3', 'test4'])
def testFailure(self): stderr = "Error: query syntax error" with osquery_test_lib.FakeOsqueryiOutput(stdout="", stderr=stderr): flow_id = flow_test_lib.StartFlow( flow_cls=osquery_flow.OsqueryFlow, client_id=self.client_id, creator=self.test_username, query="<<<FAKE OSQUERY FLOW QUERY PLACEHOLDER>>>") with self.assertRaises(RuntimeError): flow_test_lib.RunFlow( client_id=self.client_id, flow_id=flow_id, client_mock=action_mocks.OsqueryClientMock()) flow = data_store.REL_DB.ReadFlowObject(self.client_id, flow_id) self.assertEqual(flow.flow_state, rdf_flow_objects.Flow.FlowState.ERROR) self.assertIn(stderr, flow.error_message)
def testSuccess(self): stdout = """ [ { "foo": "quux", "bar": "norf", "baz": "thud" }, { "foo": "blargh", "bar": "plugh", "baz": "ztesch" } ] """ with osquery_test_lib.FakeOsqueryiOutput(stdout=stdout, stderr=""): results = self._RunFlow() self.assertLen(results, 1) table = results[0].table self.assertLen(table.header.columns, 3) self.assertEqual(table.header.columns[0].name, "foo") self.assertEqual(table.header.columns[1].name, "bar") self.assertEqual(table.header.columns[2].name, "baz") self.assertEqual(list(table.Column("foo")), ["quux", "blargh"]) self.assertEqual(list(table.Column("bar")), ["norf", "plugh"]) self.assertEqual(list(table.Column("baz")), ["thud", "ztesch"])
def testArchiveMappingsForDuplicateFilesInResult(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=["collect_column"]) flow = flow_base.FlowBase.CreateFlowInstance( flow_test_lib.GetFlowObj(self.client_id, flow_id)) results = list(flow_test_lib.GetRawFlowResults(self.client_id, flow_id)) # This is how we emulate duplicate filenames in the results duplicated_results = results + results + results mappings = list(flow.GetFilesArchiveMappings(iter(duplicated_results))) self.assertCountEqual(mappings, [ flow_base.ClientPathArchiveMapping( db.ClientPath.OS(self.client_id, temp_file_path.split("/")[1:]), f"osquery_collected_files{temp_file_path}", ), flow_base.ClientPathArchiveMapping( db.ClientPath.OS(self.client_id, temp_file_path.split("/")[1:]), f"osquery_collected_files{temp_file_path}-1", ), flow_base.ClientPathArchiveMapping( db.ClientPath.OS(self.client_id, temp_file_path.split("/")[1:]), f"osquery_collected_files{temp_file_path}-2", ), ])
def testIgnoreStderrErrors(self): query = "SELECT foo, bar, baz FROM blargh;" stdout = """ [ { "foo": "quux", "bar": "norf", "baz": "thud" } ] """ stderr = "Warning: near 'FROM': table 'blargh' might be very large" with osquery_test_lib.FakeOsqueryiOutput(stdout=stdout, stderr=stderr): results = _Query(query, ignore_stderr_errors=True) self.assertLen(results, 1) self.assertEqual(results[0].stderr, stderr) table = results[0].table self.assertLen(table.header.columns, 3) self.assertEqual(table.header.columns[0].name, "foo") self.assertEqual(table.header.columns[1].name, "bar") self.assertEqual(table.header.columns[2].name, "baz") self.assertEqual(list(table.Column("foo")), ["quux"]) self.assertEqual(list(table.Column("bar")), ["norf"]) self.assertEqual(list(table.Column("baz")), ["thud"])
def testArchiveMappingsForMultipleFiles(self): with temp.AutoTempDirPath(remove_non_empty=True) as temp_dir_path: temp_file_path1 = os.path.join(temp_dir_path, "foo") temp_file_path2 = os.path.join(temp_dir_path, "bar") with io.open(temp_file_path1, mode="w", encoding="utf-8") as fd: fd.write("Just sample text to put in the file 1.") with io.open(temp_file_path2, mode="w", encoding="utf-8") as fd: fd.write("Just sample text to put in the file 2.") table = f""" [ {{ "collect_column": "{temp_file_path1}" }}, {{ "collect_column": "{temp_file_path2}" }} ] """ with osquery_test_lib.FakeOsqueryiOutput(stdout=table, stderr=""): flow_id = self._InitializeFlow( file_collection_columns=["collect_column"]) flow = flow_base.FlowBase.CreateFlowInstance( flow_test_lib.GetFlowObj(self.client_id, flow_id)) results = flow_test_lib.GetRawFlowResults(self.client_id, flow_id) mappings = list(flow.GetFilesArchiveMappings(iter(results))) self.assertCountEqual(mappings, [ flow_base.ClientPathArchiveMapping( db.ClientPath.OS(self.client_id, temp_file_path1.split("/")[1:]), f"osquery_collected_files{temp_file_path1}", ), flow_base.ClientPathArchiveMapping( db.ClientPath.OS(self.client_id, temp_file_path2.split("/")[1:]), f"osquery_collected_files{temp_file_path2}", ), ])
def testFlowCollectFile(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=""): results = self._RunFlow("Doesn't matter", ["collect_column"]) self.assertLen(results, 2) self.assertIsInstance(results[0], rdf_osquery.OsqueryResult) self.assertIsInstance(results[1], rdf_client_fs.StatEntry) pathspec = results[1].pathspec client_path = db.ClientPath.FromPathSpec(self.client_id, pathspec) fd_rel_db = file_store.OpenFile(client_path) file_text = fd_rel_db.read().decode("utf-8") self.assertEqual(file_text, "Just sample text to put in the file.")
def testFailure(self): stderr = "Error: near '*': syntax error" with osquery_test_lib.FakeOsqueryiOutput(stdout="", stderr=stderr): with self.assertRaises(RuntimeError): self._RunQuery("SELECT * FROM *;")
def testSkipEmptyTable(self): with osquery_test_lib.FakeOsqueryiOutput(stdout="[]", stderr=""): results = self._RunFlow() self.assertEmpty(results)
def testError(self): stderr = "Error: near 'FROM': syntax error" with osquery_test_lib.FakeOsqueryiOutput(stdout="", stderr=stderr): with self.assertRaises(osquery.QueryError): _Query("FROM bar SELECT foo;")