[major] refactor agent status handling to enable target expansion pre-launch

doc/.files/createdb.sh

@@ -179,6 +179,14 @@ GRANT INSERT (name, pgpfingerprint, publickey, status, createdat, lastmodified)
 GRANT UPDATE (status, lastmodified) ON investigators TO migapi;
 GRANT USAGE ON SEQUENCE investigators_id_seq TO migapi;
 
+-- readonly user is used for things like expanding targets
+CREATE ROLE migreadonly;
+ALTER ROLE migreadonly WITH NOSUPERUSER INHERIT NOCREATEROLE NOCREATEDB NOLOGIN;
+GRANT SELECT ON actions, agents, agtmodreq, commands, invagtmodperm, modules, signatures TO migreadonly;
+GRANT SELECT (id, name, pgpfingerprint, publickey, status, createdat, lastmodified) ON investigators TO migreadonly;
+GRANT migreadonly TO migapi;
+GRANT migreadonly TO migscheduler;
+
 EOF
 
 chmod 777 $granttmp

doc/.files/createremotedb.sh

@@ -173,6 +173,14 @@ GRANT INSERT (name, pgpfingerprint, publickey, status, createdat, lastmodified)
 GRANT UPDATE (status, lastmodified) ON investigators TO migapi;
 GRANT USAGE ON SEQUENCE investigators_id_seq TO migapi;
 
+-- readonly user is used for things like expanding targets
+CREATE ROLE migreadonly;
+ALTER ROLE migreadonly WITH NOSUPERUSER INHERIT NOCREATEROLE NOCREATEDB NOLOGIN;
+GRANT SELECT ON actions, agents, agtmodreq, commands, invagtmodperm, modules, signatures TO migreadonly;
+GRANT SELECT (id, name, pgpfingerprint, publickey, status, createdat, lastmodified) ON investigators TO migreadonly;
+GRANT migreadonly TO migapi;
+GRANT migreadonly TO migscheduler;
+
 EOF
 
 psql -U $PGUSER -d $PGDATABASE -h $PGHOST -p $PGPORT -c "\i $qfile"

doc/api.rst

@@ -323,10 +323,12 @@ GET <root>/search
 	  Status depends on the type. Below are the available statuses per type:
 
 		- `action`: init, preparing, invalid, inflight, completed
-		- `agent`: heartbeating, upgraded, destroyed, inactive
+		- `agent`: online, upgraded, destroyed, offline
 		- `command`: prepared, sent, success, timeout, cancelled, expired, failed
 		- `investigator`: active, disabled
 
+	- `target`: returns agents that match a target query (only for `agent` type)
+
 	- `threatfamily`: filter results of the threat family of the action, accept
 	  `ILIKE` pattern (only for types `command` and `action`)
 

doc/api.rst.html

@@ -354,12 +354,13 @@ http://localhost:1664/api/v1/investigator/create/</code></pre>
                                     <blockquote>
                                         <ul>
                                             <li><cite>action</cite>: init, preparing, invalid, inflight, completed</li>
-                                            <li><cite>agent</cite>: heartbeating, upgraded, destroyed, inactive</li>
+                                            <li><cite>agent</cite>: online, upgraded, destroyed, offline</li>
                                             <li><cite>command</cite>: prepared, sent, success, timeout, cancelled, expired, failed</li>
                                             <li><cite>investigator</cite>: active, disabled</li>
                                         </ul>
                                     </blockquote>
                                 </li>
+                                <li><cite>target</cite>: returns agents that match a target query (only for <cite>agent</cite> type)</li>
                                 <li><cite>threatfamily</cite>: filter results of the threat family of the action, accept <cite>ILIKE</cite> pattern (only for types <cite>command</cite> and <cite>action</cite>)</li>
                             </ul>
                         </dd>

doc/data.rst.html

@@ -325,6 +325,14 @@ GRANT INSERT (name, pgpfingerprint, publickey, status, createdat, lastmodified)
 GRANT UPDATE (status, lastmodified) ON investigators TO migapi;
 GRANT USAGE ON SEQUENCE investigators_id_seq TO migapi;
 
