From d6f325f636a9fe41411a3fed5d14d70accc2e9fc Mon Sep 17 00:00:00 2001 From: Rebecca Mahany-Horton Date: Fri, 6 Oct 2023 13:26:00 -0400 Subject: [PATCH] Include messages in distributed results (#115) --- plugin/distributed/distributed.go | 11 ++++++++++- plugin/distributed/distributed_test.go | 27 ++++++++++++++++++++------ 2 files changed, 31 insertions(+), 7 deletions(-) diff --git a/plugin/distributed/distributed.go b/plugin/distributed/distributed.go index 8d2c0e6..ccecc57 100644 --- a/plugin/distributed/distributed.go +++ b/plugin/distributed/distributed.go @@ -48,6 +48,8 @@ type Result struct { Rows []map[string]string `json:"rows"` // QueryStats are the stats about the execution of the given query QueryStats *Stats `json:"stats"` + // Message is the message string indicating the status of the query + Message string `json:"message"` } // WriteResultsFunc writes the results of the executed distributed queries. The @@ -143,6 +145,7 @@ type ResultsStruct struct { Queries map[string][]map[string]string `json:"queries"` Statuses map[string]OsqueryInt `json:"statuses"` Stats map[string]Stats `json:"stats"` + Messages map[string]string `json:"messages"` } // Stats holds performance stats about the execution of a given query. @@ -158,12 +161,14 @@ func (rs *ResultsStruct) UnmarshalJSON(buff []byte) error { emptyRow := []map[string]string{} rs.Queries = make(map[string][]map[string]string) rs.Statuses = make(map[string]OsqueryInt) + rs.Messages = make(map[string]string) // Queries can be []map[string]string OR an empty string // so we need to deal with an interface to accommodate two types intermediate := struct { Queries map[string]interface{} `json:"queries"` Statuses map[string]OsqueryInt `json:"statuses"` Stats map[string]Stats `json:"stats"` + Messages map[string]string `json:"messages"` }{} if err := json.Unmarshal(buff, &intermediate); err != nil { return err @@ -192,8 +197,9 @@ func (rs *ResultsStruct) UnmarshalJSON(buff []byte) error { return fmt.Errorf("results for %q unknown type", queryName) } } - // Stats don't require any format changes + // Stats and messages don't require any format changes rs.Stats = intermediate.Stats + rs.Messages = intermediate.Messages return nil } @@ -209,6 +215,9 @@ func (rs *ResultsStruct) toResults() ([]Result, error) { if stats, ok := rs.Stats[queryName]; ok { result.QueryStats = &stats } + if msg, ok := rs.Messages[queryName]; ok { + result.Message = msg + } results = append(results, result) } return results, nil diff --git a/plugin/distributed/distributed_test.go b/plugin/distributed/distributed_test.go index 6a3a40e..ba31811 100644 --- a/plugin/distributed/distributed_test.go +++ b/plugin/distributed/distributed_test.go @@ -63,9 +63,9 @@ func TestDistributedPlugin(t *testing.T) { // Ensure correct ordering for comparison sort.Slice(results, func(i, j int) bool { return results[i].QueryName < results[j].QueryName }) assert.Equal(t, []Result{ - {"query1", 0, []map[string]string{{"iso_8601": "2017-07-10T22:08:40Z"}}, &Stats{WallTimeMs: 1, UserTime: 1, SystemTime: 1, Memory: 1}}, - {"query2", 0, []map[string]string{{"version": "2.4.0"}}, &Stats{WallTimeMs: 2, UserTime: 2, SystemTime: 2, Memory: 2}}, - {"query3", 1, []map[string]string{}, &Stats{WallTimeMs: 3, UserTime: 3, SystemTime: 3, Memory: 3}}, + {"query1", 0, []map[string]string{{"iso_8601": "2017-07-10T22:08:40Z"}}, &Stats{WallTimeMs: 1, UserTime: 1, SystemTime: 1, Memory: 1}, ""}, + {"query2", 0, []map[string]string{{"version": "2.4.0"}}, &Stats{WallTimeMs: 2, UserTime: 2, SystemTime: 2, Memory: 2}, ""}, + {"query3", 1, []map[string]string{}, &Stats{WallTimeMs: 3, UserTime: 3, SystemTime: 3, Memory: 3}, ""}, }, results) @@ -78,9 +78,24 @@ func TestDistributedPlugin(t *testing.T) { // Ensure correct ordering for comparison sort.Slice(results, func(i, j int) bool { return results[i].QueryName < results[j].QueryName }) assert.Equal(t, []Result{ - {"query1", 0, []map[string]string{{"iso_8601": "2017-07-10T22:08:40Z"}}, nil}, - {"query2", 0, []map[string]string{{"version": "2.4.0"}}, nil}, - {"query3", 1, []map[string]string{}, nil}, + {"query1", 0, []map[string]string{{"iso_8601": "2017-07-10T22:08:40Z"}}, nil, ""}, + {"query2", 0, []map[string]string{{"version": "2.4.0"}}, nil, ""}, + {"query3", 1, []map[string]string{}, nil, ""}, + }, + results) + + // Call writeResults -- with message included + getCalled = false + resp = plugin.Call(context.Background(), osquery.ExtensionPluginRequest{"action": "writeResults", "results": `{"queries":{"query1":[{"iso_8601":"2017-07-10T22:08:40Z"}],"query2":[{"version":"2.4.0"}]},"statuses":{"query1":"0","query2":"0","query3":"1"}, "stats":{"query1":{"wall_time_ms": 1, "user_time": 1, "system_time": 1, "memory": 1},"query2":{"wall_time_ms": 2, "user_time": 2, "system_time": 2, "memory": 2},"query3":{"wall_time_ms": 3, "user_time": 3, "system_time": 3, "memory": 3}}, "messages": {"query3": "distributed query is denylisted"}}`}) + assert.False(t, getCalled) + assert.True(t, writeCalled) + assert.Equal(t, &StatusOK, resp.Status) + // Ensure correct ordering for comparison + sort.Slice(results, func(i, j int) bool { return results[i].QueryName < results[j].QueryName }) + assert.Equal(t, []Result{ + {"query1", 0, []map[string]string{{"iso_8601": "2017-07-10T22:08:40Z"}}, &Stats{WallTimeMs: 1, UserTime: 1, SystemTime: 1, Memory: 1}, ""}, + {"query2", 0, []map[string]string{{"version": "2.4.0"}}, &Stats{WallTimeMs: 2, UserTime: 2, SystemTime: 2, Memory: 2}, ""}, + {"query3", 1, []map[string]string{}, &Stats{WallTimeMs: 3, UserTime: 3, SystemTime: 3, Memory: 3}, "distributed query is denylisted"}, }, results) }