fleet/tools/dibble/pkg/command/hosts.go

package command

import (
	"context"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"strings"
	"syscall"

	"github.com/AlecAivazis/survey/v2"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

// newHostsCmd wires `dibble hosts` — a thin convenience over running
// cmd/osquery-perf. dibble doesn't seed hosts (osquery-perf does that), but
// looking up the right team's enroll secret and stringing the command
// together by hand is enough friction that this subcommand pays for itself.
//
// Flow:
//  1. List teams (and the global / no-team scope) from the Fleet API.
//  2. Radio-button pick — single-select via survey.
//  3. Fetch that team's enroll secret.
//  4. Build the `go run cmd/osquery-perf/agent.go -enroll_secret … -server_url …`
//     command, print it, and offer to run it inline. While the child runs,
//     Ctrl-C is forwarded to it so the user can shut it down cleanly.
func newHostsCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "hosts",
		Short: "Pick a team and run osquery-perf against it (single-select radio + inline run)",
		Long: `Lists existing teams (fleets), lets you pick one with a radio button, fetches
that team's enroll secret, and either prints the matching osquery-perf command
or runs it for you.

Picking the global / no-team option uses the global enroll secret from
/api/latest/fleet/spec/enroll_secret.

While osquery-perf runs inline, Ctrl-C is forwarded to it so you can stop the
simulated hosts cleanly.

dibble doesn't seed hosts itself — this is a convenience around osquery-perf,
which is the canonical host simulator (see cmd/osquery-perf).`,
		RunE: func(cmd *cobra.Command, args []string) error {
			if err := requireConfig(); err != nil {
				return err
			}
			c, err := newClientFromViper()
			if err != nil {
				return err
			}

			teamName, _ := cmd.Flags().GetString("team-name")
			teamID, _ := cmd.Flags().GetInt("team-id")
			run, _ := cmd.Flags().GetBool("run")
			printOnly, _ := cmd.Flags().GetBool("print-only")
			hostCount, _ := cmd.Flags().GetInt("host-count")

			picked, err := pickHostsTeam(c, teamID, teamName)
			if err != nil {
				return err
			}

			secret, err := fetchEnrollSecret(c, picked)
			if err != nil {
				return err
			}

			repoRoot, err := findFleetRepoRoot()
			if err != nil {
				return err
			}

			fullCmd, displayCmd := buildOsqueryPerfCmd(repoRoot, secret, viper.GetString(keyFleetURL), hostCount)
			scope := picked.scopeDescription()
			printf("osquery-perf command for %s:", scope)
			fmt.Fprintln(stdoutish, "\n"+displayCmd+"\n")

			switch {
			case printOnly:
				return nil
			case run:
				return runOsqueryPerf(fullCmd)
			default:
				return maybeRunOsqueryPerf(fullCmd)
			}
		},
	}
	cmd.Flags().Int("team-id", -1, "Team to enroll into. -1 (default) = interactive pick. 0 = no team / global.")
	cmd.Flags().String("team-name", "", "Team to enroll into, by name. Overrides --team-id.")
	cmd.Flags().Int("host-count", 1, "Number of simulated hosts (passed to osquery-perf -host_count)")
	cmd.Flags().Bool("run", false, "Skip the confirm prompt and run osquery-perf immediately")
	cmd.Flags().Bool("print-only", false, "Skip the confirm prompt and just print the command")
	return cmd
}

// pickedTeam is the outcome of the team picker. ID == 0 means the global /
// no-team scope.
type pickedTeam struct {
	ID   uint
	Name string
}

func (p pickedTeam) scopeDescription() string {
	if p.ID == 0 {
		return "no team (global)"
	}
	return fmt.Sprintf("team %q (id=%d)", p.Name, p.ID)
}

