HTTP Controller updates for label queries (#96)

Add controller methods for: * Retrieving label queries * Storing results of label queries
2016-08-24 16:50:28 -07:00
parent 3f81dda638
commit 41fe404ef1
8 changed files with 584 additions and 24 deletions
@@ -32,8 +32,9 @@ type OsqueryStatusHandler interface {
 // It can be configured in a `main` function to bind the appropriate handlers
 // and it's methods can be attached to routes.
 type OsqueryHandler struct {
-	ResultHandler OsqueryResultHandler
-	StatusHandler OsqueryStatusHandler
+	LabelQueryInterval time.Duration
+	ResultHandler      OsqueryResultHandler
+	StatusHandler      OsqueryStatusHandler
 }

 // Basic implementation of the `OsqueryResultHandler` and
@@ -64,8 +65,27 @@ type OsqueryEnrollPostBody struct {
 	HostIdentifier string `json:"host_identifier" validate:"required"`
 }

+// OsqueryConfigPostBody is the generic osquery config endpoint request body
+// structure. Typically the node key and action will be parsed, and then the
+// Data JSON will be unmarshalled into one of the more specific OsqueryConfig*
+// structs.
 type OsqueryConfigPostBody struct {
-	NodeKey string `json:"node_key" validate:"required"`
+	NodeKey string           `json:"node_key" validate:"required"`
+	Action  string           `json:"action" validate:"required"`
+	Data    *json.RawMessage `json:"data" validate:"required"`
+}
+
+// OsqueryConfigDetail is the expected extra information provided with an
+// initial config request. This data will be used to determine which label
+// queries are supplied to the host.
+type OsqueryConfigDetail struct {
+	Platform string `json:"platform" validate:"required"`
+}
+
+// OsqueryConfigQueryResults contains the final information needed to generate a
+// config for the host. It containst the results of the label queries.
+type OsqueryConfigQueryResults struct {
+	Results map[string]bool `json:"results" validate:"required"`
 }

 type OsqueryLogPostBody struct {
@@ -141,31 +161,94 @@ func OsqueryEnroll(c *gin.Context) {
 		})
 }

-func OsqueryConfig(c *gin.Context) {
+func (h *OsqueryHandler) handleConfigDetail(db kolide.OsqueryStore, host *kolide.Host, data json.RawMessage) (map[string]string, error) {
+	var detail OsqueryConfigDetail
+	if err := json.Unmarshal(data, &detail); err != nil {
+		return nil, errors.NewFromError(err, http.StatusBadRequest, "JSON parse error")
+	}
+	if err := validateStruct(&detail); err != nil {
+		return nil, err
+	}
+
+	// Update fields from detail (only save if there are changes)
+	if host.Platform != detail.Platform {
+		host.Platform = detail.Platform
+		if err := db.SaveHost(host); err != nil {
+			return nil, err
+		}
+	}
+
+	return kolide.LabelQueriesForHost(db, host, h.LabelQueryInterval)
+}
+
+func (h *OsqueryHandler) handleConfigQueryResults(db kolide.OsqueryStore, host *kolide.Host, data json.RawMessage) error {
+	var results OsqueryConfigQueryResults
+	if err := json.Unmarshal(data, &results); err != nil {
+		return errors.NewFromError(err, http.StatusBadRequest, "JSON parse error")
+	}
+	if err := validateStruct(&results); err != nil {
+		return err
+	}
+
+	return db.RecordLabelQueryExecutions(host, results.Results, time.Now())
+}
+
+// Endpoint used by the osqueryd TLS logger plugin
+func (h *OsqueryHandler) OsqueryConfig(c *gin.Context) {
 	var body OsqueryConfigPostBody
 	err := ParseAndValidateJSON(c, &body)
 	if err != nil {
 		errors.ReturnOsqueryError(c, err)
 		return
 	}
-	logrus.Debugf("OsqueryConfig: %s", body.NodeKey)

-	c.JSON(http.StatusOK,
-		gin.H{
-			"schedule": map[string]map[string]interface{}{
-				"time": {
-					"query":    "select * from time;",
-					"interval": 1,
-				},
-			},
-			"node_invalid": false,
-		})
+	db := GetDB(c)
+
+	host, err := db.AuthenticateHost(body.NodeKey)
+	if err != nil {
+		errors.ReturnOsqueryError(c, err)
+		return
+	}
+
+	if body.Data == nil {
+		errors.ReturnOsqueryError(c,
+			errors.New("Missing body data", "Got nil pointer for data in config"),
+		)
+	}
+
+	var res map[string]string
+	switch body.Action {
+	case "request":
+		res, err = h.handleConfigDetail(db, host, *body.Data)
+		if err != nil {
+			c.JSON(http.StatusOK, res)
+			return
+		}
+
+	case "results":
+		err = h.handleConfigQueryResults(db, host, *body.Data)
+		// Now we should be able to calculate the appropriate config
+
+	default:
+		err = errors.NewWithStatus(
+			errors.StatusUnprocessableEntity,
+			"Unknown config request action",
+			fmt.Sprintf("Unknown config request action: %s", body.Action),
+		)
+	}
+
+	if err != nil {
+		errors.ReturnOsqueryError(c, err)
+		return
+	}
+
+	// TODO: Generate the config and return it here!
 }

 // Unmarshal the status logs before sending them to the status log handler
-func (h *OsqueryHandler) handleStatusLogs(data *json.RawMessage, nodeKey string) error {
+func (h *OsqueryHandler) handleStatusLogs(data json.RawMessage, nodeKey string) error {
 	var statuses []OsqueryStatusLog
-	if err := json.Unmarshal(*data, &statuses); err != nil {
+	if err := json.Unmarshal(data, &statuses); err != nil {
 		return errors.NewFromError(err, http.StatusBadRequest, "JSON parse error")
 	}
 	// Perhaps we should validate the unmarshalled status log
@@ -181,9 +264,9 @@ func (h *OsqueryHandler) handleStatusLogs(data *json.RawMessage, nodeKey string)
 }

 // Unmarshal the result logs before sending them to the result log handler
-func (h *OsqueryHandler) handleResultLogs(data *json.RawMessage, nodeKey string) error {
+func (h *OsqueryHandler) handleResultLogs(data json.RawMessage, nodeKey string) error {
 	var results []OsqueryResultLog
-	if err := json.Unmarshal(*data, &results); err != nil {
+	if err := json.Unmarshal(data, &results); err != nil {
 		return errors.NewFromError(err, http.StatusBadRequest, "JSON parse error")
 	}
 	// Perhaps we should validate the unmarshalled result log
@@ -215,12 +298,18 @@ func (h *OsqueryHandler) OsqueryLog(c *gin.Context) {
 		return
 	}

+	if body.Data == nil {
+		errors.ReturnOsqueryError(c,
+			errors.New("Missing body data", "Got nil pointer for data in log"),
+		)
+	}
+
 	switch body.LogType {
 	case "status":
-		err = h.handleStatusLogs(body.Data, body.NodeKey)
+		err = h.handleStatusLogs(*body.Data, body.NodeKey)

 	case "result":
-		err = h.handleResultLogs(body.Data, body.NodeKey)
+		err = h.handleResultLogs(*body.Data, body.NodeKey)

 	default:
 		err = errors.NewWithStatus(
@@ -5,7 +5,10 @@ import (
 	"fmt"
 	"net/http"
 	"testing"
+	"time"

+	"github.com/golang/mock/gomock"
+	"github.com/kolide/kolide-ose/errors"
 	"github.com/kolide/kolide-ose/kolide"
 	"github.com/stretchr/testify/assert"
 )
@@ -842,3 +845,216 @@ func TestDeleteQueryFromPack(t *testing.T) {
 	assert.Nil(t, err)
 	assert.Len(t, p.Queries, len(queriesInPack)-1)
 }
+
+func TestHandleConfigDetail(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	db := kolide.NewMockOsqueryStore(ctrl)
+
+	detail := OsqueryConfigDetail{Platform: "darwin"}
+
+	detailBytes, err := json.Marshal(detail)
+	assert.NoError(t, err)
+	detailJSON := json.RawMessage(detailBytes)
+
+	host := &kolide.Host{
+		NodeKey: "fake_key",
+	}
+
+	expectHost := &kolide.Host{
+		NodeKey:  "fake_key",
+		Platform: "darwin",
+	}
+
+	handler := OsqueryHandler{
+		LabelQueryInterval: time.Minute,
+	}
+
+	expectQueries := map[string]string{
+		"1": "query1",
+		"3": "query3",
+	}
+
+	db.EXPECT().SaveHost(expectHost)
+	db.EXPECT().LabelQueriesForHost(expectHost, gomock.Any()).
+		Return(expectQueries, nil).
+		Do(func(_ *kolide.Host, cutoff time.Time) {
+			// Check that the cutoff is in the correct interval
+			expectCutoff := time.Now().Add(-handler.LabelQueryInterval)
+			allowedDelta := 5 * time.Second
+			assert.WithinDuration(t, expectCutoff, cutoff, allowedDelta)
+		})
+
+	res, err := handler.handleConfigDetail(db, host, detailJSON)
+	assert.NoError(t, err)
+	assert.Equal(t, expectQueries, res)
+
+}
+
+func TestHandleConfigDetailNoSave(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	db := kolide.NewMockOsqueryStore(ctrl)
+
+	detail := OsqueryConfigDetail{Platform: "darwin"}
+
+	detailBytes, err := json.Marshal(detail)
+	assert.NoError(t, err)
+	detailJSON := json.RawMessage(detailBytes)
+
+	host := &kolide.Host{
+		NodeKey:  "fake_key",
+		Platform: "darwin",
+	}
+
+	expectHost := &kolide.Host{
+		NodeKey:  "fake_key",
+		Platform: "darwin",
+	}
+
+	handler := OsqueryHandler{
+		LabelQueryInterval: time.Hour,
+	}
+
+	expectQueries := map[string]string{}
+
+	// Note that we don't expect a call to save because the platform did
+	// not change
+	db.EXPECT().LabelQueriesForHost(expectHost, gomock.Any()).
+		Return(expectQueries, nil).
+		Do(func(_ *kolide.Host, cutoff time.Time) {
+			// Check that the cutoff is in the correct interval
+			expectCutoff := time.Now().Add(-handler.LabelQueryInterval)
+			allowedDelta := 5 * time.Second
+			assert.WithinDuration(t, expectCutoff, cutoff, allowedDelta)
+		})
+
+	res, err := handler.handleConfigDetail(db, host, detailJSON)
+	assert.NoError(t, err)
+	assert.Equal(t, expectQueries, res)
+
+}
+
+func marshalRawMessage(t *testing.T, obj interface{}) json.RawMessage {
+	objBytes, err := json.Marshal(obj)
+	assert.NoError(t, err)
+	objJSON := json.RawMessage(objBytes)
+	return objJSON
+}
+
+func TestHandleConfigDetailError(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	db := kolide.NewMockOsqueryStore(ctrl)
+
+	detail := OsqueryConfigDetail{Platform: "darwin"}
+
+	detailBytes, err := json.Marshal(detail)
+	assert.NoError(t, err)
+	detailJSON := json.RawMessage(detailBytes)
+
+	host := &kolide.Host{
+		NodeKey:  "fake_key",
+		Platform: "darwin",
+	}
+
+	expectHost := &kolide.Host{
+		NodeKey:  "fake_key",
+		Platform: "darwin",
+	}
+
+	handler := OsqueryHandler{
+		LabelQueryInterval: time.Hour,
+	}
+
+	// The DB call should error in this test
+	db.EXPECT().LabelQueriesForHost(expectHost, gomock.Any()).
+		Return(nil, errors.New("public", "private")).
+		Do(func(_ *kolide.Host, cutoff time.Time) {
+			// Check that the cutoff is in the correct interval
+			expectCutoff := time.Now().Add(-handler.LabelQueryInterval)
+			allowedDelta := 5 * time.Second
+			assert.WithinDuration(t, expectCutoff, cutoff, allowedDelta)
+		})
+
+	res, err := handler.handleConfigDetail(db, host, detailJSON)
+	assert.Error(t, err)
+	assert.Nil(t, res)
+
+}
+
+func TestHandleConfigQueryResults(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	db := kolide.NewMockOsqueryStore(ctrl)
+
+	results := OsqueryConfigQueryResults{
+		Results: map[string]bool{
+			"1": true,
+			"3": false,
+			"4": true,
+		},
+	}
+
+	resultsJSON := marshalRawMessage(t, results)
+
+	host := &kolide.Host{
+		NodeKey:  "fake_key",
+		Platform: "darwin",
+	}
+
+	handler := OsqueryHandler{
+		LabelQueryInterval: time.Hour,
+	}
+
+	db.EXPECT().RecordLabelQueryExecutions(host, results.Results, gomock.Any()).
+		Return(nil).
+		Do(func(_ *kolide.Host, _ map[string]bool, recordTime time.Time) {
+			// Check that the cutoff is in the correct interval
+			allowedDelta := 5 * time.Second
+			assert.WithinDuration(t, time.Now(), recordTime, allowedDelta)
+		})
+
+	assert.NoError(t, handler.handleConfigQueryResults(db, host, resultsJSON))
+}
+
+func TestHandleConfigQueryResultsError(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	db := kolide.NewMockOsqueryStore(ctrl)
+
+	results := OsqueryConfigQueryResults{
+		Results: map[string]bool{
+			"1": true,
+			"3": false,
+			"4": true,
+		},
+	}
+
+	resultsJSON := marshalRawMessage(t, results)
+
+	host := &kolide.Host{
+		NodeKey:  "fake_key",
+		Platform: "darwin",
+	}
+
+	handler := OsqueryHandler{
+		LabelQueryInterval: time.Hour,
+	}
+
+	// DB errors this time
+	db.EXPECT().RecordLabelQueryExecutions(host, results.Results, gomock.Any()).
+		Return(errors.New("public", "private")).
+		Do(func(_ *kolide.Host, _ map[string]bool, recordTime time.Time) {
+			// Check that the cutoff is in the correct interval
+			allowedDelta := 5 * time.Second
+			assert.WithinDuration(t, time.Now(), recordTime, allowedDelta)
+		})
+
+	assert.Error(t, handler.handleConfigQueryResults(db, host, resultsJSON))
+}
@@ -100,6 +100,10 @@ func ParseAndValidateJSON(c *gin.Context, obj interface{}) error {
 		return err
 	}

+	return validateStruct(obj)
+}
+
+func validateStruct(obj interface{}) error {
 	return validate.Struct(obj)
 }

@@ -209,7 +213,7 @@ func CreateServer(ds datastore.Datastore, pool kolide.SMTPConnectionPool, w io.W
 	}

 	osq.POST("/enroll", OsqueryEnroll)
-	osq.POST("/config", OsqueryConfig)
+	osq.POST("/config", osqueryHandler.OsqueryConfig)
 	osq.POST("/log", osqueryHandler.OsqueryLog)
 	osq.POST("/distributed/read", OsqueryDistributedRead)
 	osq.POST("/distributed/write", OsqueryDistributedWrite)