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")
