def testHuntIsStoppedIfAveragePerClientCpuUsageTooHigh(self): client_ids = self.SetupClients(5) hunt_id = self._CreateHunt( client_rule_set=foreman_rules.ForemanClientRuleSet(), client_rate=0, avg_cpu_seconds_per_client_limit=3, args=self.GetFileHuntArgs()) with utils.Stubber(hunt, "MIN_CLIENTS_FOR_AVERAGE_THRESHOLDS", 4): def CheckState(hunt_state, user_cpu_time, system_cpu_time): hunt_obj = data_store.REL_DB.ReadHuntObject(hunt_id) self.assertEqual(hunt_obj.hunt_state, hunt_state) self.assertEqual( hunt_obj.client_resources_stats.user_cpu_stats.sum, user_cpu_time) self.assertEqual( hunt_obj.client_resources_stats.system_cpu_stats.sum, system_cpu_time) self._RunHunt(client_ids[:2], client_mock=hunt_test_lib.SampleHuntMock( user_cpu_time=1, system_cpu_time=2, failrate=-1)) # Hunt should still be running: we need at least 3 clients to start # calculating the average. CheckState(rdf_hunt_objects.Hunt.HuntState.STARTED, 2, 4) self._RunHunt([client_ids[2]], client_mock=hunt_test_lib.SampleHuntMock( user_cpu_time=2, system_cpu_time=4, failrate=-1)) # Hunt should still be running: even though the average is higher than the # limit, number of clients is not enough. CheckState(rdf_hunt_objects.Hunt.HuntState.STARTED, 4, 8) self._RunHunt([client_ids[3]], client_mock=hunt_test_lib.SampleHuntMock( user_cpu_time=0, system_cpu_time=0, failrate=-1)) # Hunt should still be running: we got 4 clients, which is enough to check # average per-client CPU usage. But 4 user cpu + 8 system cpu seconds for # 4 clients make an average of 3 seconds per client - this is within the # limit. CheckState(rdf_hunt_objects.Hunt.HuntState.STARTED, 4, 8) self._RunHunt([client_ids[4]], client_mock=hunt_test_lib.SampleHuntMock( user_cpu_time=2, system_cpu_time=4, failrate=-1)) # Hunt should be terminated: the average is exceeded. CheckState(rdf_hunt_objects.Hunt.HuntState.STOPPED, 6, 12) self._CheckHuntStoppedNotification( "reached the average CPU seconds per client")
def testHuntIsStoppedIfAveragePerClientNetworkUsageTooHigh(self): client_ids = self.SetupClients(5) hunt_id = self._CreateHunt( client_rule_set=foreman_rules.ForemanClientRuleSet(), client_rate=0, avg_network_bytes_per_client_limit=1, args=self.GetFileHuntArgs()) with utils.Stubber(hunt, "MIN_CLIENTS_FOR_AVERAGE_THRESHOLDS", 4): def CheckState(hunt_state, network_bytes_sent): hunt_obj = data_store.REL_DB.ReadHuntObject(hunt_id) self.assertEqual(hunt_obj.hunt_state, hunt_state) hunt_counters = data_store.REL_DB.ReadHuntCounters(hunt_id) self.assertEqual(hunt_counters.total_network_bytes_sent, network_bytes_sent) self._RunHunt( client_ids[:2], client_mock=hunt_test_lib.SampleHuntMock( network_bytes_sent=1, failrate=-1)) # Hunt should still be running: we need at least 3 clients to start # calculating the average. CheckState(rdf_hunt_objects.Hunt.HuntState.STARTED, 2) self._RunHunt([client_ids[2]], client_mock=hunt_test_lib.SampleHuntMock( network_bytes_sent=2, failrate=-1)) # Hunt should still be running: even though the average is higher than the # limit, number of clients is not enough. CheckState(rdf_hunt_objects.Hunt.HuntState.STARTED, 4) self._RunHunt([client_ids[3]], client_mock=hunt_test_lib.SampleHuntMock( network_bytes_sent=0, failrate=-1)) # Hunt should still be running: we got 4 clients, which is enough to check # average per-client network bytes usage, but 4 bytes for 4 clients is # within the limit of 1 byte per client on average. CheckState(rdf_hunt_objects.Hunt.HuntState.STARTED, 4) self._RunHunt([client_ids[4]], client_mock=hunt_test_lib.SampleHuntMock( network_bytes_sent=2, failrate=-1)) # Hunt should be terminated: the limit is exceeded. CheckState(rdf_hunt_objects.Hunt.HuntState.STOPPED, 6) self._CheckHuntStoppedNotification( "reached the average network bytes per client")
def testOutputPluginsAreCorrectlyAppliedAndTheirStatusCanBeRead(self): hunt_test_lib.StatefulDummyHuntOutputPlugin.data = [] hunt_test_lib.DummyHuntOutputPlugin.num_calls = 0 hunt_test_lib.DummyHuntOutputPlugin.num_responses = 0 plugin_descriptor = rdf_output_plugin.OutputPluginDescriptor( plugin_name="DummyHuntOutputPlugin") hunt_id, _ = self._CreateAndRunHunt( num_clients=5, client_mock=hunt_test_lib.SampleHuntMock(failrate=-1), client_rule_set=foreman_rules.ForemanClientRuleSet(), client_rate=0, args=self.GetFileHuntArgs(), output_plugins=[plugin_descriptor]) self.assertEqual(hunt_test_lib.DummyHuntOutputPlugin.num_calls, 5) self.assertEqual(hunt_test_lib.DummyHuntOutputPlugin.num_responses, 5) logs = hunt.GetHuntOutputPluginLogs(hunt_id, 0, sys.maxsize) self.assertLen(logs, 5) for l in logs: self.assertEqual(l.batch_size, 1) self.assertEqual( l.status, output_plugin.OutputPluginBatchProcessingStatus.Status.SUCCESS) self.assertEqual(l.plugin_descriptor, plugin_descriptor)
def testCreatorPropagation(self): self.CreateAdminUser("adminuser") admin_token = access_control.ACLToken( username="******", reason="testing") # Start a flow that requires admin privileges in the hunt. The # parameters are not valid so the flow will error out but it's # enough to check if the flow was actually run (i.e., it passed # the label test). with implementation.StartHunt( hunt_name=standard.GenericHunt.__name__, flow_runner_args=rdf_flow_runner.FlowRunnerArgs( flow_name=administrative.UpdateClient.__name__), flow_args=administrative.UpdateClientArgs(), client_rule_set=self._CreateForemanClientRuleSet(), client_rate=0, token=admin_token) as hunt: hunt.Run() self.CreateUser("nonadmin") nonadmin_token = access_control.ACLToken( username="******", reason="testing") self.AssignTasksToClients() client_mock = hunt_test_lib.SampleHuntMock() hunt_test_lib.TestHuntHelper(client_mock, self.client_ids, False, nonadmin_token) errors = list(hunt.GetClientsErrors()) # Make sure there are errors... self.assertTrue(errors) # but they are not UnauthorizedAccess. for e in errors: self.assertNotIn("UnauthorizedAccess", e.backtrace)
def testOutputPluginsAreCorrectlyAppliedAndTheirStatusCanBeRead(self): hunt_test_lib.StatefulDummyHuntOutputPlugin.data = [] hunt_test_lib.DummyHuntOutputPlugin.num_calls = 0 hunt_test_lib.DummyHuntOutputPlugin.num_responses = 0 plugin_descriptor = rdf_output_plugin.OutputPluginDescriptor( plugin_name="DummyHuntOutputPlugin") hunt_id, client_ids = self._CreateAndRunHunt( num_clients=5, client_mock=hunt_test_lib.SampleHuntMock(failrate=-1), client_rule_set=foreman_rules.ForemanClientRuleSet(), client_rate=0, args=self.GetFileHuntArgs(), output_plugins=[plugin_descriptor]) self.assertEqual(hunt_test_lib.DummyHuntOutputPlugin.num_calls, 5) self.assertEqual(hunt_test_lib.DummyHuntOutputPlugin.num_responses, 5) logs = data_store.REL_DB.ReadHuntOutputPluginLogEntries( hunt_id, # REL_DB code uses strings for output plugin ids for consistency (as # all other DB ids are strings). At the moment plugin_id in the database # is simply an index of the plugin in Flow/Hunt.output_plugins list. output_plugin_id="0", offset=0, count=sys.maxsize, with_type=rdf_flow_objects.FlowOutputPluginLogEntry.LogEntryType. LOG) self.assertLen(logs, 5) self.assertCountEqual([l.client_id for l in logs], client_ids) for l in logs: self.assertEqual(l.hunt_id, hunt_id) self.assertGreater(l.timestamp, 0) self.assertEqual(l.message, "Processed 1 replies.")
def testHuntNotifications(self): """This tests the Hunt notification event.""" TestHuntListener.received_events = [] # Set up 10 clients. client_ids = self.SetupClients(10) client_rule_set = rdf_foreman.ForemanClientRuleSet(rules=[ rdf_foreman.ForemanClientRule( rule_type=rdf_foreman.ForemanClientRule.Type.REGEX, regex=rdf_foreman.ForemanRegexClientRule( attribute_name="GRR client", attribute_regex="GRR")) ]) with implementation.GRRHunt.StartHunt( hunt_name=BrokenSampleHunt.__name__, client_rule_set=client_rule_set, client_rate=0, notification_event="TestHuntDone", token=self.token) as hunt: hunt.GetRunner().Start() foreman = aff4.FACTORY.Open("aff4:/foreman", mode="rw", token=self.token) for client_id in client_ids: foreman.AssignTasksToClient(client_id) # Run the hunt. client_mock = hunt_test_lib.SampleHuntMock() hunt_test_lib.TestHuntHelper( client_mock, client_ids, check_flow_errors=False, token=self.token) self.assertEqual(len(TestHuntListener.received_events), 5)
def testFailingOutputPluginDoesNotAffectOtherOutputPlugins(self): failing_plugin_descriptor = rdf_output_plugin.OutputPluginDescriptor( plugin_name="FailingDummyHuntOutputPlugin") plugin_descriptor = rdf_output_plugin.OutputPluginDescriptor( plugin_name="DummyHuntOutputPlugin") hunt_id, _ = self._CreateAndRunHunt( num_clients=5, client_mock=hunt_test_lib.SampleHuntMock(failrate=-1), client_rule_set=foreman_rules.ForemanClientRuleSet(), client_rate=0, args=self.GetFileHuntArgs(), output_plugins=[failing_plugin_descriptor, plugin_descriptor]) errors = data_store.REL_DB.ReadHuntOutputPluginLogEntries( hunt_id, output_plugin_id="0", offset=0, count=sys.maxsize, with_type=rdf_flow_objects.FlowOutputPluginLogEntry.LogEntryType. ERROR) self.assertLen(errors, 5) # Check that non-failing output plugin is still correctly processed. logs = data_store.REL_DB.ReadHuntOutputPluginLogEntries( hunt_id, output_plugin_id="1", offset=0, count=sys.maxsize, with_type=rdf_flow_objects.FlowOutputPluginLogEntry.LogEntryType. LOG) self.assertLen(logs, 5)
def testOutputPluginsErrorsAreCorrectlyWrittenAndCanBeRead(self): failing_plugin_descriptor = rdf_output_plugin.OutputPluginDescriptor( plugin_name="FailingDummyHuntOutputPlugin") hunt_id, client_ids = self._CreateAndRunHunt( num_clients=5, client_mock=hunt_test_lib.SampleHuntMock(failrate=-1), client_rule_set=foreman_rules.ForemanClientRuleSet(), client_rate=0, args=self.GetFileHuntArgs(), output_plugins=[failing_plugin_descriptor]) errors = data_store.REL_DB.ReadHuntOutputPluginLogEntries( hunt_id, # REL_DB code uses strings for output plugin ids for consistency (as # all other DB ids are strings). At the moment plugin_id in the database # is simply an index of the plugin in Flow/Hunt.output_plugins list. output_plugin_id="0", offset=0, count=sys.maxsize, with_type=rdf_flow_objects.FlowOutputPluginLogEntry.LogEntryType. ERROR) self.assertLen(errors, 5) self.assertCountEqual([e.client_id for e in errors], client_ids) for e in errors: self.assertEqual(e.hunt_id, hunt_id) self.assertGreater(e.timestamp, 0) self.assertEqual(e.message, "Error while processing 1 replies: Oh no!")
def Run(self): client_ids = self.SetupClients(10) client_mock = hunt_test_lib.SampleHuntMock(failrate=2) hunt_id = self.CreateHunt(description="the hunt", creator=self.token.username) hunt.StartHunt(hunt_id) time_offset = 0 for client_id in client_ids: with test_lib.FakeTime(45 + time_offset): self.AssignTasksToClients([client_id]) hunt_test_lib.TestHuntHelper(client_mock, [client_id], self.token) time_offset += 10 replace = {hunt_id: "H:123456"} self.Check("GetHuntClientCompletionStats", args=hunt_plugin.ApiGetHuntClientCompletionStatsArgs( hunt_id=hunt_id), replace=replace) self.Check("GetHuntClientCompletionStats", args=hunt_plugin.ApiGetHuntClientCompletionStatsArgs( hunt_id=hunt_id, size=4), replace=replace) self.Check("GetHuntClientCompletionStats", args=hunt_plugin.ApiGetHuntClientCompletionStatsArgs( hunt_id=hunt_id, size=1000), replace=replace)
def SetupTestHuntView(self, client_limit=0, client_count=10): # Create some clients and a hunt to view. with self.CreateSampleHunt(client_limit=client_limit, client_count=client_count) as hunt: hunt.Log("TestLogLine") # Log an error just with some random traceback. hunt.LogClientError(self.client_ids[1], "Client Error 1", traceback.format_exc()) # Run the hunt. client_mock = hunt_test_lib.SampleHuntMock() hunt_test_lib.TestHuntHelper(client_mock, self.client_ids, False, self.token) hunt = aff4.FACTORY.Open(hunt.urn, token=self.token) all_count, _, _ = hunt.GetClientsCounts() if client_limit == 0: # No limit, so we should have all the clients self.assertEqual(all_count, client_count) else: self.assertEqual(all_count, min(client_count, client_limit)) return hunt
def testClientsTabShowsCompletedAndOutstandingClients(self): # Create some clients and a hunt to view. self.CreateSampleHunt() # Run the hunt on half the clients. finished_client_ids = self.client_ids[5:] outstanding_client_ids = self.client_ids[:5] client_mock = hunt_test_lib.SampleHuntMock() hunt_test_lib.TestHuntHelper(client_mock, finished_client_ids, False, self.token) self.Open("/#main=ManageHunts") self.Click("css=td:contains('GenericHunt')") self.Click("css=li[heading=Clients]") self.Click("css=label[name=ShowCompletedClients]") for client_id in finished_client_ids: self.WaitUntilContains(client_id.Basename(), self.GetText, "css=.tab-content") self.Click("css=label[name=ShowOutstandingClients]") for client_id in outstanding_client_ids: self.WaitUntilContains(client_id.Basename(), self.GetText, "css=.tab-content")
def testHuntIsStoppedWhenExpirationTimeIsReached(self): client_ids = self.SetupClients(5) duration = rdfvalue.Duration.From(1, rdfvalue.DAYS) expiry_time = rdfvalue.RDFDatetime.Now() + duration hunt_id = self._CreateHunt( client_rule_set=foreman_rules.ForemanClientRuleSet(), client_rate=0, duration=duration, args=self.GetFileHuntArgs()) client_mock = hunt_test_lib.SampleHuntMock(failrate=-1) foreman_obj = foreman.Foreman() for client_id in client_ids: foreman_obj.AssignTasksToClient(client_id) hunt_test_lib.TestHuntHelper(client_mock, client_ids[:3]) hunt_obj = data_store.REL_DB.ReadHuntObject(hunt_id) self.assertEqual(hunt_obj.hunt_state, rdf_hunt_objects.Hunt.HuntState.STARTED) with test_lib.FakeTime(expiry_time - rdfvalue.Duration.From(1, rdfvalue.SECONDS)): hunt_test_lib.TestHuntHelper(client_mock, client_ids[3:4]) hunt_obj = data_store.REL_DB.ReadHuntObject(hunt_id) self.assertEqual(hunt_obj.hunt_state, rdf_hunt_objects.Hunt.HuntState.STARTED) with test_lib.FakeTime(expiry_time + rdfvalue.Duration.From(1, rdfvalue.SECONDS)): hunt_test_lib.TestHuntHelper(client_mock, client_ids[4:5]) hunt_obj = data_store.REL_DB.ReadHuntObject(hunt_id) self.assertEqual(hunt_obj.hunt_state, rdf_hunt_objects.Hunt.HuntState.COMPLETED)
def testHuntFlowLogsAreCorrectlyWrittenAndCanBeRead(self): hunt_args = rdf_hunt_objects.HuntArguments.Standard( flow_name=compatibility.GetName(flow_test_lib.DummyLogFlow)) hunt_id, client_ids = self._CreateAndRunHunt( num_clients=10, client_mock=hunt_test_lib.SampleHuntMock(failrate=-1), client_rule_set=foreman_rules.ForemanClientRuleSet(), client_rate=0, args=hunt_args) hunt_logs = data_store.REL_DB.ReadHuntLogEntries( hunt_id, 0, sys.maxsize) # 4 logs for each flow. Note: DummyLogFlow also calls DummyLogFlowChild, # but children flows logs should not be present in the output. self.assertLen(hunt_logs, 4 * len(client_ids)) self.assertCountEqual(set(log.client_id for log in hunt_logs), client_ids) messages_set = set(log.message for log in hunt_logs) self.assertCountEqual(messages_set, ["First", "Second", "Third", "Fourth"]) for nested_flow_log in ["Uno", "Dos", "Tres", "Cuatro"]: self.assertNotIn(nested_flow_log, messages_set) for log in hunt_logs: self.assertEqual(log.hunt_id, hunt_id)
def testResourceUsageStatsAreReportedCorrectly(self): hunt_id, _ = self._CreateAndRunHunt( num_clients=10, client_mock=hunt_test_lib.SampleHuntMock(failrate=-1), client_rule_set=foreman_rules.ForemanClientRuleSet(), client_rate=0, args=self.GetFileHuntArgs()) usage_stats = data_store.REL_DB.ReadHuntClientResourcesStats(hunt_id) # Values below are calculated based on SampleHuntMock's behavior. self.assertEqual(usage_stats.user_cpu_stats.num, 10) self.assertAlmostEqual(usage_stats.user_cpu_stats.mean, 5.5) self.assertAlmostEqual(usage_stats.user_cpu_stats.std, 2.8722813) self.assertEqual(usage_stats.system_cpu_stats.num, 10) self.assertAlmostEqual(usage_stats.system_cpu_stats.mean, 11) self.assertAlmostEqual(usage_stats.system_cpu_stats.std, 5.7445626) self.assertEqual(usage_stats.network_bytes_sent_stats.num, 10) self.assertAlmostEqual(usage_stats.network_bytes_sent_stats.mean, 16.5) self.assertAlmostEqual(usage_stats.network_bytes_sent_stats.std, 8.61684396) # NOTE: Not checking histograms here. RunningStatsTest tests that mean, # standard deviation and histograms are calculated correctly. Therefore # if mean/stdev values are correct histograms should be ok as well. self.assertLen(usage_stats.worst_performers, 10) prev = usage_stats.worst_performers[0] for p in usage_stats.worst_performers[1:]: self.assertGreater( prev.cpu_usage.user_cpu_time + prev.cpu_usage.system_cpu_time, p.cpu_usage.user_cpu_time + p.cpu_usage.system_cpu_time) prev = p
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) 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 testOutputPluginFlushErrorIsLoggedProperly(self): plugin_descriptor = rdf_output_plugin.OutputPluginDescriptor( plugin_name="FailingInFlushDummyHuntOutputPlugin") hunt_id, client_ids = self._CreateAndRunHunt( num_clients=5, client_mock=hunt_test_lib.SampleHuntMock(failrate=-1), client_rule_set=foreman_rules.ForemanClientRuleSet(), client_rate=0, args=self.GetFileHuntArgs(), output_plugins=[plugin_descriptor]) logs = data_store.REL_DB.ReadHuntOutputPluginLogEntries( hunt_id, output_plugin_id="0", offset=0, count=sys.maxsize, with_type=rdf_flow_objects.FlowOutputPluginLogEntry.LogEntryType. LOG) self.assertEmpty(logs) errors = data_store.REL_DB.ReadHuntOutputPluginLogEntries( hunt_id, output_plugin_id="0", offset=0, count=sys.maxsize, with_type=rdf_flow_objects.FlowOutputPluginLogEntry.LogEntryType. ERROR) self.assertLen(errors, 5) self.assertCountEqual([e.client_id for e in errors], client_ids) for e in errors: self.assertEqual(e.hunt_id, hunt_id) self.assertGreater(e.timestamp, 0) self.assertEqual( e.message, "Error while processing 1 replies: Flush, oh no!")
def RunFlow(self, flow_name=None, plugins=None, flow_args=None, client_mock=None): runner_args = rdf_flows.FlowRunnerArgs(flow_name=flow_name or transfer.GetFile.__name__, output_plugins=plugins) 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() flow_urn = flow.GRRFlow.StartFlow(client_id=self.client_id, args=flow_args, runner_args=runner_args, token=self.token) for _ in flow_test_lib.TestFlowHelper(flow_urn, client_mock=client_mock, client_id=self.client_id, token=self.token): pass return flow_urn
def Run(self): client_ids = self.SetupClients(10) client_mock = hunt_test_lib.SampleHuntMock() with test_lib.FakeTime(42): with self.CreateHunt(description="the hunt") as hunt_obj: hunt_obj.Run() time_offset = 0 for client_id in client_ids: with test_lib.FakeTime(45 + time_offset): self.AssignTasksToClients([client_id]) hunt_test_lib.TestHuntHelper(client_mock, [client_id], False, self.token) time_offset += 10 replace = {hunt_obj.urn.Basename(): "H:123456"} self.Check("GetHuntClientCompletionStats", args=hunt_plugin.ApiGetHuntClientCompletionStatsArgs( hunt_id=hunt_obj.urn.Basename()), replace=replace) self.Check("GetHuntClientCompletionStats", args=hunt_plugin.ApiGetHuntClientCompletionStatsArgs( hunt_id=hunt_obj.urn.Basename(), size=4), replace=replace) self.Check("GetHuntClientCompletionStats", args=hunt_plugin.ApiGetHuntClientCompletionStatsArgs( hunt_id=hunt_obj.urn.Basename(), size=1000), replace=replace)
def Run(self): runner_args = rdf_flows.FlowRunnerArgs( flow_name=transfer.GetFile.__name__) flow_args = transfer.GetFileArgs(pathspec=rdf_paths.PathSpec( path="/tmp/evil.txt", pathtype=rdf_paths.PathSpec.PathType.OS)) client_mock = hunt_test_lib.SampleHuntMock() with test_lib.FakeTime(42): flow_urn = flow.GRRFlow.StartFlow(client_id=self.client_id, args=flow_args, runner_args=runner_args, token=self.token) for _ in flow_test_lib.TestFlowHelper(flow_urn, client_mock=client_mock, client_id=self.client_id, token=self.token): pass self.Check("ListFlowResults", args=flow_plugin.ApiListFlowResultsArgs( client_id=self.client_id.Basename(), flow_id=flow_urn.Basename()), replace={flow_urn.Basename(): "W:ABCDEF"})
def testUpdatesStatsCounterOnOutputPluginFailure(self): plugin_descriptor = rdf_output_plugin.OutputPluginDescriptor( plugin_name="FailingDummyHuntOutputPlugin") prev_success_count = stats_collector_instance.Get().GetMetricValue( "hunt_results_ran_through_plugin", fields=["FailingDummyHuntOutputPlugin"]) prev_errors_count = stats_collector_instance.Get().GetMetricValue( "hunt_output_plugin_errors", fields=["FailingDummyHuntOutputPlugin"]) self._CreateAndRunHunt( num_clients=5, client_mock=hunt_test_lib.SampleHuntMock(failrate=-1), client_rule_set=foreman_rules.ForemanClientRuleSet(), client_rate=0, args=self.GetFileHuntArgs(), output_plugins=[plugin_descriptor]) success_count = stats_collector_instance.Get().GetMetricValue( "hunt_results_ran_through_plugin", fields=["FailingDummyHuntOutputPlugin"]) errors_count = stats_collector_instance.Get().GetMetricValue( "hunt_output_plugin_errors", fields=["FailingDummyHuntOutputPlugin"]) # 1 error for each client makes it 5 errors, 0 results. self.assertEqual(success_count - prev_success_count, 0) self.assertEqual(errors_count - prev_errors_count, 5)
def Run(self): if data_store.RelationalDBReadEnabled(): clients = self.SetupTestClientObjects(10) client_ids = sorted(clients) else: client_ids = [urn.Basename() for urn in self.SetupClients(10)] client_mock = hunt_test_lib.SampleHuntMock() with test_lib.FakeTime(42): with self.CreateHunt(description="the hunt") as hunt_obj: hunt_obj.Run() time_offset = 0 for client_id in client_ids: with test_lib.FakeTime(45 + time_offset): self.AssignTasksToClients([client_id]) hunt_test_lib.TestHuntHelper(client_mock, [rdf_client.ClientURN(client_id)], False, self.token) time_offset += 10 replace = {hunt_obj.urn.Basename(): "H:123456"} self.Check("GetHuntClientCompletionStats", args=hunt_plugin.ApiGetHuntClientCompletionStatsArgs( hunt_id=hunt_obj.urn.Basename()), replace=replace) self.Check("GetHuntClientCompletionStats", args=hunt_plugin.ApiGetHuntClientCompletionStatsArgs( hunt_id=hunt_obj.urn.Basename(), size=4), replace=replace) self.Check("GetHuntClientCompletionStats", args=hunt_plugin.ApiGetHuntClientCompletionStatsArgs( hunt_id=hunt_obj.urn.Basename(), size=1000), replace=replace)
def testResourceUsageStats(self): client_ids = self.SetupClients(10) with implementation.StartHunt( hunt_name=standard.GenericHunt.__name__, flow_runner_args=rdf_flow_runner.FlowRunnerArgs( flow_name=transfer.GetFile.__name__), flow_args=transfer.GetFileArgs( pathspec=rdf_paths.PathSpec( path="/tmp/evil.txt", pathtype=rdf_paths.PathSpec.PathType.OS, )), client_rule_set=self._CreateForemanClientRuleSet(), output_plugins=[], client_rate=0, token=self.token) as hunt: hunt.Run() self.AssignTasksToClients(client_ids=client_ids) client_mock = hunt_test_lib.SampleHuntMock() hunt_test_lib.TestHuntHelper(client_mock, client_ids, False, self.token) hunt = aff4.FACTORY.Open( hunt.urn, aff4_type=standard.GenericHunt, token=self.token) # This is called once for each state method. Each flow above runs the # Start and the StoreResults methods. usage_stats = hunt.context.usage_stats self.assertEqual(usage_stats.user_cpu_stats.num, 10) self.assertTrue(math.fabs(usage_stats.user_cpu_stats.mean - 5.5) < 1e-7) self.assertTrue( math.fabs(usage_stats.user_cpu_stats.std - 2.8722813) < 1e-7) self.assertEqual(usage_stats.system_cpu_stats.num, 10) self.assertTrue(math.fabs(usage_stats.system_cpu_stats.mean - 11) < 1e-7) self.assertTrue( math.fabs(usage_stats.system_cpu_stats.std - 5.7445626) < 1e-7) self.assertEqual(usage_stats.network_bytes_sent_stats.num, 10) self.assertTrue( math.fabs(usage_stats.network_bytes_sent_stats.mean - 16.5) < 1e-7) self.assertTrue( math.fabs(usage_stats.network_bytes_sent_stats.std - 8.61684396) < 1e-7) # NOTE: Not checking histograms here. RunningStatsTest tests that mean, # standard deviation and histograms are calculated correctly. Therefore # if mean/stdev values are correct histograms should be ok as well. self.assertLen(usage_stats.worst_performers, 10) prev = usage_stats.worst_performers[0] for p in usage_stats.worst_performers[1:]: self.assertTrue( prev.cpu_usage.user_cpu_time + prev.cpu_usage.system_cpu_time > p.cpu_usage.user_cpu_time + p.cpu_usage.system_cpu_time) prev = p
def testHuntModificationWorksCorrectly(self): """This tests running the hunt on some clients.""" with implementation.StartHunt( hunt_name=standard.GenericHunt.__name__, flow_runner_args=rdf_flow_runner.FlowRunnerArgs( flow_name=transfer.GetFile.__name__), flow_args=transfer.GetFileArgs( pathspec=rdf_paths.PathSpec( path="/tmp/evil.txt", pathtype=rdf_paths.PathSpec.PathType.OS),), client_rule_set=self._CreateForemanClientRuleSet(), client_limit=1, client_rate=0, token=self.token) as hunt: hunt.Run() # Forget about hunt object, we'll use AFF4 for everything. hunt_session_id = hunt.session_id hunt = None # Pretend to be the foreman now and dish out hunting jobs to all the # client.. self.AssignTasksToClients() # Run the hunt. client_mock = hunt_test_lib.SampleHuntMock() hunt_test_lib.TestHuntHelper(client_mock, self.client_ids, False, self.token) # Re-open the hunt to get fresh data. hunt_obj = aff4.FACTORY.Open( hunt_session_id, age=aff4.ALL_TIMES, token=self.token) # There should be only one client, due to the limit started, _, _ = hunt_obj.GetClientsCounts() self.assertEqual(started, 1) # Check the hunt is paused. self.assertEqual(hunt_obj.Get(hunt_obj.Schema.STATE), "PAUSED") with aff4.FACTORY.Open( hunt_session_id, mode="rw", token=self.token) as hunt_obj: runner = hunt_obj.GetRunner() runner.runner_args.client_limit = 10 runner.Start() # Pretend to be the foreman now and dish out hunting jobs to all the # clients. self.AssignTasksToClients() hunt_test_lib.TestHuntHelper(client_mock, self.client_ids, False, self.token) hunt_obj = aff4.FACTORY.Open( hunt_session_id, age=aff4.ALL_TIMES, token=self.token) # There should be only one client, due to the limit started, _, _ = hunt_obj.GetClientsCounts() self.assertEqual(started, 10)
def _RunHunt(self, client_ids, client_mock=None, iteration_limit=None): foreman_obj = foreman.Foreman() for client_id in client_ids: foreman_obj.AssignTasksToClient(client_id) if client_mock is None: client_mock = hunt_test_lib.SampleHuntMock(failrate=2) return hunt_test_lib.TestHuntHelper( client_mock, client_ids, iteration_limit=iteration_limit)
def testHuntIsStoppedIfTotalNetworkUsageIsTooHigh(self): client_ids = self.SetupClients(5) hunt_id = self._CreateHunt( client_rule_set=foreman_rules.ForemanClientRuleSet(), client_rate=0, total_network_bytes_limit=5, args=self.GetFileHuntArgs()) def CheckState(hunt_state, network_bytes_sent): hunt_obj = data_store.REL_DB.ReadHuntObject(hunt_id) self.assertEqual(hunt_obj.hunt_state, hunt_state) hunt_counters = data_store.REL_DB.ReadHuntCounters(hunt_id) self.assertEqual(hunt_counters.total_network_bytes_sent, network_bytes_sent) self._RunHunt( client_ids[:2], client_mock=hunt_test_lib.SampleHuntMock(network_bytes_sent=2)) # 4 is lower than the total limit. The hunt should still be running. CheckState(rdf_hunt_objects.Hunt.HuntState.STARTED, 4) self._RunHunt( [client_ids[2]], client_mock=hunt_test_lib.SampleHuntMock(network_bytes_sent=1)) # 5 is equal to the total limit. Total network bytes sent should # go over the limit in order for the hunt to be stopped. CheckState(rdf_hunt_objects.Hunt.HuntState.STARTED, 5) self._RunHunt( [client_ids[3]], client_mock=hunt_test_lib.SampleHuntMock(network_bytes_sent=1)) # 6 is greater than the total limit. The hunt should be stopped now. CheckState(rdf_hunt_objects.Hunt.HuntState.STOPPED, 6) self._RunHunt([client_ids[4]], client_mock=hunt_test_lib.SampleHuntMock( network_bytes_sent=2, failrate=-1)) self._CheckHuntStoppedNotification( "reached the total network bytes sent limit")
def SetupHuntDetailView(self, failrate=2): """Create some clients and a hunt to view.""" with self.CreateSampleHunt() as hunt: hunt.LogClientError(self.client_ids[1], "Client Error 1", traceback.format_exc()) # Run the hunt. client_mock = hunt_test_lib.SampleHuntMock(failrate=failrate) hunt_test_lib.TestHuntHelper(client_mock, self.client_ids, False, self.token)
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 _RunHunt(self, client_ids, client_mock=None, iteration_limit=None): foreman_obj = foreman.GetForeman() for client_id in client_ids: foreman_obj.AssignTasksToClient(client_id.Basename()) if client_mock is None: client_mock = hunt_test_lib.SampleHuntMock() return hunt_test_lib.TestHuntHelper(client_mock, client_ids, False, iteration_limit=iteration_limit)
def testPausingAndRestartingDoesNotStartHuntTwiceOnTheSameClient(self): """This tests if the hunt completes when some clients hang or raise.""" client_ids = self.SetupClients(10) client_rule_set = rdf_foreman.ForemanClientRuleSet(rules=[ rdf_foreman.ForemanClientRule( rule_type=rdf_foreman.ForemanClientRule.Type.REGEX, regex=rdf_foreman.ForemanRegexClientRule( field="CLIENT_NAME", attribute_regex="GRR")) ]) with implementation.GRRHunt.StartHunt( hunt_name=standard.SampleHunt.__name__, client_rule_set=client_rule_set, client_rate=0, token=self.token) as hunt: hunt.GetRunner().Start() hunt_id = hunt.urn foreman = aff4.FACTORY.Open("aff4:/foreman", mode="rw", token=self.token) for client_id in client_ids: num_tasks = foreman.AssignTasksToClient(client_id.Basename()) self.assertEqual(num_tasks, 1) client_mock = hunt_test_lib.SampleHuntMock() hunt_test_lib.TestHuntHelper(client_mock, client_ids, False, self.token) # Pausing and running hunt: this leads to the fresh rules being written # to Foreman.RULES. with aff4.FACTORY.Open(hunt_id, mode="rw", token=self.token) as hunt: runner = hunt.GetRunner() runner.Pause() runner.Start() # Recreating the foreman so that it updates list of rules. foreman = aff4.FACTORY.Open("aff4:/foreman", mode="rw", token=self.token) for client_id in client_ids: num_tasks = foreman.AssignTasksToClient(client_id.Basename()) # No tasks should be assigned as this hunt ran on all the clients # before. self.assertEqual(num_tasks, 0)
def testHuntExpiration(self): """This tests that hunts with a client limit terminate correctly.""" with test_lib.FakeTime(1000): with implementation.StartHunt( hunt_name=standard.GenericHunt.__name__, flow_runner_args=rdf_flow_runner.FlowRunnerArgs( flow_name=transfer.GetFile.__name__), flow_args=transfer.GetFileArgs( pathspec=rdf_paths.PathSpec( path="/tmp/evil.txt", pathtype=rdf_paths.PathSpec.PathType.OS)), client_rule_set=self._CreateForemanClientRuleSet(), client_limit=5, expiry_time=rdfvalue.Duration("1000s"), token=self.token) as hunt: hunt.Run() # Pretend to be the foreman now and dish out hunting jobs to all the # clients (Note we have 10 clients here). self.AssignTasksToClients() hunt_obj = aff4.FACTORY.Open( hunt.session_id, age=aff4.ALL_TIMES, token=self.token) self.assertEqual(hunt_obj.Get(hunt_obj.Schema.STATE), "STARTED") # Now advance the time such that the hunt expires. time.time = lambda: 5000 # Run the hunt. client_mock = hunt_test_lib.SampleHuntMock() hunt_test_lib.TestHuntHelper( client_mock, self.client_ids, check_flow_errors=False, token=self.token) # No client should be processed since the hunt is expired. started, finished, errors = hunt_obj.GetClientsCounts() self.assertEqual(started, 0) self.assertEqual(finished, 0) self.assertEqual(errors, 0) hunt_obj = aff4.FACTORY.Open( hunt.session_id, age=aff4.ALL_TIMES, token=self.token) # Hunts are automatically stopped when they expire. self.assertEqual(hunt_obj.Get(hunt_obj.Schema.STATE), "COMPLETED")