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