def testMultiGetFileProgressReportsSkippedDuplicatesCorrectly(self):
client_mock = action_mocks.MultiGetFileClientMock()
image_path = os.path.join(self.base_path, "test_img.dd")
pathspec = rdf_paths.PathSpec(pathtype=rdf_paths.PathSpec.PathType.OS,
path=image_path)
args = transfer.MultiGetFileArgs(pathspecs=[pathspec])
# Let the flow run to make sure the file is collected.
flow_test_lib.TestFlowHelper(transfer.MultiGetFile.__name__,
client_mock,
token=self.token,
client_id=self.client_id,
args=args)
# Run the flow second time to make sure duplicates are collected.
flow_id = flow_test_lib.TestFlowHelper(transfer.MultiGetFile.__name__,
client_mock,
token=self.token,
client_id=self.client_id,
args=args)
f_obj = flow_test_lib.GetFlowObj(self.client_id, flow_id)
f_instance = transfer.MultiGetFile(f_obj)
p = f_instance.GetProgress()
self.assertEqual(p.num_collected, 0)
self.assertEqual(p.num_failed, 0)
self.assertEqual(p.num_skipped, 1)
self.assertLen(p.pathspecs_progress, 1)
self.assertEqual(p.pathspecs_progress[0].pathspec, pathspec)
self.assertEqual(p.pathspecs_progress[0].status,
transfer.PathSpecProgress.Status.SKIPPED)
def testPathInterpolation(self):
bar = rdf_client.User(username="******")
baz = rdf_client.User(username="******")
self.client_id = self.SetupClient(
0, system="foo", fqdn="norf", users=[bar, baz])
with temp.AutoTempDirPath(remove_non_empty=True) as temp_dirpath:
self._Touch(os.path.join(temp_dirpath, "foo", "bar"))
self._Touch(os.path.join(temp_dirpath, "foo", "baz"))
self._Touch(os.path.join(temp_dirpath, "foo", "quux"))
self._Touch(os.path.join(temp_dirpath, "thud", "norf", "plugh"))
self._Touch(os.path.join(temp_dirpath, "thud", "norf", "blargh"))
paths = [
os.path.join(temp_dirpath, "%%os%%", "%%users.username%%"),
os.path.join(temp_dirpath, "thud", "%%fqdn%%", "plugh"),
]
action = rdf_file_finder.FileFinderAction.Action.STAT
results, flow_id = self._RunCFF(paths, action)
result_paths = [result.stat_entry.pathspec.path for result in results]
self.assertCountEqual(result_paths, [
os.path.join(temp_dirpath, "foo", "bar"),
os.path.join(temp_dirpath, "foo", "baz"),
os.path.join(temp_dirpath, "thud", "norf", "plugh")
])
# Also check that the argument protobuf still has the original values.
flow_obj = flow_test_lib.GetFlowObj(self.client_id, flow_id)
args = flow_obj.args
self.assertCountEqual(args.paths, paths)
def testCorrectlyGeneratesArchiveMappings(self):
with mock.patch.object(collectors, "ArtifactCollectorFlow",
MockArtifactCollectorFlow):
flow_id, _, _ = self._RunCollectBrowserHistory(browsers=[
webhistory.Browser.CHROME,
webhistory.Browser.SAFARI,
])
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 = flow.GetFilesArchiveMappings(results)
self.assertCountEqual(mappings, [
flow_base.ClientPathArchiveMapping(
db.ClientPath.OS(self.client_id,
("home", "foo", "ChromeHistory")),
"chrome/ChromeHistory",
),
flow_base.ClientPathArchiveMapping(
db.ClientPath.OS(self.client_id,
("home", "foo", "SafariHistory")),
"safari/SafariHistory",
),
])
def testSendReplyWorksCorrectlyInIncrementalCallback(self):
flow_id = flow_test_lib.StartAndRunFlow(
FlowWithIncrementalCallback,
client_mock=NotSendingStatusClientMock(),
client_id=self.client_id,
# Set check_flow_errors to False, otherwise test runner will complain
# that the flow has finished in the RUNNING state.
check_flow_errors=False)
flow_obj = flow_test_lib.GetFlowObj(self.client_id, flow_id)
self.assertEqual(flow_obj.flow_state, flow_obj.FlowState.RUNNING)
results = flow_test_lib.GetFlowResults(self.client_id, flow_id)
self.assertListEqual(results, [
rdfvalue.RDFString(f"Hello World {i}") for i in range(
NotSendingStatusClientMock.NUM_INCREMENTAL_RESPONSES)
])
def testIncrementalCallbackIsNotCalledWhenStatusMessageArrivesEarly(
self, m):
# Mocks don't have names by default.
m.__name__ = "ReceiveHelloCallback"
flow_id = flow_test_lib.StartAndRunFlow(
FlowWithIncrementalCallback,
client_mock=StatusOnlyClientMock(),
client_id=self.client_id,
# Set check_flow_errors to False, otherwise test runner will complain
# that the flow has finished in the RUNNING state.
check_flow_errors=False)
flow_obj = flow_test_lib.GetFlowObj(self.client_id, flow_id)
self.assertEqual(flow_obj.flow_state, flow_obj.FlowState.RUNNING)
self.assertEqual(m.call_count, 0)
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 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 testIncrementalCallbackReturnsResultsBeforeStatus(self, m):
# Mocks don't have names by default.
m.__name__ = "ReceiveHelloCallback"
flow_id = flow_test_lib.StartAndRunFlow(
FlowWithIncrementalCallback,
client_mock=NotSendingStatusClientMock(),
client_id=self.client_id,
# Set check_flow_errors to False, otherwise test runner will complain
# that the flow has finished in the RUNNING state.
check_flow_errors=False)
flow_obj = flow_test_lib.GetFlowObj(self.client_id, flow_id)
self.assertEqual(flow_obj.flow_state, flow_obj.FlowState.RUNNING)
self.assertEqual(m.call_count,
NotSendingStatusClientMock.NUM_INCREMENTAL_RESPONSES)
for i in range(NotSendingStatusClientMock.NUM_INCREMENTAL_RESPONSES):
# Get the positional arguments of each call.
args = m.mock_calls[i][1]
# Compare the first positional argument ('responses') to the responses
# list that we expect to have been passed to the callback.
self.assertEqual(list(args[0]),
[rdfvalue.RDFString(f"Hello World {i}")])
def testMultiGetFileProgressReportsFailuresAndSuccessesCorrectly(self):
client_mock = action_mocks.MultiGetFileClientMock()
image_path = os.path.join(self.base_path, "test_img.dd")
pathspec_1 = rdf_paths.PathSpec(
pathtype=rdf_paths.PathSpec.PathType.OS, path=image_path)
pathspec_2 = rdf_paths.PathSpec(
pathtype=rdf_paths.PathSpec.PathType.OS, path="/non/existing/path")
args = transfer.MultiGetFileArgs(pathspecs=[
pathspec_1,
pathspec_2,
])
flow_id = flow_test_lib.TestFlowHelper(transfer.MultiGetFile.__name__,
client_mock,
token=self.token,
client_id=self.client_id,
args=args)
f_obj = flow_test_lib.GetFlowObj(self.client_id, flow_id)
f_instance = transfer.MultiGetFile(f_obj)
p = f_instance.GetProgress()
self.assertEqual(p.num_pending_hashes, 0)
self.assertEqual(p.num_pending_files, 0)
self.assertEqual(p.num_skipped, 0)
self.assertEqual(p.num_collected, 1)
self.assertEqual(p.num_failed, 1)
# Check that pathspecs in the progress proto are returned in the same order
# as in the args proto.
self.assertEqual(p.pathspecs_progress[0].pathspec, pathspec_1)
self.assertEqual(p.pathspecs_progress[1].pathspec, pathspec_2)
# Check that per-pathspecs statuses are correct.
self.assertEqual(p.pathspecs_progress[0].status,
transfer.PathSpecProgress.Status.COLLECTED)
self.assertEqual(p.pathspecs_progress[1].status,
transfer.PathSpecProgress.Status.FAILED)
