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)