auto load ATC config in interactive

cmd/launcher/interactive.go

@@ -3,14 +3,16 @@ package main
 import (
 	"context"
 	"errors"
-	"flag"
 	"fmt"
 	"log/slog"
 	"os"
 	"strings"
 
 	"github.com/kolide/launcher/cmd/launcher/internal"
 	"github.com/kolide/launcher/ee/agent"
+	"github.com/kolide/launcher/ee/agent/flags"
+	"github.com/kolide/launcher/ee/agent/knapsack"
+	"github.com/kolide/launcher/ee/agent/storage/inmemory"
 	"github.com/kolide/launcher/ee/tuf"
 	"github.com/kolide/launcher/pkg/autoupdate"
 	"github.com/kolide/launcher/pkg/launcher"
@@ -19,22 +21,19 @@ import (
 )
 
 func runInteractive(systemMultiSlogger *multislogger.MultiSlogger, args []string) error {
-	flagset := flag.NewFlagSet("interactive", flag.ExitOnError)
-	var (
-		flDebug        = flagset.Bool("debug", false, "enable debug logging")
-		flOsquerydPath = flagset.String("osqueryd_path", "", "The path to the oqueryd binary")
-		flOsqueryFlags launcher.ArrayFlags
-	)
-
-	flagset.Var(&flOsqueryFlags, "osquery_flag", "Flags to pass to osquery (possibly overriding Launcher defaults)")
-
-	flagset.Usage = commandUsage(flagset, "launcher interactive")
-	if err := flagset.Parse(args); err != nil {
+	opts, err := launcher.ParseOptions("interactive", args)
+	if err != nil {
 		return err
 	}
 
+	// here we are looking for the launcher "proper" root directory so that we know where
+	// to find the kv.sqlite where we can pull the auto table construction config from
+	if opts.RootDirectory == "" {
+		opts.RootDirectory = launcher.DefaultPath(launcher.RootDirectory)
+	}
+
 	slogLevel := slog.LevelInfo
-	if *flDebug {
+	if opts.Debug {
 		slogLevel = slog.LevelDebug
 	}
 
@@ -44,36 +43,36 @@ func runInteractive(systemMultiSlogger *multislogger.MultiSlogger, args []string
 		AddSource: true,
 	}))
 
-	osquerydPath := *flOsquerydPath
-	if osquerydPath == "" {
+	if opts.OsquerydPath == "" {
 		latestOsquerydBinary, err := tuf.CheckOutLatestWithoutConfig("osqueryd", systemMultiSlogger.Logger)
 		if err != nil {
-			osquerydPath = launcher.FindOsquery()
-			if osquerydPath == "" {
+			opts.OsquerydPath = launcher.FindOsquery()
+			if opts.OsquerydPath == "" {
 				return errors.New("could not find osqueryd binary")
 			}
 			// Fall back to old autoupdate library
-			osquerydPath = autoupdate.FindNewest(context.Background(), osquerydPath)
+			opts.OsquerydPath = autoupdate.FindNewest(context.Background(), opts.OsquerydPath)
 		} else {
-			osquerydPath = latestOsquerydBinary.Path
+			opts.OsquerydPath = latestOsquerydBinary.Path
 		}
 	}
 
-	// have to keep tempdir name short so we don't exceed socket length
-	rootDir, err := agent.MkdirTemp("launcher-interactive")
+	// this is a tmp root directory that launcher can use to store files it needs to run
+	// such as the osquery socket and augeas lense files
+	interactiveRootDir, err := agent.MkdirTemp("launcher-interactive")
 	if err != nil {
 		return fmt.Errorf("creating temp dir for interactive mode: %w", err)
 	}
 
 	defer func() {
-		if err := os.RemoveAll(rootDir); err != nil {
+		if err := os.RemoveAll(interactiveRootDir); err != nil {
 			fmt.Printf("error removing launcher interactive temp dir: %s\n", err)
 		}
 	}()
 
 	hasTlsServerCertsOsqueryFlag := false
 	// check to see if we were passed a tls_server_certs flag
-	for _, v := range flOsqueryFlags {
+	for _, v := range opts.OsqueryFlags {
 		if strings.HasPrefix(v, "tls_server_certs") {
 			hasTlsServerCertsOsqueryFlag = true
 			break
@@ -82,15 +81,20 @@ func runInteractive(systemMultiSlogger *multislogger.MultiSlogger, args []string
 
 	// if we were not passed a tls_server_certs flag, pass default to osquery
 	if !hasTlsServerCertsOsqueryFlag {
-		certs, err := internal.InstallCaCerts(rootDir)
+		certs, err := internal.InstallCaCerts(interactiveRootDir)
 		if err != nil {
 			return fmt.Errorf("installing CA certs: %w", err)
 		}
 
-		flOsqueryFlags = append(flOsqueryFlags, fmt.Sprintf("tls_server_certs=%s", certs))
+		opts.OsqueryFlags = append(opts.OsqueryFlags, fmt.Sprintf("tls_server_certs=%s", certs))
 	}
 
-	osqueryProc, extensionsServer, err := interactive.StartProcess(systemMultiSlogger.Logger, rootDir, osquerydPath, flOsqueryFlags)
+	fcOpts := []flags.Option{flags.WithCmdLineOpts(opts)}
+	flagController := flags.NewFlagController(systemMultiSlogger.Logger, inmemory.NewStore(), fcOpts...)
+
+	knapsack := knapsack.New(nil, flagController, nil, systemMultiSlogger, nil)
+
+	osqueryProc, extensionsServer, err := interactive.StartProcess(knapsack, interactiveRootDir)
 	if err != nil {
 		return fmt.Errorf("error starting osqueryd: %s", err)
 	}

cmd/launcher/launcher.go

@@ -247,6 +247,13 @@ func runLauncher(ctx context.Context, cancel func(), multiSlogger, systemMultiSl
 	}
 	defer s.Close()
 
+	if err := s.WriteSettings(); err != nil {
+		slogger.Log(ctx, slog.LevelError,
+			"writing startup settings",
+			"err", err,
+		)
+	}
+
 	// If we have successfully opened the DB, and written a pid,
 	// we expect we're live. Record the version for osquery to
 	// pickup

ee/agent/startupsettings/writer.go

@@ -5,6 +5,7 @@ package startupsettings
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"log/slog"
 
@@ -43,28 +44,35 @@ func OpenWriter(ctx context.Context, knapsack types.Knapsack) (*startupSettingsW
 		},
 	}
 
-	// Attempt to ensure flags are up-to-date
-	if err := s.setFlags(); err != nil {
-		s.knapsack.Slogger().Log(ctx, slog.LevelWarn, "could not set flags", "err", err)
-	}
-
 	for k := range s.storedFlags {
 		s.knapsack.RegisterChangeObserver(s, k)
 	}
 
 	return s, nil
 }
 
-// setFlags updates the flags with their values from the agent flag data store.
-func (s *startupSettingsWriter) setFlags() error {
+// WriteSettings updates the flags with their values from the agent flag data store.
+func (s *startupSettingsWriter) WriteSettings() error {
 	updatedFlags := make(map[string]string)
 	for flag, getter := range s.storedFlags {
 		updatedFlags[flag.String()] = getter()
 	}
 	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,
+			"extracting 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)
+		return fmt.Errorf("writing settings: %w", err)
 	}
 
 	return nil
@@ -74,9 +82,9 @@ func (s *startupSettingsWriter) setFlags() error {
 // that the startup database is registered for has a new value, the startup database
 // stores that updated value.
 func (s *startupSettingsWriter) FlagsChanged(flagKeys ...keys.FlagKey) {
-	if err := s.setFlags(); err != nil {
+	if err := s.WriteSettings(); err != nil {
 		s.knapsack.Slogger().Log(context.Background(), slog.LevelError,
-			"could not set flags after change",
+			"writing startup settings after flag change",
 			"err", err,
 		)
 	}
@@ -85,3 +93,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("config"))
+	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"
+	"fmt"
 	"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/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 	_ "modernc.org/sqlite"
@@ -27,11 +32,15 @@ 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)
 	require.NoError(t, err, "expected no error setting up storage db")
 
+	require.NoError(t, s.WriteSettings(), "should be able to writing settings")
+
 	// Check that all values were set
 	v1, err := s.kvStore.Get([]byte(keys.UpdateChannel.String()))
 	require.NoError(t, err, "getting startup value")
@@ -85,10 +94,15 @@ 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")
 
+	require.NoError(t, s.WriteSettings(), "should be able to writing settings")
+
 	// Now check that all values were updated
 	v1, err = s.kvStore.Get([]byte(keys.UpdateChannel.String()))
 	require.NoError(t, err, "getting startup value")
@@ -122,10 +136,25 @@ func TestFlagsChanged(t *testing.T) {
 	pinnedOsquerydVersion := "5.3.2"
 	k.On("PinnedOsquerydVersion").Return(pinnedOsquerydVersion).Once()
 
+	autoTableConstructionValue := ulid.New()
+
+	configStore := inmemory.NewStore()
+	configMap := map[string]any{
+		"auto_table_construction":      autoTableConstructionValue,
+		"something_else_not_important": ulid.New(),
+	}
+	configJson, err := json.Marshal(configMap)
+	require.NoError(t, err, "marshalling config map")
+
+	configStore.Set([]byte("config"), 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")
 
+	require.NoError(t, s.WriteSettings(), "should be able to writing settings")
+
 	// Check that all values were set
 	v1, err := s.kvStore.Get([]byte(keys.UpdateChannel.String()))
 	require.NoError(t, err, "getting startup value")
@@ -139,6 +168,10 @@ func TestFlagsChanged(t *testing.T) {
 	require.NoError(t, err, "getting startup value")
 	require.Equal(t, pinnedOsquerydVersion, string(v3), "incorrect flag value")
 
+	v4, err := s.kvStore.Get([]byte("auto_table_construction"))
+	require.NoError(t, err, "getting startup value")
+	require.Equal(t, fmt.Sprintf("{\"auto_table_construction\":\"%s\"}", autoTableConstructionValue), string(v4), "incorrect config value")
+
 	// Now, prepare for flag changes
 	newFlagValue := "alpha"
 	k.On("UpdateChannel").Return(newFlagValue).Once()

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

@@ -15,6 +15,7 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/kolide/launcher/ee/agent"
+	"github.com/kolide/launcher/ee/agent/startupsettings"
 	"github.com/kolide/launcher/ee/agent/types"
 	"github.com/kolide/launcher/pkg/backoff"
 	"github.com/kolide/launcher/pkg/osquery/runtime/history"
@@ -517,6 +518,25 @@ func (e *Extension) GenerateConfigs(ctx context.Context) (map[string]string, err
 	} else {
 		// Store good config
 		e.knapsack.ConfigStore().Set([]byte(configKey), []byte(config))
+
+		// open the start up settings writer just to trigger a write of the config,
+		// then we can immediately close it
+		startupSettingsWriter, err := startupsettings.OpenWriter(ctx, e.knapsack)
+		if err != nil {
+			e.slogger.Log(ctx, slog.LevelError,
+				"could not get startup settings writer",
+				"err", err,
+			)
+		} else {
+			defer startupSettingsWriter.Close()
+
+			if err := startupSettingsWriter.WriteSettings(); err != nil {
+				e.slogger.Log(ctx, slog.LevelError,
+					"writing startup settings",
+					"err", err,
+				)
+			}
+		}
 		// 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/extension_test.go

@@ -58,6 +58,7 @@ func makeKnapsack(t *testing.T, db *bbolt.DB) types.Knapsack {
 	m.On("ConfigStore").Return(storageci.NewStore(t, multislogger.NewNopLogger(), storage.ConfigStore.String()))
 	m.On("Slogger").Return(multislogger.NewNopLogger())
 	m.On("ReadEnrollSecret").Maybe().Return("enroll_secret", nil)
+	m.On("RootDirectory").Maybe().Return("whatever")
 	return m
 }