Exemple #1
0
func main() {

	todd_version := "0.0.1"

	cfg := config.GetConfig(arg_config)

	// Start serving collectors and testlets, and retrieve map of names and hashes
	assets := serveAssets(cfg)

	// Perform database initialization tasks
	tdb, err := db.NewToddDB(cfg)
	if err != nil {
		log.Fatalf("Error setting up database: %v\n", err)
	}

	if err := tdb.Init(); err != nil {
		log.Fatalf("Error initializing database: %v\n", err)
	}

	// Initialize API
	var tapi toddapi.ToDDApi
	go func() {
		log.Fatal(tapi.Start(cfg))
	}()

	// Start listening for agent advertisements
	var tc = comms.NewToDDComms(cfg)
	go tc.CommsPackage.ListenForAgent(assets)

	// Kick off group calculation in background
	go func() {
		for {
			log.Info("Beginning group calculation")
			grouping.CalculateGroups(cfg)
			time.Sleep(time.Second * time.Duration(cfg.Grouping.Interval))
		}
	}()

	log.Infof("ToDD server v%s. Press any key to exit...\n", todd_version)

	// Sssh, sssh, only dreams now....
	for {
		time.Sleep(time.Second * 10)
	}
}
Exemple #2
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)
	}
}
Exemple #3
0
// ListenForResponses listens for responses from an agent
func (rmq rabbitMQComms) ListenForResponses(stopListeningForResponses *chan bool) {

	queueName := "agentresponses"

	conn, err := amqp.Dial(rmq.queueUrl)
	if err != nil {
		log.Error(err)
		log.Error("Failed to connect to RabbitMQ")
		os.Exit(1)
	}
	defer conn.Close()

	ch, err := conn.Channel()
	if err != nil {
		log.Error("Failed to open a channel")
		os.Exit(1)
	}
	defer ch.Close()

	_, err = ch.QueueDeclare(
		queueName, // name
		false,     // durable
		false,     // delete when usused
		false,     // exclusive
		false,     // no-wait
		nil,       // arguments
	)
	if err != nil {
		log.Error("Failed to declare a queue")
		os.Exit(1)
	}

	msgs, err := ch.Consume(
		queueName, // queue
		queueName, // consumer
		true,      // auto-ack
		false,     // exclusive
		false,     // no-local
		false,     // no-wait
		nil,       // args
	)
	if err != nil {
		log.Error("Failed to register a consumer")
		os.Exit(1)
	}

	tdb, err := db.NewToddDB(rmq.config) // TODO(vcabbage): Consider moving this into the rabbitMQComms struct
	if err != nil {
		log.Error("Failed to connect to DB")
		os.Exit(1)
	}

	go func() {
		for d := range msgs {

			// Unmarshal into BaseResponse to determine type
			var base_msg responses.BaseResponse
			err = json.Unmarshal(d.Body, &base_msg)
			// TODO(mierdin): Need to handle this error

			log.Debugf("Agent response received: %s", d.Body)

			// call agent response method based on type
			switch base_msg.Type {
			case "AgentStatus":

				var sasr responses.SetAgentStatusResponse
				err = json.Unmarshal(d.Body, &sasr)
				// TODO(mierdin): Need to handle this error

				log.Debugf("Agent %s is '%s' regarding test %s. Writing to DB.", sasr.AgentUuid, sasr.Status, sasr.TestUuid)
				err := tdb.SetAgentTestStatus(sasr.TestUuid, sasr.AgentUuid, sasr.Status)
				if err != nil {
					log.Errorf("Error writing agent status to DB: %v", err)
				}

			case "TestData":

				var utdr responses.UploadTestDataResponse
				err = json.Unmarshal(d.Body, &utdr)
				// TODO(mierdin): Need to handle this error

				err = tdb.SetAgentTestData(utdr.TestUuid, utdr.AgentUuid, utdr.TestData)
				// TODO(mierdin): Need to handle this error

				// Send task to the agent that says to delete the entry
				var dtdt tasks.DeleteTestDataTask
				dtdt.Type = "DeleteTestData" //TODO(mierdin): This is an extra step. Maybe a factory function for the task could help here?
				dtdt.TestUuid = utdr.TestUuid
				rmq.SendTask(utdr.AgentUuid, dtdt)

				// Finally, set the status for this agent in the test to "finished"
				err := tdb.SetAgentTestStatus(dtdt.TestUuid, utdr.AgentUuid, "finished")
				if err != nil {
					log.Errorf("Error writing agent status to DB: %v", err)
				}

			default:
				log.Errorf(fmt.Sprintf("Unexpected type value for received response: %s", base_msg.Type))
			}
		}
	}()

	log.Infof(" [*] Waiting for messages. To exit press CTRL+C")
	<-*stopListeningForResponses
}
Exemple #4
0
// ListenForAgent will listen on the message queue for new agent advertisements.
// It is meant to be run as a goroutine
func (rmq rabbitMQComms) ListenForAgent(assets map[string]map[string]string) {

	// TODO(mierdin): does func param need to be a pointer?

	conn, err := amqp.Dial(rmq.queueUrl)
	if err != nil {
		log.Error(err)
		log.Error("Failed to connect to RabbitMQ")
		os.Exit(1)
	}
	defer conn.Close()

	ch, err := conn.Channel()
	if err != nil {
		log.Error("Failed to open a channel")
		os.Exit(1)
	}
	defer ch.Close()

	q, err := ch.QueueDeclare(
		"agentadvert", // name
		false,         // durable
		false,         // delete when unused
		false,         // exclusive
		false,         // no-wait
		nil,           // arguments
	)
	if err != nil {
		log.Error("Failed to declare a queue")
		os.Exit(1)
	}

	msgs, err := ch.Consume(
		q.Name,        // queue
		"agentadvert", // consumer
		true,          // auto-ack
		false,         // exclusive
		false,         // no-local
		false,         // no-wait
		nil,           // args
	)
	if err != nil {
		log.Error("Failed to register a consumer")
		os.Exit(1)
	}

	forever := make(chan bool)

	go func() {
		for d := range msgs {
			log.Debugf("Agent advertisement recieved: %s", d.Body)

			var agent defs.AgentAdvert
			err = json.Unmarshal(d.Body, &agent)
			// TODO(mierdin): Need to handle this error

			// assetList is a slice that will contain any URLs that need to be sent to an
			// agent as a response to an incorrect or incomplete list of assets
			var assetList []string

			// assets is the asset map from the SERVER's perspective
			for asset_type, asset_hashes := range assets {

				var agentAssets map[string]string

				// agentAssets is the asset map from the AGENT's perspective
				if asset_type == "factcollectors" {
					agentAssets = agent.FactCollectors
				} else if asset_type == "testlets" {
					agentAssets = agent.Testlets
				}

				for name, hash := range asset_hashes {

					// See if the hashes match (a missing asset will also result in False)
					if agentAssets[name] != hash {

						// hashes do not match, so we need to append the asset download URL to the remediate list
						var default_ip string
						if rmq.config.LocalResources.IPAddrOverride != "" {
							default_ip = rmq.config.LocalResources.IPAddrOverride
						} else {
							default_ip = hostresources.GetIPOfInt(rmq.config.LocalResources.DefaultInterface).String()
						}
						asset_url := fmt.Sprintf("http://%s:%s/%s/%s", default_ip, rmq.config.Assets.Port, asset_type, name)

						assetList = append(assetList, asset_url)

					}
				}

			}

			// Asset list is empty, so we can continue
			if len(assetList) == 0 {

				var tdb, _ = db.NewToddDB(rmq.config)
				tdb.SetAgent(agent)

				// This block of code checked that the agent time was within a certain range of the server time. If there was a large enough
				// time skew, the agent advertisement would be rejected.
				// I have disabled this for now - My plan was to use this to synchronize testrun execution amongst agents, but I have
				// a solution to that for now. May revisit this later.
				//
				// Determine difference between server and agent time
				// t1 := time.Now()
				// var diff float64
				// diff = t1.Sub(agent.LocalTime).Seconds()
				//
				// // If Agent is within half a second of server time, add insert to database
				// if diff < 0.5 && diff > -0.5 {
				// } else {
				// 	// We don't want to register an agent if there is a severe time difference,
				// 	// in order to ensure continuity during tests. So, just print log message.
				// 	log.Warn("Agent time not within boundaries.")
				// }

			} else {
				log.Warnf("Agent %s did not have the required asset files. This advertisement is ignored.", agent.Uuid)

				var task tasks.DownloadAssetTask
				task.Type = "DownloadAsset" //TODO(mierdin): This is an extra step. Maybe a factory function for the task could help here?
				task.Assets = assetList
				rmq.SendTask(agent.Uuid, task)
			}
		}
	}()

	log.Infof(" [*] Waiting for messages. To exit press CTRL+C")
	<-forever
}