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 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}",
),
])
