def testLimitPropagation(self): """This tests that client actions are limited properly.""" client_mock = action_mocks.CPULimitClientMock( user_cpu_usage=[10], system_cpu_usage=[10], network_usage=[1000], runtime_us=[rdfvalue.Duration.From(1, rdfvalue.SECONDS)]) flow_test_lib.StartAndRunFlow(flow_test_lib.CPULimitFlow, client_mock=client_mock, client_id=self.client_id, cpu_limit=1000, network_bytes_limit=10000, runtime_limit=rdfvalue.Duration.From( 5, rdfvalue.SECONDS)) self.assertEqual(client_mock.storage["cpulimit"], [1000, 980, 960]) self.assertEqual(client_mock.storage["networklimit"], [10000, 9000, 8000]) self.assertEqual(client_mock.storage["networklimit"], [10000, 9000, 8000]) self.assertEqual(client_mock.storage["runtimelimit"], [ rdfvalue.Duration.From(5, rdfvalue.SECONDS), rdfvalue.Duration.From(4, rdfvalue.SECONDS), rdfvalue.Duration.From(3, rdfvalue.SECONDS), ])
def _RunFlow(self, client_id): flow_args = transfer.GetFileArgs( pathspec=rdf_paths.PathSpec( path="/tmp/evil.txt", pathtype=rdf_paths.PathSpec.PathType.OS)) client_mock = hunt_test_lib.SampleHuntMock() if data_store.RelationalDBFlowsEnabled(): with test_lib.FakeTime(42): return flow_test_lib.StartAndRunFlow( transfer.GetFile, client_id=client_id, client_mock=client_mock, flow_args=flow_args) else: runner_args = rdf_flow_runner.FlowRunnerArgs( flow_name=transfer.GetFile.__name__) with test_lib.FakeTime(42): flow_urn = flow.StartAFF4Flow( client_id=client_id, args=flow_args, runner_args=runner_args, token=self.token) flow_test_lib.TestFlowHelper( flow_urn, client_mock=client_mock, client_id=client_id, token=self.token) return flow_urn.Basename()
def Run(self): client_id = self.SetupClient(0) failing_descriptor = rdf_output_plugin.OutputPluginDescriptor( plugin_name=hunt_test_lib.FailingDummyHuntOutputPlugin.__name__) with test_lib.FakeTime(42): if data_store.RelationalDBFlowsEnabled(): flow_id = flow_test_lib.StartAndRunFlow( flow_cls=flow_test_lib.DummyFlowWithSingleReply, client_id=client_id.Basename(), output_plugins=[failing_descriptor]) else: flow_urn = flow.StartAFF4Flow( flow_name=flow_test_lib.DummyFlowWithSingleReply.__name__, client_id=client_id, output_plugins=[failing_descriptor], token=self.token) flow_id = flow_urn.Basename() flow_test_lib.TestFlowHelper(flow_urn, token=self.token) self.Check( "ListFlowOutputPluginErrors", args=flow_plugin.ApiListFlowOutputPluginErrorsArgs( client_id=client_id.Basename(), flow_id=flow_id, plugin_id="FailingDummyHuntOutputPlugin_0"), replace={flow_id: "W:ABCDEF"})
def Run(self): client_id = self.SetupClient(0) email_descriptor = rdf_output_plugin.OutputPluginDescriptor( plugin_name=email_plugin.EmailOutputPlugin.__name__, plugin_args=email_plugin.EmailOutputPluginArgs( email_address="test@localhost", emails_limit=42)) with test_lib.FakeTime(42): if data_store.RelationalDBFlowsEnabled(): flow_id = flow_test_lib.StartAndRunFlow( flow_cls=flow_test_lib.DummyFlowWithSingleReply, client_id=client_id.Basename(), output_plugins=[email_descriptor]) else: flow_urn = flow.StartAFF4Flow( flow_name=flow_test_lib.DummyFlowWithSingleReply.__name__, client_id=client_id, output_plugins=[email_descriptor], token=self.token) flow_id = flow_urn.Basename() flow_test_lib.TestFlowHelper(flow_urn, token=self.token) self.Check( "ListFlowOutputPluginLogs", args=flow_plugin.ApiListFlowOutputPluginLogsArgs( client_id=client_id.Basename(), flow_id=flow_id, plugin_id="EmailOutputPlugin_0"), replace={flow_id: "W:ABCDEF"})
def testKeepAliveRunsSuccessfully(self): client_ids = self.SetupClients(1) client_mock = action_mocks.ActionMock(admin.Echo) flow_test_lib.StartAndRunFlow(administrative.KeepAlive, duration=rdfvalue.Duration("1s"), client_id=client_ids[0].Basename(), client_mock=client_mock)
def testCallState(self): """Test the ability to chain flows.""" CallStateFlow.success = False # Run the flow in the simulated way flow_test_lib.StartAndRunFlow(CallStateFlow, client_id=self.client_id) self.assertEqual(CallStateFlow.success, True)
def setUp(self): super().setUp() self.handler = flow_plugin.ApiListFlowResultsHandler() self.client_id = self.SetupClient(0) self.flow_id = flow_test_lib.StartAndRunFlow( DummyFlowWithTwoTaggedReplies, client_id=self.client_id)
def testChainedFlow(self): """Test the ability to chain flows.""" ParentFlow.success = False # Run the flow in the simulated way flow_test_lib.StartAndRunFlow( ParentFlow, client_mock=ClientMock(), client_id=self.client_id) self.assertEqual(ParentFlow.success, True)
def _RunBrowserHistoryFlow(self, **kwargs): flow_args = webhistory.BrowserHistoryFlowArgs(**kwargs) flow_id = flow_test_lib.StartAndRunFlow( webhistory.BrowserHistoryFlow, creator=self.token.username, client_mock=action_mocks.ActionMock(), client_id=self.client_id, flow_args=flow_args) return flow_test_lib.GetFlowResults(self.client_id, flow_id)
def testIncrementalCallbackIsCalledWhenAllResponsesArriveAtOnce(self): flow_id = flow_test_lib.StartAndRunFlow(FlowWithIncrementalCallback, client_mock=ClientMock(), client_id=self.client_id) results = flow_test_lib.GetFlowResults(self.client_id, flow_id) self.assertListEqual(results, [ rdfvalue.RDFString("Hello World"), rdfvalue.RDFString("Final: Hello World") ])
def _RunFlow(self, client_id): flow_args = transfer.GetFileArgs(pathspec=rdf_paths.PathSpec( path="/tmp/evil.txt", pathtype=rdf_paths.PathSpec.PathType.OS)) client_mock = hunt_test_lib.SampleHuntMock(failrate=2) with test_lib.FakeTime(42): return flow_test_lib.StartAndRunFlow(transfer.GetFile, client_id=client_id, client_mock=client_mock, flow_args=flow_args)
def testBrokenChainedFlow(self): BrokenParentFlow.success = False # Run the flow in the simulated way with test_lib.SuppressLogs(): flow_test_lib.StartAndRunFlow(BrokenParentFlow, client_mock=ClientMock(), client_id=self.client_id, check_flow_errors=False) self.assertEqual(BrokenParentFlow.success, True)
def testUserGetsNotificationWithNumberOfResults(self): username = "******" self.CreateUser(username) flow_test_lib.StartAndRunFlow( FlowWithMultipleResultTypes, client_id=self.client_id, creator=username) notifications = self.GetUserNotifications(username) self.assertIn("FlowWithMultipleResultTypes completed with 6 results", notifications[0].message)
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 testNestedFlowsHaveTheirResultsSaved(self): # Run the flow in the simulated way parent_flow_id = flow_test_lib.StartAndRunFlow( ParentFlow, client_mock=ClientMock(), client_id=self.client_id) child_flows = data_store.REL_DB.ReadChildFlowObjects( self.client_id, parent_flow_id) self.assertLen(child_flows, 1) child_flow_results = flow_test_lib.GetFlowResults( self.client_id, child_flows[0].flow_id) self.assertNotEmpty(child_flow_results)
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 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 testLimitPropagation(self): """This tests that client actions are limited properly.""" client_mock = action_mocks.CPULimitClientMock(user_cpu_usage=[10], system_cpu_usage=[10], network_usage=[1000]) flow_test_lib.StartAndRunFlow(flow_test_lib.CPULimitFlow, client_mock=client_mock, client_id=self.client_id, cpu_limit=1000, network_bytes_limit=10000) self.assertEqual(client_mock.storage["cpulimit"], [1000, 980, 960]) self.assertEqual(client_mock.storage["networklimit"], [10000, 9000, 8000])
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 testNetworkLimitExceeded(self): """This tests that the network limit for flows is working.""" client_mock = action_mocks.CPULimitClientMock(user_cpu_usage=[10], system_cpu_usage=[10], network_usage=[1000]) with test_lib.SuppressLogs(): flow_id = flow_test_lib.StartAndRunFlow(flow_test_lib.CPULimitFlow, client_mock=client_mock, client_id=self.client_id, cpu_limit=1000, network_bytes_limit=1500, check_flow_errors=False) rdf_flow = data_store.REL_DB.ReadFlowObject(self.client_id, flow_id) self.assertEqual(rdf_flow.flow_state, "ERROR") self.assertIn("bytes limit exceeded", rdf_flow.error_message)
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 Run(self): client_id = self.SetupClient(0) failing_descriptor = rdf_output_plugin.OutputPluginDescriptor( plugin_name=hunt_test_lib.FailingDummyHuntOutputPlugin.__name__) with test_lib.FakeTime(42): flow_id = flow_test_lib.StartAndRunFlow( flow_cls=flow_test_lib.DummyFlowWithSingleReply, client_id=client_id, output_plugins=[failing_descriptor]) self.Check("ListFlowOutputPluginErrors", args=flow_plugin.ApiListFlowOutputPluginErrorsArgs( client_id=client_id, flow_id=flow_id, plugin_id="FailingDummyHuntOutputPlugin_0"), replace={flow_id: "W:ABCDEF"})
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 testRuntimeLimitExceeded(self): client_mock = action_mocks.CPULimitClientMock( user_cpu_usage=[1], system_cpu_usage=[1], network_usage=[1], runtime_us=[rdfvalue.Duration.From(4, rdfvalue.SECONDS)]) with test_lib.SuppressLogs(): flow_id = flow_test_lib.StartAndRunFlow( flow_test_lib.CPULimitFlow, client_mock=client_mock, client_id=self.client_id, runtime_limit=rdfvalue.Duration.From(9, rdfvalue.SECONDS), check_flow_errors=False) rdf_flow = data_store.REL_DB.ReadFlowObject(self.client_id, flow_id) self.assertEqual(rdf_flow.flow_state, "ERROR") self.assertIn("Runtime limit exceeded", rdf_flow.error_message)
def Run(self): client_id = self.SetupClient(0) email_descriptor = rdf_output_plugin.OutputPluginDescriptor( plugin_name=email_plugin.EmailOutputPlugin.__name__, plugin_args=email_plugin.EmailOutputPluginArgs( email_address="test@localhost", emails_limit=42)) with test_lib.FakeTime(42): flow_id = flow_test_lib.StartAndRunFlow( flow_cls=flow_test_lib.DummyFlowWithSingleReply, client_id=client_id, output_plugins=[email_descriptor]) self.Check("ListFlowOutputPluginLogs", args=flow_plugin.ApiListFlowOutputPluginLogsArgs( client_id=client_id, flow_id=flow_id, plugin_id="EmailOutputPlugin_0"), replace={flow_id: "W:ABCDEF"})
def RunFlow(self, flow_cls=None, output_plugins=None, flow_args=None, client_mock=None): if flow_args is None: flow_args = transfer.GetFileArgs(pathspec=rdf_paths.PathSpec( path="/tmp/evil.txt", pathtype=rdf_paths.PathSpec.PathType.OS)) if client_mock is None: client_mock = hunt_test_lib.SampleHuntMock(failrate=2) flow_urn = flow_test_lib.StartAndRunFlow(flow_cls or transfer.GetFile, client_mock=client_mock, client_id=self.client_id, flow_args=flow_args, output_plugins=output_plugins) return flow_urn
def testCreatorPropagation(self): username = u"original user" data_store.REL_DB.WriteGRRUser(username) client_mock = ClientMock() flow_id = flow_test_lib.StartAndRunFlow(flow_cls=ParentFlow, client_id=self.client_id, creator=username, client_mock=client_mock) flow_obj = data_store.REL_DB.ReadFlowObject(self.client_id, flow_id) self.assertEqual(flow_obj.creator, username) child_flows = data_store.REL_DB.ReadChildFlowObjects( self.client_id, flow_id) self.assertLen(child_flows, 1) child_flow = child_flows[0] self.assertEqual(child_flow.creator, username)
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}")])