+-- readonly user is used for things like expanding targets
+CREATE ROLE migreadonly;
+ALTER ROLE migreadonly WITH NOSUPERUSER INHERIT NOCREATEROLE NOCREATEDB NOLOGIN;
+GRANT SELECT ON actions, agents, agtmodreq, commands, invagtmodperm, modules, signatures TO migreadonly;
+GRANT SELECT (id, name, pgpfingerprint, publickey, status, createdat, lastmodified) ON investigators TO migreadonly;
+GRANT migreadonly TO migapi;
+GRANT migreadonly TO migscheduler;
+
 EOF</span>
 
 chmod 777 <span class="nv">$granttmp</span>

src/mig/api/api.go

@@ -336,8 +336,7 @@ func getDashboard(respWriter http.ResponseWriter, request *http.Request) {
 		ctx.Channels.Log <- mig.Log{OpID: opid, Desc: "leaving getDashboard()"}.Debug()
 	}()
 
-	// get summary of agents active in the last 5 minutes
-	sum, err := ctx.DB.SumAgentsByVersion(time.Now().Add(-5 * time.Minute))
+	sum, err := ctx.DB.SumAgentsByVersion()
 	if err != nil {
 		panic(err)
 	}
@@ -345,7 +344,7 @@ func getDashboard(respWriter http.ResponseWriter, request *http.Request) {
 	if err != nil {
 		panic(err)
 	}
-	double, err := ctx.DB.CountDoubleAgents(time.Now().Add(-5 * time.Minute))
+	double, err := ctx.DB.CountDoubleAgents()
 	if err != nil {
 		panic(err)
 	}

src/mig/api/search.go

@@ -91,6 +91,8 @@ func search(respWriter http.ResponseWriter, request *http.Request) {
 			}
 		case "status":
 			p.Status = request.URL.Query()["status"][0]
+		case "target":
+			p.Target = request.URL.Query()["target"][0]
 		case "threatfamily":
 			p.ThreatFamily = request.URL.Query()["threatfamily"][0]
 		}
@@ -103,7 +105,11 @@ func search(respWriter http.ResponseWriter, request *http.Request) {
 		case "action":
 			results, err = ctx.DB.SearchActions(p)
 		case "agent":
-			results, err = ctx.DB.SearchAgents(p)
+			if p.Target != "" {
+				results, err = ctx.DB.ActiveAgentsByTarget(p.Target)
+			} else {
+				results, err = ctx.DB.SearchAgents(p)
+			}
 		case "command":
 			results, err = ctx.DB.SearchCommands(p, doFoundAnything)
 		case "investigator":

src/mig/client/client.go

@@ -542,3 +542,30 @@ func (cli Client) SignAction(a mig.Action) (signed_action mig.Action, err error)
 	signed_action = a
 	return
 }
+
+// EvaluateAgentTarget runs a search against the api to find all agents that match an action target string
+func (cli Client) EvaluateAgentTarget(target string) (agents []mig.Agent, err error) {
+	defer func() {
+		if e := recover(); e != nil {
+			err = fmt.Errorf("EvaluateAgentTarget() -> %v", e)
+		}
+	}()
+	query := "search?type=agent&target=" + url.QueryEscape(target)
+	resource, err := cli.GetAPIResource(query)
+	if err != nil {
+		panic(err)
+	}
+	for _, item := range resource.Collection.Items {
+		for _, data := range item.Data {
+			if data.Name != "agent" {
+				continue
+			}
+			agt, err := ValueToAgent(data.Value)
+			if err != nil {
+				panic(err)
+			}
+			agents = append(agents, agt)
+		}
+	}
+	return
+}

src/mig/client/console/action_launcher.go

@@ -44,7 +44,7 @@ func actionLauncher(tpl mig.Action, cli client.Client) (err error) {
 	}
 	hasTimes := false
 	hasSignatures := false
-
+	hasEvaluatedTarget := false
 	fmt.Println("Type \x1b[32;1mexit\x1b[0m or press \x1b[32;1mctrl+d\x1b[0m to leave. \x1b[32;1mhelp\x1b[0m may help.")
 	prompt := "\x1b[33;1mlauncher>\x1b[0m "
 	for {
@@ -183,6 +183,26 @@ times			show the various timestamps of the action
 				fmt.Println("Action has no target. Define one using 'settarget <target>'")
 				break
 			}
