update dataflatten tables to parse raw_data in query (#1732)

ee/tables/dataflattentable/tables.go

@@ -120,11 +120,12 @@ func AllTablePlugins(slogger *slog.Logger) []osquery.OsqueryPlugin {
 }
 
 func TablePlugin(slogger *slog.Logger, dataSourceType DataSourceType) osquery.OsqueryPlugin {
-	columns := Columns(table.TextColumn("path"))
+	columns := Columns(table.TextColumn("path"), table.TextColumn("raw_data"))
 
 	t := &Table{
-		tableName:       dataSourceType.TableName(),
-		flattenFileFunc: dataSourceType.FlattenFileFunc(""),
+		tableName:        dataSourceType.TableName(),
+		flattenFileFunc:  dataSourceType.FlattenFileFunc(""),
+		flattenBytesFunc: dataSourceType.FlattenBytesFunc(""),
 	}
 
 	t.slogger = slogger.With("table", t.tableName)
@@ -137,8 +138,15 @@ func (t *Table) generate(ctx context.Context, queryContext table.QueryContext) (
 	var results []map[string]string
 
 	requestedPaths := tablehelpers.GetConstraints(queryContext, "path")
-	if len(requestedPaths) == 0 {
-		return results, fmt.Errorf("The %s table requires that you specify a single constraint for path", t.tableName)
+	requestedRawDatas := tablehelpers.GetConstraints(queryContext, "raw_data")
+
+	if len(requestedPaths) == 0 && len(requestedRawDatas) == 0 {
+		return results, fmt.Errorf("The %s table requires that you specify at least one of 'path' or 'raw_data'", t.tableName)
+	}
+
+	flattenOpts := []dataflatten.FlattenOpts{
+		dataflatten.WithSlogger(t.slogger),
+		dataflatten.WithNestedPlist(),
 	}
 
 	for _, requestedPath := range requestedPaths {
@@ -151,7 +159,7 @@ func (t *Table) generate(ctx context.Context, queryContext table.QueryContext) (
 
 		for _, filePath := range filePaths {
 			for _, dataQuery := range tablehelpers.GetConstraints(queryContext, "query", tablehelpers.WithDefaults("*")) {
-				subresults, err := t.generatePath(ctx, filePath, dataQuery)
+				subresults, err := t.generatePath(ctx, filePath, dataQuery, append(flattenOpts, dataflatten.WithQuery(strings.Split(dataQuery, "/")))...)
 				if err != nil {
 					t.slogger.Log(ctx, slog.LevelInfo,
 						"failed to get data for path",
@@ -165,16 +173,43 @@ func (t *Table) generate(ctx context.Context, queryContext table.QueryContext) (
 			}
 		}
 	}
+
+	for _, rawdata := range requestedRawDatas {
+		for _, dataQuery := range tablehelpers.GetConstraints(queryContext, "query", tablehelpers.WithDefaults("*")) {
+			subresults, err := t.generateRawData(ctx, rawdata, dataQuery, append(flattenOpts, dataflatten.WithQuery(strings.Split(dataQuery, "/")))...)
+			if err != nil {
+				t.slogger.Log(ctx, slog.LevelInfo,
+					"failed to generate for raw_data",
+					"err", err,
+				)
+				continue
+			}
+
+			results = append(results, subresults...)
+		}
+	}
+
 	return results, nil
 }
 
-func (t *Table) generatePath(ctx context.Context, filePath string, dataQuery string) ([]map[string]string, error) {
-	flattenOpts := []dataflatten.FlattenOpts{
-		dataflatten.WithSlogger(t.slogger),
-		dataflatten.WithNestedPlist(),
-		dataflatten.WithQuery(strings.Split(dataQuery, "/")),
+func (t *Table) generateRawData(ctx context.Context, rawdata string, dataQuery string, flattenOpts ...dataflatten.FlattenOpts) ([]map[string]string, error) {
+	data, err := t.flattenBytesFunc([]byte(rawdata), flattenOpts...)
+	if err != nil {
+		t.slogger.Log(ctx, slog.LevelInfo,
+			"failure parsing raw data",
+			"err", err,
+		)
+		return nil, fmt.Errorf("parsing data: %w", err)
 	}
 
+	rowData := map[string]string{
+		"raw_data": rawdata,
+	}
+
+	return ToMap(data, dataQuery, rowData), nil
+}
+
+func (t *Table) generatePath(ctx context.Context, filePath string, dataQuery string, flattenOpts ...dataflatten.FlattenOpts) ([]map[string]string, error) {
 	data, err := t.flattenFileFunc(filePath, flattenOpts...)
 	if err != nil {
 		t.slogger.Log(ctx, slog.LevelInfo,

ee/tables/dataflattentable/tables_test.go

@@ -3,6 +3,7 @@ package dataflattentable
 import (
 	"context"
 	"fmt"
+	"os"
 	"path"
 	"path/filepath"
 	"sort"
@@ -24,9 +25,9 @@ func TestDataFlattenTablePlist_Animals(t *testing.T) {
 
 	// Test plist parsing both the json and xml forms
 	testTables := map[string]Table{
-		"plist": {slogger: slogger, flattenFileFunc: dataflatten.PlistFile},
-		"xml":   {slogger: slogger, flattenFileFunc: dataflatten.PlistFile},
-		"json":  {slogger: slogger, flattenFileFunc: dataflatten.JsonFile},
+		"plist": {slogger: slogger, flattenFileFunc: dataflatten.PlistFile, flattenBytesFunc: dataflatten.Plist},
+		"xml":   {slogger: slogger, flattenFileFunc: dataflatten.PlistFile, flattenBytesFunc: dataflatten.Plist},
+		"json":  {slogger: slogger, flattenFileFunc: dataflatten.JsonFile, flattenBytesFunc: dataflatten.Json},
 	}
 
 	var tests = []struct {
@@ -57,13 +58,14 @@ func TestDataFlattenTablePlist_Animals(t *testing.T) {
 	for _, tt := range tests {
 		for dataType, tableFunc := range testTables {
 			testFile := filepath.Join("testdata", "animals."+dataType)
-			mockQC := tablehelpers.MockQueryContext(map[string][]string{
+
+			// test file path
+			mockPathQC := tablehelpers.MockQueryContext(map[string][]string{
 				"path":  {testFile},
 				"query": tt.queries,
 			})
 
-			rows, err := tableFunc.generate(context.TODO(), mockQC)
-
+			rows, err := tableFunc.generate(context.TODO(), mockPathQC)
 			require.NoError(t, err)
 
 			// delete the path and query keys, so we don't need to enumerate them in the test case
@@ -77,6 +79,30 @@ func TestDataFlattenTablePlist_Animals(t *testing.T) {
 			sort.SliceStable(rows, func(i, j int) bool { return rows[i]["fullkey"] < rows[j]["fullkey"] })
 
 			require.EqualValues(t, tt.expected, rows, "table type %s test", dataType)
+
+			// test bytes path
+			raw_data, err := os.ReadFile(testFile)
+			require.NoError(t, err)
+
+			mockBytesQC := tablehelpers.MockQueryContext(map[string][]string{
+				"raw_data": {string(raw_data)},
+				"query":    tt.queries,
+			})
+
+			rows, err = tableFunc.generate(context.TODO(), mockBytesQC)
+			require.NoError(t, err)
+
+			// delete the query keys, so we don't need to enumerate them in the test case
+			for _, row := range rows {
+				delete(row, "query")
+				delete(row, "raw_data")
+			}
+
+			// Despite being an array. data is returned unordered. Sort it.
+			sort.SliceStable(tt.expected, func(i, j int) bool { return tt.expected[i]["fullkey"] < tt.expected[j]["fullkey"] })
+			sort.SliceStable(rows, func(i, j int) bool { return rows[i]["fullkey"] < rows[j]["fullkey"] })
+
+			require.EqualValues(t, tt.expected, rows, "table type %s test", dataType)
 		}
 	}
 

ee/tables/dataflattentable/testdata/animals.json

@@ -7,24 +7,24 @@
   "users": [
     {
       "favorites": [
-	"ants"
+        "ants"
       ],
       "uuid": "abc123",
       "name": "Alex Aardvark",
       "id": 1
     },
     {
       "favorites": [
-	"mice",
-	"birds"
+        "mice",
+        "birds"
       ],
       "uuid": "def456",
       "name": "Bailey Bobcat",
       "id": 2
     },
     {
       "favorites": [
-	"seeds"
+        "seeds"
       ],
       "uuid": "ghi789",
       "name": "Cam Chipmunk",