auto load ATC config in interactive

ee/agent/startupsettings/writer.go

@@ -5,12 +5,14 @@ package startupsettings
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"log/slog"
 
 	"github.com/kolide/launcher/ee/agent/flags/keys"
 	agentsqlite "github.com/kolide/launcher/ee/agent/storage/sqlite"
 	"github.com/kolide/launcher/ee/agent/types"
+	"github.com/kolide/launcher/pkg/osquery"
 	"github.com/kolide/launcher/pkg/traces"
 )
 
@@ -63,6 +65,18 @@ func (s *startupSettingsWriter) setFlags() error {
 	}
 	updatedFlags["use_tuf_autoupdater"] = "enabled" // Hardcode for backwards compatibility circa v1.5.3
 
+	// Set flags will only be called when a flag value changes. The osquery config that contains the atc config
+	// comes in via osquery extension. So a new config will not trigger a re-write.
+	atcConfig, err := s.extractAutoTableConstructionConfig()
+	if err != nil {
+		s.knapsack.Slogger().Log(context.TODO(), slog.LevelDebug,
+			"could not extract auto_table_construction config",
+			"err", err,
+		)
+	} else {
+		updatedFlags["auto_table_construction"] = atcConfig
+	}
+
 	if _, err := s.kvStore.Update(updatedFlags); err != nil {
 		return fmt.Errorf("updating flags: %w", err)
 	}