+			if !hasEvaluatedTarget {
+				agents, err := cli.EvaluateAgentTarget(a.Target)
+				if err != nil {
+					panic(err)
+				}
+				count := len(agents)
+				if count == 0 {
+					fmt.Println("0 agents match this target. launch aborted")
+					break
+				}
+				fmt.Printf("%d agents will be targetted by search \"%s\"\n", count, a.Target)
+				input, err = readline.String("continue? (y/n)> ")
+				if err != nil {
+					panic(err)
+				}
+				if input != "y" {
+					fmt.Println("launch aborted")
+					break
+				}
+			}
 			if !hasTimes {
 				fmt.Printf("Times are not defined. Setting validity from now until +%s\n", defaultExpiration)
 				// for immediate execution, set validity one minute in the past
@@ -251,6 +271,13 @@ times			show the various timestamps of the action
 				break
 			}
 			a.Target = strings.Join(orders[1:], " ")
+			agents, err := cli.EvaluateAgentTarget(a.Target)
+			if err != nil {
+				fmt.Println(err)
+				break
+			}
+			fmt.Printf("%d agents will be targetted. To get the list, use 'expandtarget'\n", len(agents))
+			hasEvaluatedTarget = true
 		case "settimes":
 			// set the dates
 			if len(orders) != 3 {

src/mig/client/console/action_reader.go

@@ -283,7 +283,7 @@ Times          valid from %s until %s
 	}
 	fmt.Printf("\n")
 	fmt.Printf("Counters       sent=%d; done=%d; in flight=%d\n"+
-		"               sucess=%d; cancelled=%d; expired=%d; failed=%d; timeout=%d\n",
+		"               success=%d; cancelled=%d; expired=%d; failed=%d; timeout=%d\n",
 		a.Counters.Sent, a.Counters.Done, a.Counters.InFlight, a.Counters.Success,
 		a.Counters.Cancelled, a.Counters.Expired, a.Counters.Failed, a.Counters.TimeOut)
 	return

src/mig/database/agents.go

