[Kolide ATC] Construct KATC tables and add support for Firefox extension data

ee/katc/config.go

@@ -0,0 +1,120 @@
+package katc
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"log/slog"
+	"runtime"
+
+	"github.com/osquery/osquery-go"
+	"github.com/osquery/osquery-go/plugin/table"
+)
+
+// katcSourceType defines a source of data for a KATC table. The `name` is the
+// identifier parsed from the JSON KATC config, and the `dataFunc` is the function
+// that performs the query against the source.
+type katcSourceType struct {
+	name     string
+	dataFunc func(ctx context.Context, slogger *slog.Logger, path string, query string, sourceConstraints *table.ConstraintList) ([]sourceData, error)
+}
+
+// sourceData holds the result of calling `katcSourceType.dataFunc`. It maps the
+// source to the query results. (A config may have wildcards in the source,
+// allowing for querying against multiple sources.)
+type sourceData struct {
+	path string
+	rows []map[string][]byte
+}
+
+const (
+	sqliteSourceType = "sqlite"
+)
+
+func (kst *katcSourceType) UnmarshalJSON(data []byte) error {
+	var s string
+	err := json.Unmarshal(data, &s)
+	if err != nil {
+		return fmt.Errorf("unmarshalling string: %w", err)
+	}
+
+	switch s {
+	case sqliteSourceType:
+		kst.name = sqliteSourceType
+		kst.dataFunc = sqliteData
+		return nil
+	default:
+		return fmt.Errorf("unknown table type %s", s)
+	}
+}
+
+// rowTransformStep defines an operation performed against a row of data
+// returned from a source. The `name` is the identifier parsed from the
+// JSON KATC config.
+type rowTransformStep struct {
+	name          string
+	transformFunc func(ctx context.Context, slogger *slog.Logger, row map[string][]byte) (map[string][]byte, error)
+}
+
+const (
+	snappyDecodeTransformStep       = "snappy"
+	deserializeFirefoxTransformStep = "deserialize_firefox"
+)
+
+func (r *rowTransformStep) UnmarshalJSON(data []byte) error {
+	var s string
+	err := json.Unmarshal(data, &s)
+	if err != nil {
+		return fmt.Errorf("unmarshalling string: %w", err)
+	}
+
+	switch s {
+	case snappyDecodeTransformStep:
+		r.name = snappyDecodeTransformStep
+		r.transformFunc = snappyDecode
+		return nil
+	case deserializeFirefoxTransformStep:
+		r.name = deserializeFirefoxTransformStep
+		r.transformFunc = deserializeFirefox
+		return nil
+	default:
+		return fmt.Errorf("unknown data processing step %s", s)
+	}
+}
+
+// katcTableConfig is the configuration for a specific KATC table. The control server
+// sends down these configurations.
+type katcTableConfig struct {
+	SourceType        katcSourceType     `json:"source_type"`
+	Source            string             `json:"source"` // Describes how to connect to source (e.g. path to db) -- % and _ wildcards supported
+	Platform          string             `json:"platform"`
+	Columns           []string           `json:"columns"`
+	Query             string             `json:"query"` // Query to run against `path`
+	RowTransformSteps []rowTransformStep `json:"row_transform_steps"`
+}
+
+// ConstructKATCTables takes stored configuration of KATC tables, parses the configuration,
+// and returns the constructed tables.
+func ConstructKATCTables(config map[string]string, slogger *slog.Logger) []osquery.OsqueryPlugin {
+	plugins := make([]osquery.OsqueryPlugin, 0)
+	for tableName, tableConfigStr := range config {
+		var cfg katcTableConfig
+		if err := json.Unmarshal([]byte(tableConfigStr), &cfg); err != nil {
+			slogger.Log(context.TODO(), slog.LevelWarn,
+				"unable to unmarshal config for Kolide ATC table, skipping",
+				"table_name", tableName,
+				"err", err,
+			)
+			continue
+		}
+
+		if cfg.Platform != runtime.GOOS {
+			continue
+		}
+
+		t, columns := newKatcTable(tableName, cfg, slogger)
+		plugins = append(plugins, table.NewPlugin(tableName, columns, t.generate))
+	}
+
+	return plugins
+}

ee/katc/config_test.go

@@ -0,0 +1,100 @@
+package katc
+
+import (
+	_ "embed"
+	"fmt"
+	"runtime"
+	"testing"
+
+	"github.com/kolide/launcher/pkg/log/multislogger"
+	"github.com/stretchr/testify/require"
+)
+
+func TestConstructKATCTables(t *testing.T) {
+	t.Parallel()
+
+	for _, tt := range []struct {
+		testCaseName        string
+		katcConfig          map[string]string
+		expectedPluginCount int
+	}{
+		{
+			testCaseName: "snappy_sqlite",
+			katcConfig: map[string]string{
+				"kolide_snappy_sqlite_test": fmt.Sprintf(`{
+					"source_type": "sqlite",
+					"platform": "%s",
+					"columns": ["data"],
+					"source": "/some/path/to/db.sqlite",
+					"query": "SELECT data FROM object_data JOIN object_store ON (object_data.object_store_id = object_store.id) WHERE object_store.name=\"testtable\";",
+					"row_transform_steps": ["snappy"]
+				}`, runtime.GOOS),
+			},
+			expectedPluginCount: 1,
+		},
+		{
+			testCaseName: "multiple plugins",
+			katcConfig: map[string]string{
+				"test_1": fmt.Sprintf(`{
+					"source_type": "sqlite",
+					"platform": "%s",
+					"columns": ["data"],
+					"source": "/some/path/to/db.sqlite",
+					"query": "SELECT data FROM object_data;",
+					"row_transform_steps": ["snappy"]
+				}`, runtime.GOOS),
+				"test_2": fmt.Sprintf(`{
+					"source_type": "sqlite",
+					"platform": "%s",
+					"columns": ["col1", "col2"],
+					"source": "/some/path/to/a/different/db.sqlite",
+					"query": "SELECT col1, col2 FROM some_table;",
+					"row_transform_steps": []
+				}`, runtime.GOOS),
+			},
+			expectedPluginCount: 2,
+		},
+		{
+			testCaseName: "malformed config",
+			katcConfig: map[string]string{
+				"malformed_test": "this is not a config",
+			},
+			expectedPluginCount: 0,
+		},
+		{
+			testCaseName: "invalid table source",
+			katcConfig: map[string]string{
+				"kolide_snappy_test": fmt.Sprintf(`{
+					"source_type": "unknown_source",
+					"platform": "%s",
+					"columns": ["data"],
+					"source": "/some/path/to/db.sqlite",
+					"query": "SELECT data FROM object_data;"
+				}`, runtime.GOOS),
+			},
+			expectedPluginCount: 0,
+		},
+		{
+			testCaseName: "invalid data processing step type",
+			katcConfig: map[string]string{
+				"kolide_snappy_test": fmt.Sprintf(`{
+					"source_type": "sqlite",
+					"platform": "%s",
+					"columns": ["data"],
+					"source": "/some/path/to/db.sqlite",
+					"query": "SELECT data FROM object_data;",
+					"row_transform_steps": ["unknown_step"]
+				}`, runtime.GOOS),
+			},
+			expectedPluginCount: 0,
+		},
+	} {
+		tt := tt
+		t.Run(tt.testCaseName, func(t *testing.T) {
+			t.Parallel()
+
+			plugins := ConstructKATCTables(tt.katcConfig, multislogger.NewNopLogger())
+			require.Equal(t, tt.expectedPluginCount, len(plugins), "unexpected number of plugins")
+		})
+	}
+}