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 _ArgsToHuntArgs(
self, args: ApiCreatePerClientFileCollectionHuntArgs
) -> rdf_hunt_objects.HuntArguments:
flow_groups = []
for client_arg in args.per_client_args:
pathspecs = []
for p in client_arg.paths:
pathspecs.append(
rdf_paths.PathSpec(path=p, pathtype=client_arg.path_type))
flow_name = transfer.MultiGetFile.__name__
flow_args = transfer.MultiGetFileArgs(pathspecs=pathspecs)
flow_group = rdf_hunt_objects.VariableHuntFlowGroup(
client_ids=[client_arg.client_id],
flow_name=flow_name,
flow_args=rdf_structs.AnyValue.Pack(flow_args))
flow_groups.append(flow_group)
return rdf_hunt_objects.HuntArguments(
hunt_type=rdf_hunt_objects.HuntArguments.HuntType.VARIABLE,
variable=rdf_hunt_objects.HuntArgumentsVariable(
flow_groups=flow_groups))
def testStartVariableHuntRaisesIfMoreThanOneFlowPerClient(self):
client_id = self.SetupClients(1)[0]
hunt_obj = rdf_hunt_objects.Hunt(client_rate=0)
hunt_obj.args.hunt_type = hunt_obj.args.HuntType.VARIABLE
for index in range(2):
hunt_obj.args.variable.flow_groups.append(
rdf_hunt_objects.VariableHuntFlowGroup(
client_ids=[client_id],
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)
with self.assertRaises(hunt.CanStartAtMostOneFlowPerClientError):
hunt.StartHunt(hunt_obj.hunt_id)
# Check that no flows were scheduled on the client.
flows = data_store.REL_DB.ReadAllFlowObjects(client_id)
self.assertEmpty(flows)