// pickHostsTeam resolves which team to enroll into. Precedence mirrors
// resolveSoftwareTeamID in software.go:
//
//  1. --team-name (looked up against Fleet)
//  2. --team-id (0 = global, ≥1 = that team)
//  3. Interactive single-select via survey
func pickHostsTeam(c *Client, flagID int, flagName string) (pickedTeam, error) {
	teams, err := listExistingTeams(c)
	if err != nil {
		return pickedTeam{}, fmt.Errorf("list teams: %w", err)
	}

	if flagName != "" {
		for _, t := range teams {
			if strings.EqualFold(t.Name, flagName) {
				return pickedTeam{ID: t.ID, Name: t.Name}, nil
			}
		}
		return pickedTeam{}, fmt.Errorf("no team named %q (try --team-id or run `dibble teams` first)", flagName)
	}
	if flagID == 0 {
		return pickedTeam{ID: 0, Name: ""}, nil
	}
	if flagID > 0 {
		for _, t := range teams {
			if t.ID == uint(flagID) {
				return pickedTeam{ID: t.ID, Name: t.Name}, nil
			}
		}
		// Honor explicit ID even if we couldn't resolve a name.
		return pickedTeam{ID: uint(flagID), Name: ""}, nil
	}

	options := []string{"(no team — global)"}
	for _, t := range teams {
		options = append(options, fmt.Sprintf("%s (id=%d)", t.Name, t.ID))
	}
	def := options[0]
	if len(options) > 1 {
		def = options[1]
	}
	var picked string
	if err := survey.AskOne(&survey.Select{
		Message: "Which team should the simulated hosts enroll into?",
		Options: options,
		Default: def,
	}, &picked); err != nil {
		return pickedTeam{}, err
	}
	if picked == options[0] {
		return pickedTeam{ID: 0, Name: ""}, nil
	}
	for _, t := range teams {
		if picked == fmt.Sprintf("%s (id=%d)", t.Name, t.ID) {
			return pickedTeam{ID: t.ID, Name: t.Name}, nil
		}
	}
	return pickedTeam{}, errors.New("no team selected")
}

// fetchEnrollSecret returns the first secret for the given team, or the
// first global secret when picked.ID == 0.
func fetchEnrollSecret(c *Client, picked pickedTeam) (string, error) {
	if picked.ID == 0 {
		var resp struct {
			Spec struct {
				Secrets []struct {
					Secret string `json:"secret"`
				} `json:"secrets"`
			} `json:"spec"`
		}
		if err := c.Get("/api/latest/fleet/spec/enroll_secret", &resp); err != nil {
			return "", fmt.Errorf("fetch global enroll secret: %w", err)
		}
		if len(resp.Spec.Secrets) == 0 {
			return "", errors.New("global enroll secret is empty — set one in Fleet settings or run `dibble enroll-secrets`")
		}
		return resp.Spec.Secrets[0].Secret, nil
	}

	var resp struct {
		Secrets []struct {
			Secret string `json:"secret"`
		} `json:"secrets"`
	}
	if err := c.Get(fmt.Sprintf("/api/latest/fleet/fleets/%d/secrets", picked.ID), &resp); err != nil {
		return "", fmt.Errorf("fetch enroll secret for team %d: %w", picked.ID, err)
	}
	if len(resp.Secrets) == 0 {
		return "", fmt.Errorf("team %d has no enroll secret — run `dibble enroll-secrets` to add one", picked.ID)
	}
	return resp.Secrets[0].Secret, nil
}

// findFleetRepoRoot walks up from the current working directory until it
// finds the parent fleet checkout (the one with cmd/osquery-perf/agent.go).
// Returned path is absolute.
func findFleetRepoRoot() (string, error) {
	cwd, err := os.Getwd()
	if err != nil {
		return "", fmt.Errorf("get cwd: %w", err)
	}
	dir := cwd
	for {
		if _, err := os.Stat(filepath.Join(dir, "cmd", "osquery-perf", "agent.go")); err == nil {
			return dir, nil
		}
		parent := filepath.Dir(dir)
		if parent == dir {
			return "", fmt.Errorf("could not find fleet repo root (cmd/osquery-perf/agent.go) walking up from %s", cwd)
		}
		dir = parent
	}
}