@@ -85,3 +99,31 @@ func (s *startupSettingsWriter) FlagsChanged(flagKeys ...keys.FlagKey) {
 func (s *startupSettingsWriter) Close() error {
 	return s.kvStore.Close()
 }
+
+func (s *startupSettingsWriter) extractAutoTableConstructionConfig() (string, error) {
+	osqConfig, err := s.knapsack.ConfigStore().Get([]byte(osquery.ConfigKey))
+	if err != nil {
+		return "", fmt.Errorf("could not get osquery config from store: %w", err)
+	}
+
+	// convert osquery config to map
+	var configUnmarshalled map[string]json.RawMessage
+	if err := json.Unmarshal(osqConfig, &configUnmarshalled); err != nil {
+		return "", fmt.Errorf("could not unmarshal osquery config: %w", err)
+	}
+
+	// delete what we don't want
+	for k := range configUnmarshalled {
+		if k == "auto_table_construction" {
+			continue
+		}
+		delete(configUnmarshalled, k)
+	}
+
+	atcJson, err := json.Marshal(configUnmarshalled)
+	if err != nil {
+		return "", fmt.Errorf("could not marshal auto_table_construction: %w", err)
+	}
+
+	return string(atcJson), nil
+}

ee/agent/startupsettings/writer_test.go

@@ -2,12 +2,17 @@ package startupsettings
 
 import (
 	"context"
+	"encoding/json"
 	"testing"
 
 	_ "github.com/golang-migrate/migrate/v4/database/sqlite"
+	"github.com/kolide/kit/ulid"
 	"github.com/kolide/launcher/ee/agent/flags/keys"
+	"github.com/kolide/launcher/ee/agent/storage/inmemory"
 	agentsqlite "github.com/kolide/launcher/ee/agent/storage/sqlite"
 	typesmocks "github.com/kolide/launcher/ee/agent/types/mocks"
+	"github.com/kolide/launcher/pkg/log/multislogger"
+	"github.com/kolide/launcher/pkg/osquery"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 	_ "modernc.org/sqlite"
@@ -27,6 +32,8 @@ func TestOpenWriter_NewDatabase(t *testing.T) {
 	k.On("UpdateChannel").Return(updateChannelVal)
 	k.On("PinnedLauncherVersion").Return("")
 	k.On("PinnedOsquerydVersion").Return("5.11.0")
+	k.On("ConfigStore").Return(inmemory.NewStore())
+	k.On("Slogger").Return(multislogger.NewNopLogger())
 
 	// Set up storage db, which should create the database and set all flags
 	s, err := OpenWriter(context.TODO(), k)
@@ -85,6 +92,9 @@ func TestOpenWriter_DatabaseAlreadyExists(t *testing.T) {
 	k.On("PinnedLauncherVersion").Return(pinnedLauncherVersion)
 	k.On("PinnedOsquerydVersion").Return(pinnedOsquerydVersion)
 
+	k.On("ConfigStore").Return(inmemory.NewStore())
+	k.On("Slogger").Return(multislogger.NewNopLogger())
+
 	// Set up storage db, which should create the database and set all flags
 	s, err := OpenWriter(context.TODO(), k)
 	require.NoError(t, err, "expected no error setting up storage db")
@@ -122,6 +132,17 @@ func TestFlagsChanged(t *testing.T) {
 	pinnedOsquerydVersion := "5.3.2"
 	k.On("PinnedOsquerydVersion").Return(pinnedOsquerydVersion).Once()
 
+	configStore := inmemory.NewStore()
+	configMap := map[string]any{
+		"auto_table_construction":      ulid.New(),
+		"something_else_not_important": ulid.New(),
+	}
+	configJson, err := json.Marshal(configMap)
+	require.NoError(t, err, "marshalling config map")
+
+	configStore.Set([]byte(osquery.ConfigKey), configJson)
+	k.On("ConfigStore").Return(configStore)
+
 	// Set up storage db, which should create the database and set all flags
 	s, err := OpenWriter(context.TODO(), k)
 	require.NoError(t, err, "expected no error setting up storage db")

pkg/launcher/paths.go

@@ -1,6 +1,7 @@
 package launcher
 
 import (
+	"os"
 	"path/filepath"
 	"runtime"
 )
@@ -56,7 +57,14 @@ func DefaultPath(path defaultPath) string {
 	// not windows
 	switch path {
 	case RootDirectory:
-		return "/var/kolide-k2/k2device.kolide.com/"
+		const defaultRootDir = "/var/kolide-k2/k2device.kolide.com"
+
+		// see if default root dir exists, if not assume it's a preprod install
+		if _, err := os.Stat(defaultRootDir); err != nil {
+			return "/var/kolide-k2/k2device-preprod.kolide.com"
+		}
+
+		return defaultRootDir
 	case EtcDirectory:
 		return "/etc/kolide-k2/"
 	case BinDirectory:

pkg/osquery/extension.go

@@ -51,7 +51,7 @@ const (
 	// DB key for node key
 	nodeKeyKey = "nodeKey"
 	// DB key for last retrieved config
-	configKey = "config"
+	ConfigKey = "config"
 	// DB keys for the rsa keys
 	privateKeyKey = "privateKey"
 
@@ -332,7 +332,7 @@ func NodeKey(getter types.Getter) (string, error) {
 
 // Config returns the device config from the storage layer
 func Config(getter types.Getter) (string, error) {
-	key, err := getter.Get([]byte(configKey))
+	key, err := getter.Get([]byte(ConfigKey))
 	if err != nil {
 		return "", fmt.Errorf("error getting config key: %w", err)
 	}
@@ -504,7 +504,7 @@ func (e *Extension) GenerateConfigs(ctx context.Context) (map[string]string, err
 		)
 		// Try to use cached config
 		var confBytes []byte
-		confBytes, _ = e.knapsack.ConfigStore().Get([]byte(configKey))
+		confBytes, _ = e.knapsack.ConfigStore().Get([]byte(ConfigKey))
 
 		if len(confBytes) == 0 {
 			if !e.enrolled() {
@@ -516,7 +516,7 @@ func (e *Extension) GenerateConfigs(ctx context.Context) (map[string]string, err
 		config = string(confBytes)
 	} else {
 		// Store good config
-		e.knapsack.ConfigStore().Set([]byte(configKey), []byte(config))
+		e.knapsack.ConfigStore().Set([]byte(ConfigKey), []byte(config))
 		// TODO log or record metrics when caching config fails? We
 		// would probably like to return the config and not an error in
 		// this case.

pkg/osquery/interactive/interactive.go

@@ -1,24 +1,31 @@
 package interactive
 
 import (
+	"context"
 	"fmt"
 	"log/slog"
 	"os"
 	"path/filepath"
 	"runtime"
+	"strings"
 	"time"
 
 	"github.com/kolide/kit/fsutil"
+	"github.com/kolide/launcher/ee/agent/startupsettings"
 	"github.com/kolide/launcher/pkg/augeas"
+	"github.com/kolide/launcher/pkg/launcher"
 	osqueryRuntime "github.com/kolide/launcher/pkg/osquery/runtime"
 	"github.com/kolide/launcher/pkg/osquery/table"
 	osquery "github.com/osquery/osquery-go"
+	"github.com/osquery/osquery-go/plugin/config"
 )
 
-const extensionName = "com.kolide.launcher_interactive"
+const (
+	extensionName           = "com.kolide.launcher_interactive"
+	defaultConfigPluginName = "interactive_config"
+)
 
 func StartProcess(slogger *slog.Logger, rootDir, osquerydPath string, osqueryFlags []string) (*os.Process, *osquery.ExtensionManagerServer, error) {
-
 	if err := os.MkdirAll(rootDir, fsutil.DirMode); err != nil {
 		return nil, nil, fmt.Errorf("creating root dir for interactive mode: %w", err)
 	}
@@ -37,6 +44,35 @@ func StartProcess(slogger *slog.Logger, rootDir, osquerydPath string, osqueryFla
 		}
 	}
 
+	// check to see if a config flag path was given,
+	// we need to check this before loading the default config plugin,
+	// passing 2 configs to osquery will result in an error
+	haveConfigPathOsqFlag := false
+	for _, flag := range osqueryFlags {
+		if strings.HasPrefix(flag, "config_path") {
+			haveConfigPathOsqFlag = true
+			break
+		}
+	}
+
+	// start building list of osq plugins with the kolide tables
+	osqPlugins := table.PlatformTables(slogger, osquerydPath)
+
+	// if we were not provided a config path flag, try to add default config
+	if !haveConfigPathOsqFlag {
+		// check to see if we can actually get a config plugin
+		configPlugin, err := configPlugin()
+		if err != nil {
+			slogger.Log(context.TODO(), slog.LevelDebug,
+				"error creating config plugin",
+				"err", err,
+			)
+		} else {
+			osqPlugins = append(osqPlugins, configPlugin)
+			osqueryFlags = append(osqueryFlags, fmt.Sprintf("config_plugin=%s", defaultConfigPluginName))
+		}
+	}
+
 	proc, err := os.StartProcess(osquerydPath, buildOsqueryFlags(socketPath, augeasLensesPath, osqueryFlags), &os.ProcAttr{
 		// Transfer stdin, stdout, and stderr to the new process
 		Files: []*os.File{os.Stdin, os.Stdout, os.Stderr},
@@ -46,7 +82,7 @@ func StartProcess(slogger *slog.Logger, rootDir, osquerydPath string, osqueryFla
 		return nil, nil, fmt.Errorf("error starting osqueryd in interactive mode: %w", err)
 	}
 
-	// while developing for windows it was found that it will sometimes take osquey a while
+	// while developing for windows it was found that it will sometimes take osquery a while
 	// to create the socket, so we wait for it to exist before continuing
 	if err := waitForFile(socketPath, time.Second/4, time.Second*10); err != nil {
 
@@ -58,7 +94,7 @@ func StartProcess(slogger *slog.Logger, rootDir, osquerydPath string, osqueryFla
 		return nil, nil, fmt.Errorf("error waiting for osquery to create socket: %w", err)
 	}
 
-	extensionServer, err := loadExtensions(slogger, socketPath, osquerydPath)
+	extensionServer, err := loadExtensions(socketPath, osqPlugins...)
 	if err != nil {
 		err = fmt.Errorf("error loading extensions: %w", err)
 
@@ -74,7 +110,6 @@ func StartProcess(slogger *slog.Logger, rootDir, osquerydPath string, osqueryFla
 }
 
 func buildOsqueryFlags(socketPath, augeasLensesPath string, osqueryFlags []string) []string {
-
 	// putting "-S" (the interactive flag) first because the behavior is inconsistent
 	// when it's in the middle, found this during development on M1 macOS monterey 12.4
 	// ~James Pickett 07/05/2022
@@ -100,7 +135,7 @@ func buildOsqueryFlags(socketPath, augeasLensesPath string, osqueryFlags []strin
 	return flags
 }
 
-func loadExtensions(slogger *slog.Logger, socketPath string, osquerydPath string) (*osquery.ExtensionManagerServer, error) {
+func loadExtensions(socketPath string, plugins ...osquery.OsqueryPlugin) (*osquery.ExtensionManagerServer, error) {
 	client, err := osquery.NewClient(socketPath, 10*time.Second, osquery.MaxWaitTime(10*time.Second))
 	if err != nil {
 		return nil, fmt.Errorf("error creating osquery client: %w", err)
@@ -117,7 +152,7 @@ func loadExtensions(slogger *slog.Logger, socketPath string, osquerydPath string
 		return extensionManagerServer, fmt.Errorf("error creating extension manager server: %w", err)
 	}
 
-	extensionManagerServer.RegisterPlugin(table.PlatformTables(slogger, osquerydPath)...)
+	extensionManagerServer.RegisterPlugin(plugins...)
 
 	if err := extensionManagerServer.Start(); err != nil {
 		return nil, fmt.Errorf("error starting extension manager server: %w", err)
@@ -153,3 +188,20 @@ func waitForFile(path string, interval, timeout time.Duration) error {
 		}
 	}
 }
+
+func configPlugin() (*config.Plugin, error) {
+	r, err := startupsettings.OpenReader(context.TODO(), launcher.DefaultPath(launcher.RootDirectory))
+	if err != nil {
+		return nil, fmt.Errorf("error opening startup settings reader: %w", err)
+	}
+	defer r.Close()
+
+	atcConfig, err := r.Get("auto_table_construction")
+	if err != nil {
+		return nil, fmt.Errorf("error getting auto_table_construction from startup settings: %w", err)
+	}
+
+	return config.NewPlugin(defaultConfigPluginName, func(ctx context.Context) (map[string]string, error) {
+		return map[string]string{defaultConfigPluginName: atcConfig}, nil
+	}), nil
+}

pkg/osquery/interactive/interactive_test.go

@@ -15,6 +15,7 @@ import (
 	"testing"
 
 	"github.com/kolide/kit/fsutil"
+	"github.com/kolide/kit/ulid"
 	"github.com/kolide/launcher/pkg/log/multislogger"
 	"github.com/kolide/launcher/pkg/packaging"
 	"github.com/stretchr/testify/require"
@@ -69,6 +70,13 @@ func TestProc(t *testing.T) {
 			},
 			wantProc: true,
 		},
+		{
+			name: "config path",
+			osqueryFlags: []string{
+				fmt.Sprintf("config_path=%s", ulid.New()),
+			},
+			wantProc: true,
+		},
 		{
 			name:           "socket path too long, the name of the test causes the socket path to be to long to be created, resulting in timeout waiting for the socket",
 			wantProc:       false,