Add orbit_info table extension (#4587)

orbit/cmd/orbit/orbit.go

@@ -23,6 +23,7 @@ import (
 	"github.com/fleetdm/fleet/v4/pkg/certificate"
 	"github.com/fleetdm/fleet/v4/pkg/file"
 	"github.com/fleetdm/fleet/v4/pkg/secure"
+	"github.com/google/uuid"
 	"github.com/oklog/run"
 	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
@@ -181,6 +182,11 @@ func main() {
 			return fmt.Errorf("initialize root dir: %w", err)
 		}
 
+		deviceAuthToken, err := loadOrGenerateToken(c.String("root-dir"))
+		if err != nil {
+			return fmt.Errorf("load identifier file: %w", err)
+		}
+
 		localStore, err := filestore.New(filepath.Join(c.String("root-dir"), "tuf-metadata.json"))
 		if err != nil {
 			log.Fatal().Err(err).Msg("failed to create local metadata store")
@@ -395,7 +401,9 @@ func main() {
 		}
 		g.Add(r.Execute, r.Interrupt)
 
-		ext := table.NewRunner(r.ExtensionSocketPath())
+		ext := table.NewRunner(r.ExtensionSocketPath(), table.WithExtension(orbitInfoExtension{
+			deviceAuthToken: deviceAuthToken,
+		}))
 		g.Add(ext.Execute, ext.Interrupt)
 
 		// Install a signal handler
@@ -415,6 +423,26 @@ func main() {
 	}
 }
 
+func loadOrGenerateToken(rootDir string) (string, error) {
+	filePath := filepath.Join(rootDir, "identifier")
+	id, err := ioutil.ReadFile(filePath)
+	switch {
+	case err == nil:
+		return string(id), nil
+	case errors.Is(err, os.ErrNotExist):
+		id, err := uuid.NewRandom()
+		if err != nil {
+			return "", fmt.Errorf("generate identifier: %w", err)
+		}
+		if err := ioutil.WriteFile(filePath, []byte(id.String()), constant.DefaultFileMode); err != nil {
+			return "", fmt.Errorf("write identifier file %q: %w", filePath, err)
+		}
+		return id.String(), nil
+	default:
+		return "", fmt.Errorf("load identifier file %q: %w", filePath, err)
+	}
+}
+
 var versionCommand = &cli.Command{
 	Name:  "version",
 	Usage: "Get the orbit version",

orbit/cmd/orbit/orbit_info.go

@@ -0,0 +1,42 @@
+package main
+
+import (
+	"context"
+
+	orbit_table "github.com/fleetdm/fleet/v4/orbit/pkg/table"
+	"github.com/kolide/osquery-go/plugin/table"
+)
+
+// orbitInfoExtension implements an extension table that provides info about Orbit.
+type orbitInfoExtension struct {
+	deviceAuthToken string
+}
+
+var _ orbit_table.Extension = orbitInfoExtension{}
+
+// Name partially implements orbit_table.Extension.
+func (o orbitInfoExtension) Name() string {
+	return "orbit_info"
+}
+
+// Columns partially implements orbit_table.Extension.
+func (o orbitInfoExtension) Columns() []table.ColumnDefinition {
+	return []table.ColumnDefinition{
+		table.TextColumn("version"),
+		table.TextColumn("device_auth_token"),
+	}
+}
+
+// GenerateFunc partially implements orbit_table.Extension.
+func (o orbitInfoExtension) GenerateFunc(ctx context.Context, _ table.QueryContext) ([]map[string]string, error) {
+	v := version
+	if v == "" {
+		v = "unknown"
+	}
+	return []map[string]string{
+		{
+			"version":           v,
+			"device_auth_token": o.deviceAuthToken,
+		},
+	}, nil
+}

orbit/cmd/orbit/orbit_test.go

@@ -0,0 +1,22 @@
+package main
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestLoadOrGenerateToken(t *testing.T) {
+	dir := t.TempDir()
+	token, err := loadOrGenerateToken(dir)
+	require.NoError(t, err)
+	require.NotEmpty(t, token)
+	token2, err := loadOrGenerateToken(dir)
+	require.NoError(t, err)
+	require.Equal(t, token, token2)
+	dir2 := t.TempDir()
+	token3, err := loadOrGenerateToken(dir2)
+	require.NoError(t, err)
+	require.NotEmpty(t, token3)
+	require.NotEqual(t, token, token3)
+}

orbit/pkg/table/extension.go

@@ -17,14 +17,38 @@ import (
 
 // Runner wraps the osquery extension manager with okglog/run Execute and Interrupt functions.
 type Runner struct {
-	socket string
+	socket          string
-	srv    *osquery.ExtensionManagerServer
+	tableExtensions []Extension
-	cancel func()
+	srv             *osquery.ExtensionManagerServer
+	cancel          func()
+}
+
+// Extension implements a osquery-go table extension.
+type Extension interface {
+	// Name returns the name of the table.
+	Name() string
+	// Column returns the definition of the table columns.
+	Columns() []table.ColumnDefinition
+	// GenerateFunc generates results for a query.
+	GenerateFunc(ctx context.Context, queryContext table.QueryContext) ([]map[string]string, error)
+}
+
+// Opt allows configuring a Runner.
+type Opt func(*Runner)
+
+// WithExtension registers the given Extension on the Runner.
+func WithExtension(t Extension) Opt {
+	return func(r *Runner) {
+		r.tableExtensions = append(r.tableExtensions, t)
+	}
 }
 
 // NewRunner creates an extension runner.
-func NewRunner(socket string) *Runner {
+func NewRunner(socket string, opts ...Opt) *Runner {
 	r := &Runner{socket: socket}
+	for _, fn := range opts {
+		fn(r)
+	}
 	return r
 }
 
@@ -72,6 +96,13 @@ func (r *Runner) Execute() error {
 		table.NewPlugin("file_lines", fileline.FileLineColumns(), fileline.FileLineGenerate),
 	}
 	plugins = append(plugins, platformTables()...)
+	for _, t := range r.tableExtensions {
+		plugins = append(plugins, table.NewPlugin(
+			t.Name(),
+			t.Columns(),
+			t.GenerateFunc,
+		))
+	}
 	r.srv.RegisterPlugin(plugins...)
 
 	if err := r.srv.Run(); err != nil {