Ejemplo n.º 1
0
// CalculateGroups is a function designed to ingest a list of group objects, a collection of agents, and
// return a map that contains the resulting group for each agent UUID.
func CalculateGroups(cfg config.Config) {

	tdb, err := db.NewToddDB(cfg)
	if err != nil {
		log.Fatalf("Error connecting to DB: %v", err)
	}

	// Retrieve all currently active agents
	agents, err := tdb.GetAgents()
	if err != nil {
		log.Fatalf("Error retrieving agents: %v", err)
	}

	// Retrieve all objects with type "group"
	group_objs, err := tdb.GetObjects("group")
	if err != nil {
		log.Fatalf("Error retrieving groups: %v", err)
	}

	// Cast retrieved slice of ToddObject interfaces to actual GroupObjects
	groups := make([]objects.GroupObject, len(group_objs))
	for i, gobj := range group_objs {
		groups[i] = gobj.(objects.GroupObject)
	}

	// groupmap contains the uuid-to-groupname mappings to be used for test runs
	groupmap := map[string]string{}

	// This slice will hold all of the agents that are sad because they didn't get into a group
	var lonelyAgents []defs.AgentAdvert

next:
	for x := range agents {

		for i := range groups {

			// See if this agent is in this group
			if isInGroup(groups[i].Spec.Matches, agents[x].Facts) {

				// Insert this group name ("Label") into groupmap under the key of the UUID for the agent that belongs to it
				log.Debugf("Agent %s is in group %s\n", agents[x].Uuid, groups[i].Label)
				groupmap[agents[x].Uuid] = groups[i].Label
				continue next

			}
		}

		// The "continue next" should prohibit all agents that have a group from getting to this point,
		// so the only ones left do not have a group.
		lonelyAgents = append(lonelyAgents, agents[x])

	}

	// Write results to database
	err = tdb.SetGroupMap(groupmap)
	if err != nil {
		log.Fatalf("Error setting group map: %v", err)
	}

	// Send notifications to each agent to let them know what group they're in, so they can cache it
	var tc = comms.NewToDDComms(cfg)
	for uuid, groupName := range groupmap {
		setGroupTask := tasks.SetGroupTask{
			GroupName: groupName,
		}

		setGroupTask.Type = "SetGroup" //TODO(mierdin): Apparently this is necessary because inner type promotion doesn't apply for struct literals?
		tc.CommsPackage.SendTask(uuid, setGroupTask)
	}

	// need to send a message to all agents that weren't in groupmap to set their group to nothing
	for x := range lonelyAgents {
		setGroupTask := tasks.SetGroupTask{
			GroupName: "",
		}

		setGroupTask.Type = "SetGroup" //TODO(mierdin): Apparently this is necessary because inner type promotion doesn't apply for struct literals?
		tc.CommsPackage.SendTask(lonelyAgents[x].Uuid, setGroupTask)
	}
}