// buildOsqueryPerfCmd assembles the slice of arguments passed to exec.Cmd
// alongside a human-readable rendering for printing back to the user.
func buildOsqueryPerfCmd(repoRoot, secret, serverURL string, hostCount int) (full []string, display string) {
	if serverURL == "" {
		serverURL = "https://localhost:8080"
	}
	if hostCount < 1 {
		hostCount = 1
	}
	agentPath := filepath.Join(repoRoot, "cmd", "osquery-perf", "agent.go")
	full = []string{
		"go", "run", agentPath,
		"-enroll_secret", secret,
		"-server_url", serverURL,
		"-host_count", fmt.Sprintf("%d", hostCount),
	}

	// Quote the secret for display so copy-paste works in shells.
	display = fmt.Sprintf("go run %s -enroll_secret %q -server_url %s -host_count %d",
		agentPath, secret, serverURL, hostCount)
	return full, display
}

// maybeRunOsqueryPerf prompts the user, and if confirmed runs osquery-perf
// inline with Ctrl-C forwarded. Non-interactive sessions (no TTY, survey
// returns error) default to printing only.
func maybeRunOsqueryPerf(args []string) error {
	wantRun := true
	if err := survey.AskOne(&survey.Confirm{
		Message: "Run osquery-perf now? (Ctrl-C to stop)",
		Default: true,
	}, &wantRun); err != nil {
		// Non-interactive — fall back to print-only behavior.
		printf("not running osquery-perf (interactive prompt failed: %v)", err)
		return nil
	}
	if !wantRun {
		return nil
	}
	return runOsqueryPerf(args)
}

// runOsqueryPerf executes the command with stdout / stderr passthrough and
// forwards SIGINT / SIGTERM to the child so Ctrl-C stops the simulated
// hosts cleanly. Returns whatever exit error the child raises, except that
// a signal-initiated shutdown is treated as success.
//
// Used to use exec.CommandContext + signal.Notify with manual Process.Signal
// calls so the child gets the signal directly rather than the cancel-style
// kill exec.CommandContext does by default.
func runOsqueryPerf(args []string) error {
	if len(args) < 2 {
		return errors.New("runOsqueryPerf: empty command")
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	cmd := exec.CommandContext(ctx, args[0], args[1:]...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin

	// Disable the default ctx-cancel kill so we control shutdown ourselves.
	// (Go 1.20+ exec wires Cancel to SIGKILL; we want SIGINT so osquery-perf
	// can run its own deferred cleanup.)
	cmd.Cancel = func() error {
		if cmd.Process == nil {
			return nil
		}
		return cmd.Process.Signal(syscall.SIGINT)
	}

	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
	defer signal.Stop(sigCh)

	if err := cmd.Start(); err != nil {
		return fmt.Errorf("start osquery-perf: %w", err)
	}

	// Forward the first signal we receive; further signals from the user
	// keep getting forwarded so a hung child can be killed harder by
	// repeated Ctrl-C.
	go func() {
		for sig := range sigCh {
			if cmd.Process == nil {
				return
			}
			_ = cmd.Process.Signal(sig)
		}
	}()

	err := cmd.Wait()
	// A SIGINT-induced exit looks like an error to exec.Wait; treat it as
	// success when we asked the child to stop.
	if err != nil {
		var exitErr *exec.ExitError
		if errors.As(err, &exitErr) {
			if status, ok := exitErr.Sys().(syscall.WaitStatus); ok {
				if status.Signaled() && (status.Signal() == syscall.SIGINT || status.Signal() == syscall.SIGTERM) {
					return nil
				}
			}
		}
		return err
	}
	return nil
}