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)