@@ -101,7 +101,7 @@ func (db *DB) InsertAgent(agt mig.Agent) (err error) {
 // UpdateAgentHeartbeat updates the heartbeat timestamp of an agent in the database
 func (db *DB) UpdateAgentHeartbeat(agt mig.Agent) (err error) {
 	_, err = db.c.Exec(`UPDATE agents
-		SET heartbeattime=$2 WHERE id=$1`, agt.ID, agt.HeartBeatTS)
+		SET status='online', heartbeattime=$2 WHERE id=$1`, agt.ID, agt.HeartBeatTS)
 	if err != nil {
 		return fmt.Errorf("Failed to update agent in database: '%v'", err)
 	}
@@ -114,7 +114,7 @@ func (db *DB) InsertOrUpdateAgent(agt mig.Agent) (err error) {
 	agent, err := db.AgentByQueueAndPID(agt.QueueLoc, agt.PID)
 	if err != nil {
 		agt.DestructionTime = time.Date(9998, time.January, 11, 11, 11, 11, 11, time.UTC)
-		agt.Status = "heartbeating"
+		agt.Status = "online"
 		// create a new agent
 		return db.InsertAgent(agt)
 	} else {
@@ -129,8 +129,7 @@ func (db *DB) ActiveAgentsByQueue(queueloc string, pointInTime time.Time) (agent
 	rows, err := db.c.Query(`SELECT agents.id, agents.name, agents.queueloc, agents.os,
 		agents.version, agents.pid, agents.starttime, agents.heartbeattime, agents.status
 		FROM agents
-		WHERE agents.heartbeattime >= $1 AND agents.heartbeattime <= NOW()
-		AND agents.queueloc=$2`, pointInTime, queueloc)
+		WHERE agents.heartbeattime > $1 AND agents.queueloc=$2`, pointInTime, queueloc)
 	if err != nil {
 		err = fmt.Errorf("Error while finding agents: '%v'", err)
 		return
@@ -152,15 +151,31 @@ func (db *DB) ActiveAgentsByQueue(queueloc string, pointInTime time.Time) (agent
 	return
 }
 
-// ActiveAgentsByTarget runs a search for all agents that match a given target string
-func (db *DB) ActiveAgentsByTarget(target string, pointInTime time.Time) (agents []mig.Agent, err error) {
-	rows, err := db.c.Query(`SELECT DISTINCT ON (queueloc) id, name, queueloc, os, version, pid,
+// ActiveAgentsByTarget runs a search for all agents that match a given target string.
+// For safety, it does so in a transaction that runs as a readonly user.
+func (db *DB) ActiveAgentsByTarget(target string) (agents []mig.Agent, err error) {
+	// save current user
+	var dbuser string
+	err = db.c.QueryRow("SELECT CURRENT_USER").Scan(&dbuser)
+	if err != nil {
+		return
+	}
+	txn, err := db.c.Begin()
+	if err != nil {
+		return
+	}
+	_, err = txn.Exec(`SET ROLE migreadonly`)
+	if err != nil {
+		_ = txn.Rollback()
+		return
+	}
+	rows, err := txn.Query(`SELECT DISTINCT ON (queueloc) id, name, queueloc, os, version, pid,
 		starttime, destructiontime, heartbeattime, status
 		FROM agents
-		WHERE agents.heartbeattime >= $1 AND agents.heartbeattime <= NOW()
-		AND (`+target+`)
-		ORDER BY agents.queueloc, agents.heartbeattime DESC`, pointInTime)
+		WHERE agents.status = 'online' AND (` + target + `)
+		ORDER BY agents.queueloc, agents.heartbeattime DESC`)
 	if err != nil {
+		_ = txn.Rollback()
 		err = fmt.Errorf("Error while finding agents: '%v'", err)
 		return
 	}
@@ -179,6 +194,16 @@ func (db *DB) ActiveAgentsByTarget(target string, pointInTime time.Time) (agents
 	if err := rows.Err(); err != nil {
 		err = fmt.Errorf("Failed to complete database query: '%v'", err)
 	}
+	_, err = txn.Exec(`SET ROLE ` + dbuser)
+	if err != nil {
+		_ = txn.Rollback()
+		return
+	}
+	err = txn.Commit()
+	if err != nil {
+		_ = txn.Rollback()
+		return
+	}
 	return
 }
 
@@ -211,10 +236,9 @@ type AgentsSum struct {
 }
 
 // SumAgentsByVersion retrieves a sum of agents grouped by version
-func (db *DB) SumAgentsByVersion(pointInTime time.Time) (sum []AgentsSum, err error) {
+func (db *DB) SumAgentsByVersion() (sum []AgentsSum, err error) {
 	rows, err := db.c.Query(`SELECT COUNT(*), version FROM agents
-		WHERE agents.heartbeattime >= $1 AND agents.heartbeattime <= NOW()
-		GROUP BY version`, pointInTime)
+		WHERE agents.status='online' GROUP BY version`)
 	if err != nil {
 		err = fmt.Errorf("Error while counting agents: '%v'", err)
 		return
@@ -237,7 +261,7 @@ func (db *DB) SumAgentsByVersion(pointInTime time.Time) (sum []AgentsSum, err er
 
 // CountNewAgents retrieves a count of agents that started after `pointInTime`
 func (db *DB) CountNewAgents(pointInTime time.Time) (sum float64, err error) {
-	err = db.c.QueryRow(`SELECT COUNT(name) FROM agents
+	err = db.c.QueryRow(`SELECT COUNT(DISTINCT(queueloc)) FROM agents
 		WHERE starttime >= $1 AND starttime <= NOW()`, pointInTime).Scan(&sum)
 	if err != nil {
 		err = fmt.Errorf("Error while counting agents: '%v'", err)
@@ -250,13 +274,13 @@ func (db *DB) CountNewAgents(pointInTime time.Time) (sum float64, err error) {
 }
 
 // CountDoubleAgents counts the number of endpoints that run more than one agent
-func (db *DB) CountDoubleAgents(pointInTime time.Time) (sum float64, err error) {
+func (db *DB) CountDoubleAgents() (sum float64, err error) {
 	err = db.c.QueryRow(`SELECT COUNT(DISTINCT(queueloc)) FROM agents
 		WHERE queueloc IN (
 			SELECT queueloc FROM agents
-			WHERE heartbeattime >= $1
+			WHERE heartbeattime >= NOW() - INTERVAL '10 minutes'
 			GROUP BY queueloc HAVING count(queueloc) > 1
-		)`, pointInTime).Scan(&sum)
+		)`).Scan(&sum)
 	if err != nil {
 		err = fmt.Errorf("Error while counting double agents: '%v'", err)
 		return
@@ -286,3 +310,13 @@ func (db *DB) CountDisappearedAgents(seenSince, activeSince time.Time) (sum floa
 	}
 	return
 }
+
+// MarkOfflineAgents updates the status of agents that have not sent a heartbeat since pointInTime
+func (db *DB) MarkOfflineAgents(pointInTime time.Time) (err error) {
+	_, err = db.c.Exec(`UPDATE agents SET status='offline'
+		WHERE heartbeattime<$1 AND status!='offline'`, pointInTime)
+	if err != nil {
+		return fmt.Errorf("Failed to mark agents as offline in database: '%v'", err)
+	}
+	return
+}

src/mig/database/searches.go

@@ -31,6 +31,7 @@ type SearchParameters struct {
 	Limit            float64   `json:"limit"`
 	Report           string    `json:"report"`
 	Status           string    `json:"status"`
+	Target           string    `json:"target"`
 	ThreatFamily     string    `json:"threatfamily"`
 	Type             string    `json:"type"`
 }

src/mig/scheduler/collector.go

@@ -46,7 +46,10 @@ func spoolInspection(ctx Context) (err error) {
 	if err != nil {
 		panic(err)
 	}
-
+	err = timeoutAgents(ctx)
+	if err != nil {
+		panic(err)
+	}
 	return
 }
 
@@ -209,3 +212,23 @@ func cleanDir(ctx Context, targetDir string) (err error) {
 	dir.Close()
 	return
 }
+
+// timeoutAgents updates the status of agents that are no longer heartbeating to "offline"
+func timeoutAgents(ctx Context) (err error) {
+	defer func() {
+		if e := recover(); e != nil {
+			err = fmt.Errorf("timeoutAgents() -> %v", e)
+		}
+		ctx.Channels.Log <- mig.Log{OpID: ctx.OpID, Desc: "leaving timeoutAgents()"}.Debug()
+	}()
+	timeOutPeriod, err := time.ParseDuration(ctx.Agent.TimeOut)
+	if err != nil {
+		panic(err)
+	}
+	pointInTime := time.Now().Add(-timeOutPeriod)
+	err = ctx.DB.MarkOfflineAgents(pointInTime)
+	if err != nil {
+		panic(err)
+	}
+	return
+}

src/mig/scheduler/scheduler.go

@@ -364,7 +364,10 @@ func processNewAction(actionPath string, ctx Context) (err error) {
 		return
 	}
 	// find target agents for the action
-	agents, err := getTargetAgents(action, ctx)
+	agents, err := ctx.DB.ActiveAgentsByTarget(action.Target)
+	if err != nil {
+		panic(err)
+	}
 	action.Counters.Sent = len(agents)
 	if action.Counters.Sent == 0 {
 		err = fmt.Errorf("No agents found for target '%s'. invalidating action.", action.Target)
@@ -437,26 +440,6 @@ func processNewAction(actionPath string, ctx Context) (err error) {
 	return
 }
 
-// getTargetAgents retrieves an array of agents from the target of an action
-func getTargetAgents(action mig.Action, ctx Context) (agents []mig.Agent, err error) {
-	defer func() {
-		if e := recover(); e != nil {
-			err = fmt.Errorf("getTargetAgents() -> %v", e)
-		}
-		ctx.Channels.Log <- mig.Log{OpID: ctx.OpID, ActionID: action.ID, Desc: "leaving getTargetAgents()"}.Debug()
-	}()
-	timeOutPeriod, err := time.ParseDuration(ctx.Agent.TimeOut)
-	if err != nil {
-		panic(err)
-	}
-	pointInTime := time.Now().Add(-timeOutPeriod)
-	agents, err = ctx.DB.ActiveAgentsByTarget(action.Target, pointInTime)
-	if err != nil {
-		panic(err)
-	}
-	return
-}
-
 func createCommand(ctx Context, action mig.Action, agent mig.Agent, emptyResults []mig.ModuleResult) (err error) {
 	cmdid := mig.GenID()
 	defer func() {