def testForemanRulesAreCorrectlyPropagatedWhenHuntStarts(self):
client_rule_set = foreman_rules.ForemanClientRuleSet(rules=[
foreman_rules.ForemanClientRule(
rule_type=foreman_rules.ForemanClientRule.Type.REGEX,
regex=foreman_rules.ForemanRegexClientRule(
field="CLIENT_NAME", attribute_regex="HUNT")),
foreman_rules.ForemanClientRule(
rule_type=foreman_rules.ForemanClientRule.Type.INTEGER,
integer=foreman_rules.ForemanIntegerClientRule(
field="CLIENT_CLOCK",
operator=foreman_rules.ForemanIntegerClientRule.Operator
.GREATER_THAN,
value=1336650631137737))
])
self.assertEmpty(data_store.REL_DB.ReadAllForemanRules())
hunt_obj = rdf_hunt_objects.Hunt(client_rule_set=client_rule_set)
hunt_obj.args.hunt_type = hunt_obj.args.HuntType.STANDARD
data_store.REL_DB.WriteHuntObject(hunt_obj)
hunt_obj = hunt.StartHunt(hunt_obj.hunt_id)
rules = data_store.REL_DB.ReadAllForemanRules()
self.assertLen(rules, 1)
rule = rules[0]
self.assertEqual(rule.client_rule_set, client_rule_set)
self.assertEqual(rule.hunt_id, hunt_obj.hunt_id)
self.assertEqual(rule.expiration_time, hunt_obj.expiry_time)
# Running a second time should not change the rules any more.
with self.assertRaises(hunt.OnlyPausedHuntCanBeStartedError):
hunt_obj = hunt.StartHunt(hunt_obj.hunt_id)
rules = data_store.REL_DB.ReadAllForemanRules()
self.assertEqual(len(rules), 1)
def testVariableHuntCanNotBeScheduledWithNonZeroClientRate(self):
hunt_obj = rdf_hunt_objects.Hunt(client_rate=42)
hunt_obj.args.hunt_type = hunt_obj.args.HuntType.VARIABLE
data_store.REL_DB.WriteHuntObject(hunt_obj)
with self.assertRaises(hunt.VariableHuntCanNotHaveClientRateError):
hunt.StartHunt(hunt_obj.hunt_id)
hunt_obj.client_rate = 0
data_store.REL_DB.WriteHuntObject(hunt_obj)
hunt.StartHunt(hunt_obj.hunt_id)
def Run(self):
failing_descriptor = rdf_output_plugin.OutputPluginDescriptor(
plugin_name=hunt_test_lib.FailingDummyHuntOutputPlugin.__name__)
with test_lib.FakeTime(42, increment=1):
if data_store.RelationalDBEnabled():
hunt_id = self.CreateHunt(description="the hunt",
output_plugins=[failing_descriptor])
hunt.StartHunt(hunt_id)
else:
hunt_urn = self.StartHunt(description="the hunt",
output_plugins=[failing_descriptor])
hunt_id = hunt_urn.Basename()
self.client_ids = self.SetupClients(2)
for index, client_id in enumerate(self.client_ids):
self.RunHunt(client_ids=[client_id], failrate=-1)
with test_lib.FakeTime(100042 + index * 100):
try:
self.ProcessHuntOutputPlugins()
except process_results.ResultsProcessingError:
if flags.FLAGS.pdb_post_mortem:
pdb.post_mortem()
self.Check("ListHuntOutputPluginErrors",
args=hunt_plugin.ApiListHuntOutputPluginErrorsArgs(
hunt_id=hunt_id,
plugin_id="FailingDummyHuntOutputPlugin_0"),
replace={hunt_id: "H:123456"})
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 testForemanRulesAreCorrectlyRemovedWhenHuntIsStopped(self):
client_rule_set = foreman_rules.ForemanClientRuleSet(rules=[
foreman_rules.ForemanClientRule(
rule_type=foreman_rules.ForemanClientRule.Type.REGEX,
regex=foreman_rules.ForemanRegexClientRule(
field="CLIENT_NAME", attribute_regex="HUNT")),
foreman_rules.ForemanClientRule(
rule_type=foreman_rules.ForemanClientRule.Type.INTEGER,
integer=foreman_rules.ForemanIntegerClientRule(
field="CLIENT_CLOCK",
operator=foreman_rules.ForemanIntegerClientRule.Operator.
GREATER_THAN,
value=1336650631137737))
])
hunt_obj = rdf_hunt_objects.Hunt(client_rule_set=client_rule_set)
hunt_obj.args.hunt_type = hunt_obj.args.HuntType.STANDARD
data_store.REL_DB.WriteHuntObject(hunt_obj)
hunt_obj = hunt.StartHunt(hunt_obj.hunt_id)
rules = data_store.REL_DB.ReadAllForemanRules()
self.assertLen(rules, 1)
hunt.StopHunt(hunt_obj.hunt_id)
rules = data_store.REL_DB.ReadAllForemanRules()
self.assertEmpty(rules)
def testForemanRulesWorkCorrectlyWithStandardHunt(self):
client_rule_set = foreman_rules.ForemanClientRuleSet(rules=[
foreman_rules.ForemanClientRule(
rule_type=foreman_rules.ForemanClientRule.Type.OS,
os=foreman_rules.ForemanOsClientRule(os_windows=True))
])
hunt_obj = rdf_hunt_objects.Hunt(client_rule_set=client_rule_set,
client_rate=0,
args=self.GetFileHuntArgs())
hunt_obj.args.hunt_type = hunt_obj.args.HuntType.STANDARD
data_store.REL_DB.WriteHuntObject(hunt_obj)
hunt.StartHunt(hunt_obj.hunt_id)
# Check matching client.
client_id = self.SetupClient(0, system="Windows")
foreman_obj = foreman.Foreman()
foreman_obj.AssignTasksToClient(client_id)
flows = data_store.REL_DB.ReadAllFlowObjects(client_id=client_id)
self.assertLen(flows, 1)
# Check non-matching client.
client_id = self.SetupClient(1, system="Linux")
foreman_obj.AssignTasksToClient(client_id)
flows = data_store.REL_DB.ReadAllFlowObjects(client_id=client_id)
self.assertEmpty(flows)
def Run(self):
with test_lib.FakeTime(42):
if data_store.RelationalDBEnabled():
hunt_id = self.CreateHunt(description="the hunt")
hunt.StartHunt(hunt_id)
else:
hunt_urn = self.StartHunt(description="the hunt")
hunt_id = hunt_urn.Basename()
if data_store.RelationalDBEnabled():
client = self.SetupTestClientObject(0)
client_ids = [rdf_client.ClientURN(client.client_id)]
else:
client_ids = self.SetupClients(1)
self.RunHunt(failrate=2, client_ids=client_ids)
# Create replace dictionary.
replace = {hunt_id: "H:123456"}
if data_store.RelationalDBEnabled():
stats = data_store.REL_DB.ReadHuntClientResourcesStats(hunt_id)
for performance in stats.worst_performers:
session_id = unicode(performance.session_id)
replace[session_id] = "<replaced session value>"
else:
with aff4.FACTORY.Open(hunt_urn, mode="r",
token=self.token) as hunt_obj:
stats = hunt_obj.GetRunner().context.usage_stats
for performance in stats.worst_performers:
session_id = unicode(performance.session_id)
replace[session_id] = "<replaced session value>"
self.Check("GetHuntStats",
args=hunt_plugin.ApiGetHuntStatsArgs(hunt_id=hunt_id),
replace=replace)
def testRaisesWhenStoppingHuntInTheWrongState(self):
hunt.StartHunt(self.hunt_id)
hunt.StopHunt(self.hunt_id)
self.args.state = "STOPPED"
with self.assertRaises(hunt_plugin.HuntNotStoppableError):
self.handler.Handle(self.args, token=self.token)
def Run(self):
output_plugins = [
rdf_output_plugin.OutputPluginDescriptor(
plugin_name=test_plugins.DummyHuntTestOutputPlugin.__name__,
plugin_args=test_plugins.DummyHuntTestOutputPlugin.args_type(
filename_regex="blah!", fetch_binaries=True))
]
with test_lib.FakeTime(42, increment=1):
if data_store.RelationalDBEnabled():
hunt_id = self.CreateHunt(description="the hunt",
output_plugins=output_plugins)
hunt.StartHunt(hunt_id)
else:
hunt_urn = self.StartHunt(description="the hunt",
output_plugins=output_plugins)
hunt_id = hunt_urn.Basename()
self.client_ids = self.SetupClients(2)
for index, client_id in enumerate(self.client_ids):
self.RunHunt(client_ids=[client_id], failrate=-1)
with test_lib.FakeTime(100042 + index * 100):
self.ProcessHuntOutputPlugins()
self.Check("ListHuntOutputPluginLogs",
args=hunt_plugin.ApiListHuntOutputPluginLogsArgs(
hunt_id=hunt_id,
plugin_id="DummyHuntTestOutputPlugin_0"),
replace={hunt_id: "H:123456"})
def Run(self):
with test_lib.FakeTime(42):
hunt_id = self.CreateHunt(description="the hunt",
creator=self.token.username)
hunt.StartHunt(hunt_id)
client_ids = self.SetupClients(5)
self.AssignTasksToClients(client_ids=client_ids[:-1])
# Only running the hunt on a single client, as SampleMock
# implementation is non-deterministic in terms of resources
# usage that gets reported back to the hunt.
self.RunHunt(client_ids=[client_ids[-1]], failrate=0)
# Create replace dictionary.
replace = {hunt_id: "H:123456", hunt_id + ":hunt": "H:123456"}
self.Check("ListHuntClients",
args=hunt_plugin.ApiListHuntClientsArgs(
hunt_id=hunt_id, client_status="STARTED"),
replace=replace)
self.Check("ListHuntClients",
args=hunt_plugin.ApiListHuntClientsArgs(
hunt_id=hunt_id, client_status="OUTSTANDING"),
replace=replace)
self.Check("ListHuntClients",
args=hunt_plugin.ApiListHuntClientsArgs(
hunt_id=hunt_id, client_status="COMPLETED"),
replace=replace)
def testVariableHuntSchedulesAllFlowsOnStart(self):
client_ids = self.SetupClients(10)
hunt_obj = rdf_hunt_objects.Hunt(client_rate=0)
hunt_obj.args.hunt_type = hunt_obj.args.HuntType.VARIABLE
for index, pair in enumerate(collection.Batch(client_ids, 2)):
hunt_obj.args.variable.flow_groups.append(
rdf_hunt_objects.VariableHuntFlowGroup(
client_ids=pair,
flow_name=compatibility.GetName(transfer.GetFile),
flow_args=transfer.GetFileArgs(pathspec=rdf_paths.PathSpec(
path="/tmp/evil_%d.txt" % index,
pathtype=rdf_paths.PathSpec.PathType.OS,
))))
data_store.REL_DB.WriteHuntObject(hunt_obj)
hunt.StartHunt(hunt_obj.hunt_id)
hunt_counters = data_store.REL_DB.ReadHuntCounters(hunt_obj.hunt_id)
self.assertEqual(hunt_counters.num_clients, 10)
all_flows = data_store.REL_DB.ReadHuntFlows(hunt_obj.hunt_id, 0,
sys.maxsize)
self.assertCountEqual(client_ids, [f.client_id for f in all_flows])
for index, pair in enumerate(collection.Batch(client_ids, 2)):
for client_id in pair:
all_flows = data_store.REL_DB.ReadAllFlowObjects(client_id)
self.assertLen(all_flows, 1)
self.assertEqual(all_flows[0].flow_class_name,
compatibility.GetName(transfer.GetFile))
self.assertEqual(all_flows[0].args.pathspec.path,
"/tmp/evil_%d.txt" % index)
def testRaisesWhenStartingHuntInTheWrongState(self):
hunt.StartHunt(self.hunt_id)
hunt.StopHunt(self.hunt_id)
self.args.state = "STARTED"
with self.assertRaises(hunt_plugin.HuntNotStartableError):
self.handler.Handle(self.args, context=self.context)
def Handle(self, args, context=None):
hunt_id = str(args.hunt_id)
has_change = False
for field_name in ["client_limit", "client_rate", "expires"]:
if args.HasField(field_name):
has_change = True
break
try:
hunt_obj = data_store.REL_DB.ReadHuntObject(hunt_id)
if has_change:
kw_args = {}
if hunt_obj.hunt_state != hunt_obj.HuntState.PAUSED:
raise HuntNotModifiableError(
"Hunt's client limit/client rate/expiry time attributes "
"can only be changed if hunt's current state is "
"PAUSED")
if args.HasField("client_limit"):
kw_args["client_limit"] = args.client_limit
if args.HasField("client_rate"):
kw_args["client_rate"] = args.client_rate
if args.duration is None:
kw_args["duration"] = None
else:
kw_args["duration"] = rdfvalue.Duration(args.duration)
data_store.REL_DB.UpdateHuntObject(hunt_id, **kw_args)
hunt_obj = data_store.REL_DB.ReadHuntObject(hunt_id)
except db.UnknownHuntError:
raise HuntNotFoundError("Hunt with id %s could not be found" %
args.hunt_id)
if args.HasField("state"):
if args.state == ApiHunt.State.STARTED:
if hunt_obj.hunt_state != hunt_obj.HuntState.PAUSED:
raise HuntNotStartableError(
"Hunt can only be started from PAUSED state.")
hunt_obj = hunt.StartHunt(hunt_obj.hunt_id)
elif args.state == ApiHunt.State.STOPPED:
if hunt_obj.hunt_state not in [
hunt_obj.HuntState.PAUSED, hunt_obj.HuntState.STARTED
]:
raise HuntNotStoppableError(
"Hunt can only be stopped from STARTED or "
"PAUSED states.")
hunt_obj = hunt.StopHunt(hunt_obj.hunt_id)
else:
raise InvalidHuntStateError(
"Hunt's state can only be updated to STARTED or STOPPED")
hunt_counters = data_store.REL_DB.ReadHuntCounters(hunt_obj.hunt_id)
return ApiHunt().InitFromHuntObject(
hunt_obj, hunt_counters=hunt_counters, with_full_summary=True)
def testRaisesIfHuntIsRunning(self):
if data_store.RelationalDBEnabled():
hunt.StartHunt(self.hunt_id)
else:
self.hunt_obj.Run()
with self.assertRaises(hunt_plugin.HuntNotDeletableError):
self.handler.Handle(self.args, token=self.token)
def testRaisesWhenModifyingHuntInNonPausedState(self):
if data_store.RelationalDBEnabled():
hunt.StartHunt(self.hunt_id)
else:
self.hunt_obj.Run()
self.args.client_rate = 100
with self.assertRaises(hunt_plugin.HuntNotModifiableError):
self.handler.Handle(self.args, token=self.token)
def testHuntListShowsStartAndExpirationTime(self):
hunt_1_start_time = rdfvalue.RDFDatetime.FromHumanReadable(
"1992-11-11")
hunt_2_start_time = rdfvalue.RDFDatetime.FromHumanReadable(
"2001-05-03")
hunt_1_duration = rdfvalue.Duration("3d")
hunt_2_duration = rdfvalue.Duration("5h")
hunt_1_expiration_time = hunt_1_start_time + hunt_1_duration
hunt_2_expiration_time = hunt_2_start_time + hunt_2_duration
hunt_1_id = self.CreateHunt(duration=hunt_1_duration)
hunt_2_id = self.CreateHunt(duration=hunt_2_duration)
# Navigate to the hunt list.
self.Open("/")
self.WaitUntil(self.IsElementPresent, "client_query")
self.Click("css=a[grrtarget=hunts]")
self.WaitUntil(self.IsTextPresent, hunt_1_id)
self.WaitUntil(self.IsTextPresent, hunt_2_id)
self.assertFalse(self.IsTextPresent(str(hunt_1_start_time)))
self.assertFalse(self.IsTextPresent(str(hunt_1_expiration_time)))
self.assertFalse(self.IsTextPresent(str(hunt_2_start_time)))
self.assertFalse(self.IsTextPresent(str(hunt_2_expiration_time)))
with test_lib.FakeTime(hunt_1_start_time):
hunt.StartHunt(hunt_1_id)
self.Refresh()
self.WaitUntil(self.IsTextPresent, str(hunt_1_start_time))
self.WaitUntil(self.IsTextPresent, str(hunt_1_expiration_time))
self.assertFalse(self.IsTextPresent(str(hunt_2_start_time)))
self.assertFalse(self.IsTextPresent(str(hunt_2_expiration_time)))
with test_lib.FakeTime(hunt_2_start_time):
hunt.StartHunt(hunt_2_id)
self.Refresh()
self.WaitUntil(self.IsTextPresent, str(hunt_1_start_time))
self.WaitUntil(self.IsTextPresent, str(hunt_1_expiration_time))
self.WaitUntil(self.IsTextPresent, str(hunt_2_start_time))
self.WaitUntil(self.IsTextPresent, str(hunt_2_expiration_time))
def Run(self):
if data_store.RelationalDBEnabled():
client_obj = self.SetupTestClientObject(0)
client_id = client_obj.client_id
else:
client_id = self.SetupClient(0).Basename()
client_mocks = {
client_id: flow_test_lib.CrashClientMock(client_id, self.token)
}
if data_store.RelationalDBEnabled():
hunt_id = self.CreateHunt(description="the hunt")
hunt.StartHunt(hunt_id)
else:
with test_lib.FakeTime(42):
with self.CreateHunt(description="the hunt") as hunt_obj:
hunt_obj.Run()
hunt_id = hunt_obj.urn.Basename()
with test_lib.FakeTime(45):
self.AssignTasksToClients([client_id])
hunt_test_lib.TestHuntHelperWithMultipleMocks(
client_mocks, False, self.token)
if data_store.RelationalDBEnabled():
crash = data_store.REL_DB.ReadHuntFlows(
hunt_id,
0,
1,
filter_condition=db.HuntFlowsCondition.CRASHED_FLOWS_ONLY,
)[0].client_crash_info
else:
crashes = implementation.GRRHunt.CrashCollectionForHID(
hunt_obj.urn)
crash = list(crashes)[0]
replace = {
hunt_id:
"H:123456",
unicode(crash.session_id):
"aff4:/hunts/H:123456/C.1000000000000000/H:11223344"
}
self.Check("ListHuntCrashes",
args=hunt_plugin.ApiListHuntCrashesArgs(hunt_id=hunt_id),
replace=replace)
self.Check("ListHuntCrashes",
args=hunt_plugin.ApiListHuntCrashesArgs(hunt_id=hunt_id,
count=1),
replace=replace)
self.Check("ListHuntCrashes",
args=hunt_plugin.ApiListHuntCrashesArgs(hunt_id=hunt_id,
offset=1,
count=1),
replace=replace)
def testRaisesWhenStoppingHuntInTheWrongState(self):
if data_store.RelationalDBEnabled():
hunt.StartHunt(self.hunt_id)
hunt.StopHunt(self.hunt_id)
else:
self.hunt_obj.Run()
self.hunt_obj.Stop()
self.args.state = "STOPPED"
with self.assertRaises(hunt_plugin.HuntNotStoppableError):
self.handler.Handle(self.args, token=self.token)
def StartHunt(self, paused=False, **kwargs):
if data_store.RelationalDBEnabled():
hunt_id = self.CreateHunt(**kwargs)
if not paused:
hunt.StartHunt(hunt_id)
return rdfvalue.RDFURN("hunts").Add(hunt_id)
else:
with self.CreateHunt(**kwargs) as hunt_obj:
if not paused:
hunt_obj.Run()
return hunt_obj.urn
def testHuntOverviewShowsStartAndExpirationTime(self):
duration = rdfvalue.Duration("3d")
init_start_time = rdfvalue.RDFDatetime.FromHumanReadable(
"1973-01-01 08:34")
last_start_time = rdfvalue.RDFDatetime.FromHumanReadable(
"1981-03-04 12:52")
expiration_time = init_start_time + duration
hunt_id = self.CreateHunt(duration=duration)
# Navigate to the hunt view.
self.Open("/")
self.WaitUntil(self.IsElementPresent, "client_query")
self.Click("css=a[grrtarget=hunts]")
self.WaitUntil(self.IsTextPresent, hunt_id)
# Select the hunt.
self.Click("css=td:contains('{}')".format(hunt_id))
self.RequestAndGrantHuntApproval(hunt_id)
self.assertFalse(self.IsTextPresent(str(init_start_time)))
self.assertFalse(self.IsTextPresent(str(expiration_time)))
self.assertFalse(self.IsTextPresent(str(last_start_time)))
with test_lib.FakeTime(init_start_time):
hunt.StartHunt(hunt_id)
self.Refresh()
self.WaitUntil(self.IsTextPresent, str(init_start_time))
self.WaitUntil(self.IsTextPresent, str(expiration_time))
self.assertFalse(self.IsTextPresent(str(last_start_time)))
with test_lib.FakeTime(last_start_time):
hunt.PauseHunt(hunt_id)
hunt.StartHunt(hunt_id)
self.Refresh()
self.WaitUntil(self.IsTextPresent, str(init_start_time))
self.WaitUntil(self.IsTextPresent, str(expiration_time))
self.WaitUntil(self.IsTextPresent, str(last_start_time))
def testHuntWithInvalidForemanRulesDoesNotStart(self):
client_rule_set = foreman_rules.ForemanClientRuleSet(rules=[
foreman_rules.ForemanClientRule(
rule_type=foreman_rules.ForemanClientRule.Type.REGEX,
regex=foreman_rules.ForemanRegexClientRule(
field="UNSET", attribute_regex="HUNT"))
])
hunt_obj = rdf_hunt_objects.Hunt(client_rule_set=client_rule_set)
data_store.REL_DB.WriteHuntObject(hunt_obj)
with self.assertRaises(ValueError):
hunt.StartHunt(hunt_obj.hunt_id)
def Run(self):
output_plugins = [
rdf_output_plugin.OutputPluginDescriptor(
plugin_name=test_plugins.DummyHuntTestOutputPlugin.__name__,
plugin_args=test_plugins.DummyHuntTestOutputPlugin.args_type(
filename_regex="blah!", fetch_binaries=True))
]
with test_lib.FakeTime(42):
hunt_id = self.CreateHunt(description="the hunt",
output_plugins=output_plugins)
hunt.StartHunt(hunt_id)
self.Check(
"ListHuntOutputPlugins",
args=hunt_plugin.ApiListHuntOutputPluginsArgs(hunt_id=hunt_id),
replace={hunt_id: "H:123456"})
def testPausingAndRestartingDoesNotStartHuntTwiceOnTheSameClient(self):
hunt_id, client_ids = self._CreateAndRunHunt(
num_clients=10,
client_rule_set=foreman_rules.ForemanClientRuleSet(),
client_rate=0,
args=self.GetFileHuntArgs())
for client_id in client_ids:
flows = data_store.REL_DB.ReadAllFlowObjects(client_id=client_id)
self.assertLen(flows, 1)
hunt.PauseHunt(hunt_id)
hunt.StartHunt(hunt_id)
self._RunHunt(client_ids)
for client_id in client_ids:
flows = data_store.REL_DB.ReadAllFlowObjects(client_id=client_id)
self.assertLen(flows, 1)
def Run(self):
with test_lib.FakeTime(42):
hunt_id = self.CreateHunt(description="the hunt",
creator=self.token.username)
hunt.StartHunt(hunt_id)
client_id = self.SetupClient(0)
self.RunHunt(failrate=2, client_ids=[client_id])
# Create replace dictionary.
replace = {hunt_id: "H:123456"}
stats = data_store.REL_DB.ReadHuntClientResourcesStats(hunt_id)
for performance in stats.worst_performers:
session_id = str(performance.session_id)
replace[session_id] = "<replaced session value>"
self.Check("GetHuntStats",
args=hunt_plugin.ApiGetHuntStatsArgs(hunt_id=hunt_id),
replace=replace)
def Run(self):
failing_descriptor = rdf_output_plugin.OutputPluginDescriptor(
plugin_name=hunt_test_lib.FailingDummyHuntOutputPlugin.__name__)
with test_lib.FakeTime(42, increment=1):
hunt_id = self.CreateHunt(description="the hunt",
output_plugins=[failing_descriptor],
creator=self.token.username)
hunt.StartHunt(hunt_id)
self.client_ids = self.SetupClients(2)
for client_id in self.client_ids:
self.RunHunt(client_ids=[client_id], failrate=-1)
self.Check("ListHuntOutputPluginErrors",
args=hunt_plugin.ApiListHuntOutputPluginErrorsArgs(
hunt_id=hunt_id,
plugin_id="FailingDummyHuntOutputPlugin_0"),
replace={hunt_id: "H:123456"})
def Run(self):
client_id = self.SetupClient(0)
client_mocks = {
client_id: flow_test_lib.CrashClientMock(client_id, self.token)
}
hunt_id = self.CreateHunt(description="the hunt",
creator=self.token.username)
hunt.StartHunt(hunt_id)
with test_lib.FakeTime(45):
self.AssignTasksToClients([client_id])
hunt_test_lib.TestHuntHelperWithMultipleMocks(
client_mocks, self.token)
crash = data_store.REL_DB.ReadHuntFlows(
hunt_id,
0,
1,
filter_condition=db.HuntFlowsCondition.CRASHED_FLOWS_ONLY,
)[0].client_crash_info
replace = {
hunt_id:
"H:123456",
str(crash.session_id):
"aff4:/hunts/H:123456/C.1000000000000000/H:11223344"
}
self.Check("ListHuntCrashes",
args=hunt_plugin.ApiListHuntCrashesArgs(hunt_id=hunt_id),
replace=replace)
self.Check("ListHuntCrashes",
args=hunt_plugin.ApiListHuntCrashesArgs(hunt_id=hunt_id,
count=1),
replace=replace)
self.Check("ListHuntCrashes",
args=hunt_plugin.ApiListHuntCrashesArgs(hunt_id=hunt_id,
offset=1,
count=1),
replace=replace)
def testHangingClientsAreCorrectlyAccountedFor(self):
client_ids = self.SetupClients(10)
hunt_obj = rdf_hunt_objects.Hunt(
client_rule_set=foreman_rules.ForemanClientRuleSet(),
client_rate=0,
args=self.GetFileHuntArgs())
hunt.CreateHunt(hunt_obj)
hunt_obj = hunt.StartHunt(hunt_obj.hunt_id)
foreman_obj = foreman.Foreman()
for client_id in client_ids:
foreman_obj.AssignTasksToClient(client_id)
client_mock = hunt_test_lib.SampleHuntMock(failrate=2)
hunt_test_lib.TestHuntHelper(client_mock, client_ids[1:9])
hunt_counters = data_store.REL_DB.ReadHuntCounters(hunt_obj.hunt_id)
self.assertEqual(hunt_counters.num_clients, 10)
self.assertEqual(hunt_counters.num_successful_clients, 4)
self.assertEqual(hunt_counters.num_failed_clients, 4)
def Run(self):
if data_store.RelationalDBEnabled():
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(failrate=2)
if data_store.RelationalDBEnabled():
hunt_id = self.CreateHunt(description="the hunt")
hunt.StartHunt(hunt_id)
else:
with test_lib.FakeTime(42):
with self.CreateHunt(description="the hunt") as hunt_obj:
hunt_obj.Run()
hunt_id = hunt_obj.urn.Basename()
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_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 Run(self):
with test_lib.FakeTime(42):
hunt_id = self.CreateHunt(description="the hunt")
hunt.StartHunt(hunt_id)
if data_store.RelationalDBEnabled():
client = self.SetupTestClientObject(0)
client_ids = [rdf_client.ClientURN(client.client_id)]
else:
client_ids = self.SetupClients(1)
self.RunHunt(failrate=2, client_ids=client_ids)
# Create replace dictionary.
replace = {hunt_id: "H:123456"}
stats = data_store.REL_DB.ReadHuntClientResourcesStats(hunt_id)
for performance in stats.worst_performers:
session_id = unicode(performance.session_id)
replace[session_id] = "<replaced session value>"
self.Check("GetHuntStats",
args=hunt_plugin.ApiGetHuntStatsArgs(hunt_id=hunt_id),
replace=replace)
def Run(self):
output_plugins = [
rdf_output_plugin.OutputPluginDescriptor(
plugin_name=test_plugins.DummyHuntTestOutputPlugin.__name__,
plugin_args=test_plugins.DummyHuntTestOutputPlugin.args_type(
filename_regex="blah!", fetch_binaries=True))
]
with test_lib.FakeTime(42, increment=1):
hunt_id = self.CreateHunt(description="the hunt",
output_plugins=output_plugins,
creator=self.token.username)
hunt.StartHunt(hunt_id)
self.client_ids = self.SetupClients(2)
for client_id in self.client_ids:
self.RunHunt(client_ids=[client_id], failrate=-1)
self.Check("ListHuntOutputPluginLogs",
args=hunt_plugin.ApiListHuntOutputPluginLogsArgs(
hunt_id=hunt_id,
plugin_id="DummyHuntTestOutputPlugin_0"),
replace={hunt_id: "H:123456"})