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;")