Policy Result Control in osquery-perf (#17649)

Author: mostlikelee

cmd/osquery-perf/README.md

@@ -56,7 +56,21 @@ go run agent.go --enroll_secret hgh4hk3434l2jjf --os_templates ubuntu_22.04,wind
 
 would start 3 Ubuntu hosts and 3 Windows hosts. See the `os_templates` flag description in `go run agent.go --help` for the list of supported template names.
 
-### Running Locally (Development Environment)
+## Controlling Agent Behavior From the Fleet UI
+
+### Specify Query Results
+
+Using the naming convention `MyQuery_10` (name separated by `_number`) will instruct agents to
+return 10 rows for that query
+
+### Control policy pass/fail per policy
+
+In the Policy SQL:
+
+- `select 1` will instruct agents to send back only passing responses
+- `select 0` will instruct agents to send back only failing responses
+
+## Running Locally (Development Environment)
 
 First, ensure your Fleet local development environment is up and running. Refer to [Building Fleet](../../docs/Contributing/Building-Fleet.md) for details. Once this is done:
 
@@ -68,13 +82,13 @@ Alternatively, you can retrieve the enroll secret from the command-line using `f
 
 The agent will start. You can connect to MySQL to view changes made to the development database by the agent (e.g., at the terminal, with `docker-compose exec mysql mysql -uroot -ptoor -Dfleet`). Remember that frequency of the reported data depends on the configuration of the Fleet instance, so you may want to start it with shorter delays for some cases and enable debug logging (e.g., `./build/fleet serve --dev --logging_debug --osquery_detail_update_interval 1m`).
 
-### Resource Limits
+## Resource Limits
 
 On many systems, trying to simulate a large number of hosts will result in hitting system resource limits (such as number of open file descriptors).
 
 If you see errors such as `dial tcp: lookup localhost: no such host` or `read: connection reset by peer`, try increasing these limits.
 
-#### macOS
+### macOS
 
 Run the following command in the shell before running the Fleet server _and_ before running `agent.go` (run it once in each shell):
 

cmd/osquery-perf/agent.go

@@ -1093,6 +1093,7 @@ func (a *agent) orbitEnroll() error {
 	return nil
 }
 
+// This is an osquery enroll as opposed to an orbit enroll
 func (a *agent) enroll(i int, onlyAlreadyEnrolled bool) error {
 	a.nodeKey = a.nodeKeyManager.Get(i)
 	if a.nodeKey != "" {
@@ -1181,6 +1182,9 @@ func (a *agent) config() error {
 			q := scheduledQuery{}
 			q.packName = packName
 			q.Name = queryName
+
+			// This allows us to set the number of rows returned by the query
+			// by appending a number to the query name, e.g. "queryName_10"
 			q.numRows = 1
 			parts := strings.Split(q.Name, "_")
 			if len(parts) == 2 {
@@ -1190,6 +1194,7 @@ func (a *agent) config() error {
 				}
 				q.numRows = uint(num)
 			}
+
 			q.ScheduleInterval = m["interval"].(float64)
 			q.Query = m["query"].(string)
 
@@ -1421,6 +1426,21 @@ func (a *agent) genLastOpenedAt(count *int) *time.Time {
 }
 
 func (a *agent) runPolicy(query string) []map[string]string {
+	// Used to control the pass or fail of a policy
+	// in the UI by setting the query to "select 1"(pass)
+	// or "select 0"(fail)
+	query = strings.TrimRight(query, ";")
+	query = strings.ToLower(query)
+
+	switch query {
+	case "select 1":
+		return []map[string]string{
+			{"1": "1"},
+		}
+	case "select 0":
+		return []map[string]string{}
+	}
+
 	if rand.Float64() <= a.policyPassProb {
 		return []map[string]string{
 			{"1": "1"},
@@ -1994,7 +2014,7 @@ func main() {
 		// Flag logger_tls_period defines how often to check for sending scheduled query results.
 		// osquery-perf will send log requests with results only if there are scheduled queries configured AND it's their time to run.
 		logInterval         = flag.Duration("logger_tls_period", 10*time.Second, "Interval for scheduled queries log requests")
-		queryInterval       = flag.Duration("query_interval", 10*time.Second, "Interval for live query requests")
+		queryInterval       = flag.Duration("query_interval", 10*time.Second, "Interval for distributed query requests")
 		mdmCheckInInterval  = flag.Duration("mdm_check_in_interval", 10*time.Second, "Interval for performing MDM check-ins (applies to both macOS and Windows)")
 		onlyAlreadyEnrolled = flag.Bool("only_already_enrolled", false, "Only start agents that are already enrolled")
 		nodeKeyFile         = flag.String("node_key_file", "", "File with node keys to use")