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)